jp: fix file resolving

Original-upstream-commit: c933446486e6923c8468a979b6fab890b46f5cdb
Upstream: https://github.com/commaai/openpilot/commit/c933446486e6923c8468a979b6fab890b46f5cdb

openpilot/tools/jotpluggler/sketch_layout.cc

@@ -320,20 +320,18 @@ std::map<int, SegmentLogs> trim_segments(std::map<int, SegmentLogs> segments, co
 
 std::map<int, SegmentLogs> load_segments_from_json(const json11::Json &json) {
   std::map<int, SegmentLogs> segments;
-  static const std::regex rx(R"(\/(\d+)\/)");
-  for (const auto &value : json.object_items()) {
-    for (const auto &url : value.second.array_items()) {
-      const std::string url_str = url.string_value();
-      std::smatch match;
-      if (!std::regex_search(url_str, match, rx)) continue;
-      add_log_file_to_segments(&segments, std::stoi(match[1].str()), url_str);
-    }
+  for (const auto &[number, files] : json.object_items()) {
+    segments[std::stoi(number)] = {
+      files["rlog"].string_value(), files["qlog"].string_value(), files["narrow_road"].string_value(),
+      files["cabin"].string_value(), files["wide_road"].string_value(), files["qcamera"].string_value(),
+    };
   }
   return segments;
 }
 
 std::map<int, SegmentLogs> load_segments_from_server(const RouteSelection &route) {
-  const std::string result = PyDownloader::getRouteFiles(route.canonical_name);
+  const std::string selector = route.selector == LogSelector::RLog ? "r" : route.selector == LogSelector::QLog ? "q" : "a";
+  const std::string result = PyDownloader::resolveRouteFiles(route.canonical_name, route.begin_segment, route.end_segment, selector);
   if (result.empty()) throw std::runtime_error("Failed to fetch route files for " + route.canonical_name);
 
   std::string parse_error;

openpilot/tools/lib/file_downloader.py

@@ -5,6 +5,7 @@ Called by C++ replay/cabana via subprocess.
 
 Subcommands:
   route-files <route>    - Get route file URLs as JSON
+  resolve-route-files <route> - Resolve logs and cameras from API and internal storage
   download <url>         - Download/decompress URL to local cache, print local path
   decompress <path>      - Decompress a local log file, print temporary path
   devices                - List user's devices as JSON
@@ -18,6 +19,8 @@ import os
 import shutil
 import sys
 import tempfile
+import traceback
+from urllib.parse import urlparse
 
 from openpilot.common.hardware.hw import Paths
 
@@ -100,6 +103,100 @@ def cmd_route_files(args):
   api_call(lambda api: api.get(f"v1/route/{args.route}/files"))
 
 
+def resolve_route_files(route_name: str, begin: int = 0, end: int = -1, selector: str = "a") -> dict[int, dict[str, str]]:
+  from openpilot.tools.lib.filereader import DATA_ENDPOINT
+  from openpilot.tools.lib.file_sources import InternalUnavailableException, internal_source
+  from openpilot.tools.lib.route import FileName, SegmentRange
+  from openpilot.tools.lib.api import CommaApi
+  from openpilot.tools.lib.auth_config import get_token
+
+  sr = SegmentRange(route_name)
+  api_error = None
+  try:
+    route_files = CommaApi(get_token()).get(f"v1/route/{sr.route_name}/files")
+  except Exception as e:
+    api_error = e
+    route_files = {}
+  fields = {"rlog": FileName.RLOG, "qlog": FileName.QLOG,
+            "narrow_road": FileName.FCAMERA, "cabin": FileName.DCAMERA, "wide_road": FileName.ECAMERA, "qcamera": FileName.QCAMERA}
+  names = {name: key for key, fns in fields.items() for name in fns}
+  names.update({"rlog": "rlog", "qlog": "qlog"})
+  api_files: dict[int, dict[str, str]] = {}
+  for url in (url for urls in route_files.values() for url in urls):
+    _, number, name = urlparse(url).path.rsplit("/", 2)
+    if name in names:
+      api_files.setdefault(int(number), {})[names[name]] = url
+  # Preserve replay's inclusive segment bounds.
+  open_ended = end < 0
+  local_internal = os.path.isdir(DATA_ENDPOINT)
+  if open_ended:
+    end = max(begin, max(api_files, default=0))
+    if local_internal:
+      internal_dir = os.path.join(DATA_ENDPOINT, sr.dongle_id, sr.log_id)
+      if os.path.isdir(internal_dir):
+        end = max(end, max((int(name) for name in os.listdir(internal_dir) if name.isdigit()), default=0))
+  indices = list(range(max(0, begin), end + 1))
+  segments: dict[int, dict[str, str]] = {seg: {} for seg in indices}
+  internal_errors: dict[str, Exception] = {}
+
+  def resolve(key, requested):
+    files = {seg: api_files[seg][key] for seg in requested if key in api_files.get(seg, {})}
+    missing = [seg for seg in requested if seg not in files]
+    if missing:
+      try:
+        files.update(internal_source(sr, missing, fields[key]))
+      except Exception as e:
+        if e is not InternalUnavailableException:
+          internal_errors[key] = e
+    for seg, path in files.items():
+      segments[seg][key] = path
+
+  log_types = ["qlog"] if selector == "q" else ["rlog"]
+  if selector == "a":
+    log_types.append("qlog")
+  for key in log_types:
+    resolve(key, [seg for seg in indices if not segments[seg]])
+  if open_ended and not local_internal:
+    # Remote internal storage exposes individual files, not a segment listing.
+    # Probe the contiguous tail beyond the API's last segment until no selected log exists.
+    number = end + 1
+    while True:
+      internal_errors.clear()
+      segments[number] = {}
+      for key in log_types:
+        if not segments[number]:
+          resolve(key, [number])
+      if not segments[number]:
+        if internal_errors:
+          print(f"Internal discovery stopped at segment {number}; keeping resolved logs", file=sys.stderr)
+          traceback.print_exception(next(iter(internal_errors.values())))
+        del segments[number]
+        break
+      indices.append(number)
+      number += 1
+  available = [seg for seg in indices if segments[seg]]
+  if not available:
+    if api_error is not None:
+      raise api_error
+    raise FileNotFoundError(f"No {selector} logs found for {sr.route_name} in segments {indices} through the API or internal storage")
+  for key in ("narrow_road", "cabin", "wide_road", "qcamera"):
+    resolve(key, available)
+  return {seg: segments[seg] for seg in available}
+
+
+def cmd_resolve_route_files(args):
+  try:
+    result = resolve_route_files(args.route, args.begin, args.end, args.selector)
+  except Exception as e:
+    error = f"{type(e).__name__} resolving {args.route} (begin={args.begin}, end={args.end}, selector={args.selector}): {e}"
+    print(error, file=sys.stderr)
+    traceback.print_exc()
+    result = {"error": error}
+  json.dump(result, sys.stdout)
+  sys.stdout.write("\n")
+  sys.stdout.flush()
+
+
 def cmd_download(args):
   url = args.url
   use_cache = not args.no_cache
@@ -232,6 +329,13 @@ def main():
   p_rf.add_argument("route")
   p_rf.set_defaults(func=cmd_route_files)
 
+  p_resolve = subparsers.add_parser("resolve-route-files")
+  p_resolve.add_argument("route")
+  p_resolve.add_argument("--begin", type=int, default=0)
+  p_resolve.add_argument("--end", type=int, default=-1, help="Inclusive last segment; negative means open-ended")
+  p_resolve.add_argument("--selector", choices=("a", "r", "q"), default="a")
+  p_resolve.set_defaults(func=cmd_resolve_route_files)
+
   p_dl = subparsers.add_parser("download")
   p_dl.add_argument("url")
   p_dl.add_argument("--no-cache", action="store_true")

openpilot/tools/replay/py_downloader.cc

@@ -229,6 +229,11 @@ std::string getRouteFiles(const std::string &route) {
   return runPython(DOWNLOADER_MODULE, {"route-files", route});
 }
 
+std::string resolveRouteFiles(const std::string &route, int begin, int end, const std::string &selector) {
+  return runPython(DOWNLOADER_MODULE, {"resolve-route-files", route, "--begin", std::to_string(begin),
+                                     "--end", std::to_string(end), "--selector", selector});
+}
+
 std::string authenticate(const std::string &provider, std::atomic<bool> *abort) {
   return runPython(AUTH_MODULE, {provider, "--json"}, abort);
 }

openpilot/tools/replay/py_downloader.h

@@ -18,6 +18,9 @@ std::string decompress(const std::string &path, std::atomic<bool> *abort = nullp
 // Returns JSON string of route files (same format as /v1/route/.../files API)
 std::string getRouteFiles(const std::string &route);
 
+// Resolves API/internal files into a JSON manifest keyed by segment number. End is inclusive.
+std::string resolveRouteFiles(const std::string &route, int begin, int end, const std::string &selector);
+
 // Browser sign-in; abort closes the local callback server. Returns a JSON status.
 std::string authenticate(const std::string &provider, std::atomic<bool> *abort);