#!/usr/bin/env python3
import os
import argparse
import coverage
import pyray as rl

from tqdm import tqdm
from typing import Literal
from collections.abc import Callable
from cereal.messaging import PubMaster
from openpilot.common.api import Api
from openpilot.common.basedir import BASEDIR
from openpilot.common.params import Params
from openpilot.common.prefix import OpenpilotPrefix
from openpilot.selfdrive.ui.tests.diff.diff import DIFF_OUT_DIR
from openpilot.system.updated.updated import parse_release_notes
from openpilot.system.version import terms_version, training_version

LayoutVariant = Literal["mici", "tizi"]

FPS = 60
HEADLESS = os.getenv("WINDOWED", "0") != "1"


def setup_state():
  params = Params()
  params.put("HasAcceptedTerms", terms_version, block=True)
  params.put("CompletedTrainingVersion", training_version, block=True)
  params.put("DongleId", "test123456789", block=True)
  # Combined description for layouts that still use it (BIG home, settings/software)
  params.put("UpdaterCurrentDescription", "0.10.1 / test-branch / abc1234 / Nov 30", block=True)
  params.put("UpdaterCurrentReleaseNotes", parse_release_notes(BASEDIR), block=True)
  # Params for mici home
  params.put("Version", "0.10.1", block=True)
  params.put("GitBranch", "test-branch", block=True)
  params.put("GitCommit", "abc12340ff9131237ba23a1d0fbd8edf9c80e87", block=True)
  params.put("GitCommitDate", "'1732924800 2024-11-30 00:00:00 +0000'", block=True)

  # Patch Api.get_token to return a static token so the pairing QR code is deterministic across runs
  Api.get_token = lambda self, payload_extra=None, expiry_hours=0: "test_token"


def run_replay(variant: LayoutVariant) -> None:
  if HEADLESS:
    rl.set_config_flags(rl.ConfigFlags.FLAG_WINDOW_HIDDEN)
    os.environ["OFFSCREEN"] = "1"  # Run UI without FPS limit (set before importing gui_app)

  setup_state()
  os.makedirs(DIFF_OUT_DIR, exist_ok=True)

  from openpilot.selfdrive.ui.ui_state import ui_state, device  # Import within OpenpilotPrefix context so param values are setup correctly
  from openpilot.system.ui.lib.application import gui_app  # Import here for accurate coverage
  from openpilot.selfdrive.ui.tests.diff.replay_script import build_script

  gui_app.init_window("ui diff test", fps=FPS)

  # Dynamically import main layout based on variant
  if variant == "mici":
    from openpilot.selfdrive.ui.mici.layouts.main import MiciMainLayout as MainLayout
  else:
    from openpilot.selfdrive.ui.layouts.main import MainLayout
  main_layout = MainLayout()

  # Disable interactive timeout — replay clicks use left_down=False so they never reset the timer,
  # and after 30s of real wall-clock time the settings panel would close automatically.
  device.set_override_interactive_timeout(99999)

  pm = PubMaster(["deviceState", "pandaStates", "driverStateV2", "selfdriveState"])
  script = build_script(pm, main_layout, variant)
  script_index = 0

  send_fn: Callable | None = None
  frame = 0
  # Override raylib timing functions to return deterministic values based on frame count instead of real time
  rl.get_frame_time = lambda: 1.0 / FPS
  rl.get_time = lambda: frame / FPS

  # Main loop to replay events and render frames
  with tqdm(total=script[-1][0] + 1, desc="Replaying", unit="frame", disable=bool(os.getenv("CI"))) as pbar:
    for _ in gui_app.render():
      # Handle all events for the current frame
      while script_index < len(script) and script[script_index][0] == frame:
        _, event = script[script_index]
        # Call setup function, if any
        if event.setup:
          event.setup()
        # Send mouse events to the application
        if event.mouse_events:
          with gui_app._mouse._lock:
            gui_app._mouse._events.extend(event.mouse_events)
        # Update persistent send function
        if event.send_fn is not None:
          send_fn = event.send_fn
        # Move to next script event
        script_index += 1

      # Keep sending cereal messages for persistent states (onroad, alerts)
      if send_fn:
        send_fn()

      ui_state.update()

      frame += 1
      pbar.update(1)

      if script_index >= len(script):
        break

  gui_app.close()

  print(f"Total frames: {frame}")
  print(f"Video saved to: {os.environ['RECORD_OUTPUT']}")


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('--big', action='store_true', help='Use big UI layout (tizi/tici) instead of mici layout')
  args = parser.parse_args()

  variant: LayoutVariant = 'tizi' if args.big else 'mici'

  if args.big:
    os.environ["BIG"] = "1"
  os.environ["RECORD"] = "1"
  os.environ["RECORD_QUALITY"] = "0"  # Use CRF 0 ("lossless" encode) for deterministic output across different machines
  os.environ["RECORD_OUTPUT"] = os.path.join(DIFF_OUT_DIR, os.environ.get("RECORD_OUTPUT", f"{variant}_ui_replay.mp4"))

  print(f"Running {variant} UI replay...")
  with OpenpilotPrefix():
    sources = ["openpilot.system.ui"]
    if variant == "mici":
      sources.append("openpilot.selfdrive.ui.mici")
      omit = ["**/*tizi*", "**/*tici*"]  # exclude files containing "tizi" or "tici"
    else:
      sources.extend(["openpilot.selfdrive.ui.layouts", "openpilot.selfdrive.ui.onroad", "openpilot.selfdrive.ui.widgets"])
      omit = ["**/*mici*"]  # exclude files containing "mici"
    cov = coverage.Coverage(source=sources, omit=omit)
    with cov.collect():
      run_replay(variant)
    cov.save()
    cov.report()
    directory = os.path.join(DIFF_OUT_DIR, f"htmlcov-{variant}")
    cov.html_report(directory=directory)
    print(f"HTML report: {directory}/index.html")


if __name__ == "__main__":
  main()