#!/usr/bin/env python3
import os
import sys
import time
import logging
import subprocess
import threading
import queue
import multiprocessing
import itertools
import numpy as np
import tqdm
from argparse import ArgumentParser
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

from openpilot.tools.lib.route import Route
from openpilot.tools.lib.logreader import LogReader
from openpilot.tools.lib.filereader import FileReader
from openpilot.tools.lib.framereader import FrameReader, ffprobe
from openpilot.selfdrive.test.process_replay.migration import migrate_all
from openpilot.common.prefix import OpenpilotPrefix
from openpilot.common.utils import Timer
from msgq.visionipc import VisionIpcServer, VisionStreamType

FRAMERATE = 20
DEMO_ROUTE, DEMO_START, DEMO_END = 'a2a0ccea32023010/2023-07-27--13-01-19', 90, 105

logger = logging.getLogger('clip')


def parse_args():
  parser = ArgumentParser(description="Direct clip renderer")
  parser.add_argument("route", nargs="?", help="Route ID (dongle/route or dongle/route/start/end)")
  parser.add_argument("-s", "--start", type=int, help="Start time in seconds")
  parser.add_argument("-e", "--end", type=int, help="End time in seconds")
  parser.add_argument("-o", "--output", default="output.mp4", help="Output file path")
  parser.add_argument("-d", "--data-dir", help="Local directory with route data")
  parser.add_argument("-t", "--title", help="Title overlay text")
  parser.add_argument("-f", "--file-size", type=float, default=9.0, help="Target file size in MB")
  parser.add_argument("-x", "--speed", type=int, default=1, help="Speed multiplier")
  parser.add_argument("--demo", action="store_true", help="Use demo route with default timing")
  ui_group = parser.add_mutually_exclusive_group()
  ui_group.add_argument("--big", dest="big", action="store_true", default=None, help="Force big UI (2160x1080)")
  ui_group.add_argument("--mici", dest="big", action="store_false", help="Force mici UI (536x240)")
  parser.add_argument("--qcam", action="store_true", help="Use qcamera instead of fcamera")
  parser.add_argument("--windowed", action="store_true", help="Show window")
  parser.add_argument("--no-metadata", action="store_true", help="Disable metadata overlay")
  parser.add_argument("--no-time-overlay", action="store_true", help="Disable time overlay")
  args = parser.parse_args()

  if args.demo:
    args.route, args.start, args.end = args.route or DEMO_ROUTE, args.start or DEMO_START, args.end or DEMO_END
  elif not args.route:
    parser.error("route is required (or use --demo)")

  if args.route and args.route.count('/') == 3:
    parts = args.route.split('/')
    args.route, args.start, args.end = '/'.join(parts[:2]), args.start or int(parts[2]), args.end or int(parts[3])

  if args.start is None or args.end is None:
    parser.error("--start and --end are required")
  if args.end <= args.start:
    parser.error(f"end ({args.end}) must be greater than start ({args.start})")
  return args


def setup_env(output_path: str, big: bool = False, speed: int = 1, target_mb: float = 0, duration: int = 0):
  os.environ.update({"RECORD": "1", "OFFSCREEN": "1", "RECORD_OUTPUT": str(Path(output_path).with_suffix(".mp4"))})
  if speed > 1:
    os.environ["RECORD_SPEED"] = str(speed)
  if target_mb > 0 and duration > 0:
    os.environ["RECORD_BITRATE"] = f"{int(target_mb * 8 * 1024 / (duration / speed))}k"
  if big:
    os.environ["BIG"] = "1"
  else:
    os.environ["BIG"] = "0"


def _download_segment(path: str) -> bytes:
  with FileReader(path) as f:
    return bytes(f.read())


def _parse_and_chunk_segment(args: tuple) -> list[dict]:
  raw_data, fps = args
  from openpilot.tools.lib.logreader import _LogFileReader
  messages = migrate_all(list(_LogFileReader("", dat=raw_data, sort_by_time=True)))
  if not messages:
    return []

  dt_ns, chunks, current, next_time = 1e9 / fps, [], {}, messages[0].logMonoTime + 1e9 / fps  # type: ignore[var-annotated]
  for msg in messages:
    if msg.logMonoTime >= next_time:
      chunks.append(current)
      current, next_time = {}, next_time + dt_ns * ((msg.logMonoTime - next_time) // dt_ns + 1)
    current[msg.which()] = msg
  return chunks + [current] if current else chunks


def load_logs_parallel(log_paths: list[str], fps: int = 20) -> list[dict]:
  num_workers = min(16, len(log_paths), (multiprocessing.cpu_count() or 1))
  logger.info(f"Downloading {len(log_paths)} segments with {num_workers} workers...")

  with ThreadPoolExecutor(max_workers=num_workers) as pool:
    futures = {pool.submit(_download_segment, path): idx for idx, path in enumerate(log_paths)}
    raw_data = {futures[f]: f.result() for f in as_completed(futures)}

  logger.info("Parsing and chunking segments...")
  with multiprocessing.Pool(num_workers) as pool:
    return list(itertools.chain.from_iterable(pool.map(_parse_and_chunk_segment, [(raw_data[i], fps) for i in range(len(log_paths))])))


def patch_submaster(message_chunks, ui_state):
  # Reset started_frame so alerts render correctly (recv_frame must be >= started_frame)
  ui_state.started_frame = 0
  ui_state.started_time = time.monotonic()

  def mock_update(timeout=None):
    sm, t = ui_state.sm, time.monotonic()
    sm.updated = dict.fromkeys(sm.services, False)
    if sm.frame < len(message_chunks):
      for svc, msg in message_chunks[sm.frame].items():
        if svc in sm.data:
          sm.seen[svc] = sm.updated[svc] = sm.alive[svc] = sm.valid[svc] = True
          sm.data[svc] = getattr(msg.as_builder(), svc)
          sm.logMonoTime[svc], sm.recv_time[svc], sm.recv_frame[svc] = msg.logMonoTime, t, sm.frame
    sm.frame += 1
  ui_state.sm.update = mock_update


def get_frame_dimensions(camera_path: str) -> tuple[int, int]:
  """Get frame dimensions from a video file using ffprobe."""
  probe = ffprobe(camera_path)
  stream = probe["streams"][0]
  return stream["width"], stream["height"]


def iter_segment_frames(camera_paths, start_time, end_time, fps=20, use_qcam=False,
                        frame_size: tuple[int, int] | None = None, on_segment_open=None):
  frames_per_seg = fps * 60
  start_frame, end_frame = int(start_time * fps), int(end_time * fps)
  current_seg: int = -1
  seg_frames: FrameReader | np.ndarray | None = None

  for global_idx in range(start_frame, end_frame):
    seg_idx, local_idx = global_idx // frames_per_seg, global_idx % frames_per_seg

    if seg_idx != current_seg:
      current_seg = seg_idx
      path = camera_paths[seg_idx] if seg_idx < len(camera_paths) else None
      if not path:
        raise RuntimeError(f"No camera file for segment {seg_idx}")
      if on_segment_open is not None:
        on_segment_open(seg_idx, path)

      if use_qcam:
        w, h = frame_size or get_frame_dimensions(path)
        with FileReader(path) as f:
          result = subprocess.run(["ffmpeg", "-v", "quiet", "-i", "-", "-f", "rawvideo", "-pix_fmt", "nv12", "-"],
                                  input=f.read(), capture_output=True)
        if result.returncode != 0:
          raise RuntimeError(f"ffmpeg failed: {result.stderr.decode()}")
        seg_frames = np.frombuffer(result.stdout, dtype=np.uint8).reshape(-1, w * h * 3 // 2)
      else:
        seg_frames = FrameReader(path, pix_fmt="nv12")

    assert seg_frames is not None
    frame = seg_frames[local_idx] if use_qcam else seg_frames.get(local_idx)  # type: ignore[index, union-attr]
    yield global_idx, frame


class FrameQueue:
  def __init__(self, camera_paths, start_time, end_time, fps=20, prefetch_count=60, use_qcam=False):
    # Probe first valid camera file for dimensions
    first_path = next((p for p in camera_paths if p), None)
    if not first_path:
      raise RuntimeError("No valid camera paths")
    self.frame_w, self.frame_h = get_frame_dimensions(first_path)

    self._queue, self._stop, self._error = queue.Queue(maxsize=prefetch_count), threading.Event(), None
    self._use_qcam = use_qcam
    self._current_seg_idx = None
    self._current_path = first_path
    self._state_lock = threading.Lock()
    self._thread = threading.Thread(target=self._worker,
                                    args=(camera_paths, start_time, end_time, fps, use_qcam, (self.frame_w, self.frame_h)), daemon=True)
    self._thread.start()

  def _set_current_source(self, seg_idx, path):
    with self._state_lock:
      self._current_seg_idx = seg_idx
      self._current_path = path

  def _worker(self, camera_paths, start_time, end_time, fps, use_qcam, frame_size):
    try:
      for idx, data in iter_segment_frames(camera_paths, start_time, end_time, fps, use_qcam, frame_size, self._set_current_source):
        if self._stop.is_set():
          break
        self._queue.put((idx, data.tobytes()))
    except Exception as e:
      logger.exception("Decode error")
      self._error = e
    finally:
      self._queue.put(None)

  def get(self, timeout=60.0):
    deadline = time.monotonic() + timeout
    while True:
      if self._error:
        raise self._error

      remaining = max(0.0, deadline - time.monotonic())
      if remaining == 0.0:
        break

      try:
        result = self._queue.get(timeout=min(0.5, remaining))
      except queue.Empty:
        continue

      if result is None:
        if self._error:
          raise self._error
        raise StopIteration("No more frames")
      return result

    if self._error:
      raise self._error

    with self._state_lock:
      seg_idx = self._current_seg_idx
      path = self._current_path

    camera_kind = "qcamera" if self._use_qcam else "fcamera"
    source = f"segment {seg_idx}" if seg_idx is not None else "the initial segment"
    if path:
      source = f"{source} ({path})"
    hint = ""
    if path and path.startswith(("http://", "https://")):
      hint = " Try downloading the route locally with --data-dir or verify the remote camera endpoint supports timely range reads."
    raise TimeoutError(f"Timed out after {timeout:.0f}s waiting for {camera_kind} frames from {source}; camera fetch or decode is stalled.{hint}")

  def stop(self):
    self._stop.set()
    while not self._queue.empty():
      try:
        self._queue.get_nowait()
      except queue.Empty:
        break
    self._thread.join(timeout=2.0)


def load_route_metadata(route):
  from openpilot.common.params import Params, UnknownKeyName
  lr = LogReader(route.log_paths()[0])
  init_data, car_params = lr.first('initData'), lr.first('carParams')

  params = Params()
  for entry in init_data.params.entries:
    try:
      value = params.cpp2python(entry.key, entry.value)
      if value is None:
        logger.warning("Skipping malformed route param %s while loading clip metadata", entry.key)
        continue
      params.put(entry.key, value)
    except UnknownKeyName:
      pass
    except TypeError:
      logger.warning("Skipping route param %s due to type mismatch while loading clip metadata", entry.key)

  origin = init_data.gitRemote.split('/')[3] if len(init_data.gitRemote.split('/')) > 3 else 'unknown'
  return {
    'version': init_data.version, 'route': route.name.canonical_name,
    'car': car_params.carFingerprint if car_params else 'unknown', 'origin': origin,
    'commit': init_data.gitCommit[:7],
  }


def detect_big_ui(route: Route) -> bool:
  try:
    init_data = LogReader(route.log_paths()[0]).first('initData')
    git_branch = (init_data.gitBranch or "").lower()
    device_type = str(init_data.deviceType).lower()
    if device_type in ("mici", "tizi", "tici"):
      big = device_type != "mici"
      reason = f"route device type {device_type}"
    elif "mici" in git_branch:
      big = False
      reason = f"route branch {git_branch}"
    else:
      big = True
      reason = f"route branch {git_branch or 'unknown'}"
    logger.info("Auto-detected %s UI from %s", "big" if big else "mici", reason)
    return big
  except Exception:
    logger.warning("Falling back to big UI; failed to auto-detect UI mode", exc_info=True)
    return True


def draw_text_box(rl, text, x, y, size, gui_app, font, font_scale, color=None, center=False):
  box_color, text_color = rl.Color(0, 0, 0, 85), color or rl.WHITE
  # measure_text_ex is NOT auto-scaled, so multiply by font_scale
  # draw_text_ex IS auto-scaled, so pass size directly
  text_size = rl.measure_text_ex(font, text, size * font_scale, 0)
  text_width, text_height = int(text_size.x), int(text_size.y)
  if center:
    x = (gui_app.width - text_width) // 2
  rl.draw_rectangle(x - 8, y - 4, text_width + 16, text_height + 8, box_color)
  rl.draw_text_ex(font, text, rl.Vector2(x, y), size, 0, text_color)


def render_overlays(rl, gui_app, font, font_scale, metadata, title, start_time, frame_idx, show_metadata, show_time):
  if show_metadata and metadata and frame_idx < FRAMERATE * 5:
    m = metadata
    text = ", ".join([f"IQ.Pilot v{m['version']}", f"route: {m['route']}", f"car: {m['car']}", f"origin: {m['origin']}",
                      f"commit: {m['commit']}"])
    # Truncate if too wide (leave 20px margin on each side)
    max_width = gui_app.width - 40
    while rl.measure_text_ex(font, text, 15 * font_scale, 0).x > max_width and len(text) > 20:
      text = text[:-4] + "..."
    draw_text_box(rl, text, 0, 8, 15, gui_app, font, font_scale, center=True)

  if title:
    draw_text_box(rl, title, 0, 60, 32, gui_app, font, font_scale, center=True)

  if show_time:
    t = start_time + frame_idx / FRAMERATE
    time_text = f"{int(t)//60:02d}:{int(t)%60:02d}"
    time_width = int(rl.measure_text_ex(font, time_text, 24 * font_scale, 0).x)
    draw_text_box(rl, time_text, gui_app.width - time_width - 45, 45, 24, gui_app, font, font_scale)


def prefetch_nav_tiles(road_view, ui_state, message_chunks) -> None:
  """If the route drove with on-screen maps enabled, feed the panel the first messages and block
  until the opening viewport's tiles are fetched, so the clip doesn't start on the placeholder grid."""
  nav_panel = getattr(getattr(road_view, "_hud_renderer", None), "nav_map_panel", None)
  if nav_panel is None or not getattr(nav_panel, "maps_enabled", lambda: False)():
    return

  logger.info("Route has on-screen maps enabled, prefetching map tiles...")
  for _ in range(min(len(message_chunks), FRAMERATE * 2)):
    ui_state.sm.update()
    nav_panel.update()
    if nav_panel.active:
      break
  if nav_panel.active:
    if not nav_panel.warm_up_tiles():
      logger.warning("Map tiles incomplete after warmup (missing MapboxToken or slow network); map may render partially")
  else:
    logger.warning("No nav position in the first seconds of the clip; map tiles will load mid-clip")
  ui_state.sm.frame = 0


def clip(route: Route, output: str, start: int, end: int, headless: bool = True, big: bool = False,
         title: str | None = None, show_metadata: bool = True, show_time: bool = True, use_qcam: bool = False):
  timer, duration = Timer(), end - start

  # The prefix must wrap the UI imports: ui_state and the nav map panel bind Params() to the
  # params path active when they're constructed, and load_route_metadata below seeds the route's
  # params (on-screen maps gate, units, Mapbox token) into the prefixed dir.
  with OpenpilotPrefix(shared_download_cache=True):
    import pyray as rl
    if big:
      from openpilot.selfdrive.ui.onroad.augmented_road_view import AugmentedRoadView
    else:
      from openpilot.selfdrive.ui.mici.onroad.augmented_road_view import AugmentedRoadView  # type: ignore[assignment]
    from openpilot.selfdrive.ui.ui_state import ui_state
    from openpilot.system.ui.lib.application import gui_app, FontWeight, FONT_SCALE
    timer.lap("import")

    logger.info(f"Clipping {route.name.canonical_name}, {start}s-{end}s ({duration}s)")
    seg_start, seg_end = start // 60, (end - 1) // 60 + 1
    all_chunks = load_logs_parallel(route.log_paths()[seg_start:seg_end], fps=FRAMERATE)
    timer.lap("logs")

    frame_start = (start - seg_start * 60) * FRAMERATE
    message_chunks = all_chunks[frame_start:frame_start + duration * FRAMERATE]
    if not message_chunks:
      logger.error("No messages to render")
      sys.exit(1)

    metadata = load_route_metadata(route)
    if not show_metadata:
      metadata = None
    if headless:
      rl.set_config_flags(rl.ConfigFlags.FLAG_WINDOW_HIDDEN)

    camera_paths = route.qcamera_paths() if use_qcam else route.camera_paths()
    frame_queue = FrameQueue(camera_paths, start, end, fps=FRAMERATE, use_qcam=use_qcam)

    ecamera_paths = route.ecamera_paths() if not use_qcam else []
    wide_frame_queue: FrameQueue | None = None
    if any(p for p in ecamera_paths[seg_start:seg_end] if p):
      wide_frame_queue = FrameQueue(ecamera_paths, start, end, fps=FRAMERATE)

    vipc = VisionIpcServer("camerad")
    vipc.create_buffers(VisionStreamType.VISION_STREAM_ROAD, 4, frame_queue.frame_w, frame_queue.frame_h)
    if wide_frame_queue:
      vipc.create_buffers(VisionStreamType.VISION_STREAM_WIDE_ROAD, 4, wide_frame_queue.frame_w, wide_frame_queue.frame_h)
    vipc.start_listener()

    patch_submaster(message_chunks, ui_state)
    gui_app.init_window("clip", fps=FRAMERATE)

    road_view = AugmentedRoadView()
    road_view.set_rect(rl.Rectangle(0, 0, gui_app.width, gui_app.height))
    font = gui_app.font(FontWeight.NORMAL)
    prefetch_nav_tiles(road_view, ui_state, message_chunks)
    timer.lap("setup")

    frame_idx = 0
    with tqdm.tqdm(total=len(message_chunks), desc="Rendering", unit="frame") as pbar:
      for should_render in gui_app.render():
        if frame_idx >= len(message_chunks):
          break
        _, frame_bytes = frame_queue.get()
        vipc.send(VisionStreamType.VISION_STREAM_ROAD, frame_bytes, frame_idx, int(frame_idx * 5e7), int(frame_idx * 5e7))
        if wide_frame_queue:
          _, wide_bytes = wide_frame_queue.get()
          vipc.send(VisionStreamType.VISION_STREAM_WIDE_ROAD, wide_bytes, frame_idx, int(frame_idx * 5e7), int(frame_idx * 5e7))
        ui_state.update()
        if should_render:
          road_view.render()
          render_overlays(rl, gui_app, font, FONT_SCALE, metadata, title, start, frame_idx, show_metadata, show_time)
        frame_idx += 1
        pbar.update(1)
    timer.lap("render")

    frame_queue.stop()
    if wide_frame_queue:
      wide_frame_queue.stop()
    gui_app.close()
    timer.lap("ffmpeg")

  logger.info(f"Clip saved to: {Path(output).resolve()}")
  logger.info(f"Generated {timer.fmt(duration)}")


def main():
  logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s\t%(message)s")
  args = parse_args()
  route = Route(args.route, data_dir=args.data_dir)
  big = args.big if args.big is not None else detect_big_ui(route)

  setup_env(args.output, big=big, speed=args.speed, target_mb=args.file_size, duration=args.end - args.start)
  try:
    clip(route, args.output, args.start, args.end, not args.windowed,
         big, args.title, not args.no_metadata, not args.no_time_overlay, args.qcam)
  except TimeoutError as e:
    logger.error("%s", e)
    sys.exit(1)


if __name__ == "__main__":
  main()