#!/usr/bin/env python3
"""
USB camera format and FPS tester for Jetson.

- Lists available formats/resolutions via v4l2-ctl.
- For selected modes, opens the camera with OpenCV and measures per-frame
  capture time (ms) to show real delivered FPS/latency.
"""

import argparse
import subprocess
import sys
import time
import statistics
from typing import List, Tuple

import cv2


def run_v4l2_list_formats(device: str) -> str:
    """Return the raw output of v4l2-ctl --list-formats-ext."""
    try:
        out = subprocess.check_output(
            ["v4l2-ctl", f"--device={device}", "--list-formats-ext"],
            stderr=subprocess.STDOUT,
            text=True,
        )
        return out
    except (subprocess.CalledProcessError, FileNotFoundError) as e:
        print(f"[WARN] Failed to run v4l2-ctl: {e}")
        return ""


def parse_formats(output: str) -> List[Tuple[str, int, int, List[float]]]:
    """
    Parse v4l2-ctl --list-formats-ext output.

    Returns list of (fourcc, width, height, fps_list).
    """
    lines = output.splitlines()
    formats: List[Tuple[str, int, int, List[float]]] = []

    current_fourcc = None
    for line in lines:
        line = line.strip()
        if line.startswith("[") and "):" in line and "'" in line:
            # Example: [0]: 'MJPG' (Motion-JPEG, compressed)
            try:
                fourcc = line.split("'")[1]
                current_fourcc = fourcc
            except Exception:
                current_fourcc = None
        elif line.startswith("Size: Discrete") and current_fourcc:
            # Example: Size: Discrete 1280x720
            try:
                parts = line.split()
                wh = parts[-1]
                w_str, h_str = wh.split("x")
                width = int(w_str)
                height = int(h_str)
                formats.append((current_fourcc, width, height, []))
            except Exception:
                pass
        elif "Interval:" in line and "fps" in line and formats:
            # Example: Interval: Discrete 0.033s (30.000 fps)
            try:
                fps_part = line.split("(")[-1].split("fps")[0]
                fps = float(fps_part)
                fourcc, w, h, fps_list = formats[-1]
                fps_list.append(fps)
                formats[-1] = (fourcc, w, h, fps_list)
            except Exception:
                pass

    return formats


def fourcc_to_int(code: str) -> int:
    return cv2.VideoWriter_fourcc(*code)


def benchmark_mode(
    device_index: int,
    fourcc: str,
    width: int,
    height: int,
    target_fps: float,
    warmup: int = 10,
    measure_frames: int = 60,
) -> None:
    """Open a specific mode and measure per-frame capture time."""
    print(f"\n=== Testing {fourcc} {width}x{height} @ {target_fps:.1f} fps ===")

    cap = cv2.VideoCapture(device_index, cv2.CAP_V4L2)
    if not cap.isOpened():
        print("  [ERROR] Failed to open device with V4L2 backend.")
        return

    cap.set(cv2.CAP_PROP_FOURCC, fourcc_to_int(fourcc))
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
    cap.set(cv2.CAP_PROP_FPS, target_fps)

    # Query negotiated parameters
    negotiated_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    negotiated_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    negotiated_fps = cap.get(cv2.CAP_PROP_FPS)
    negotiated_fourcc_val = int(cap.get(cv2.CAP_PROP_FOURCC))
    negotiated_fourcc = "".join(
        chr((negotiated_fourcc_val >> 8 * i) & 0xFF) for i in range(4)
    )
    print(
        f"  Negotiated: {negotiated_w}x{negotiated_h} {negotiated_fourcc} "
        f"@ {negotiated_fps:.1f} fps (backend={cap.getBackendName()})"
    )

    # Warm-up frames (to let auto-exposure, buffers, etc. settle)
    for i in range(warmup):
        ret, _ = cap.read()
        if not ret:
            print(f"  [WARN] Warmup frame {i} failed.")
            time.sleep(0.05)

    times_ms: List[float] = []
    for i in range(measure_frames):
        t0 = time.perf_counter()
        ret, _ = cap.read()
        t1 = time.perf_counter()
        if not ret:
            print(f"  [WARN] Capture failed at frame {i}.")
            break
        times_ms.append((t1 - t0) * 1000.0)

    cap.release()

    if not times_ms:
        print("  [ERROR] No successful frames captured.")
        return

    mean_ms = statistics.mean(times_ms)
    med_ms = statistics.median(times_ms)
    min_ms = min(times_ms)
    max_ms = max(times_ms)
    eff_fps = 1000.0 / mean_ms

    print(
        f"  Frames measured: {len(times_ms)}\n"
        f"  Capture time (ms): mean={mean_ms:.2f}, median={med_ms:.2f}, "
        f"min={min_ms:.2f}, max={max_ms:.2f}\n"
        f"  Effective FPS from capture timing: {eff_fps:.2f}"
    )


def main() -> None:
    parser = argparse.ArgumentParser(
        description="List USB camera formats and benchmark capture times."
    )
    parser.add_argument(
        "--device",
        default="/dev/video4",
        help="V4L2 device path (default: /dev/video0)",
    )
    parser.add_argument(
        "--index",
        type=int,
        default=0,
        help="OpenCV device index (default: 0)",
    )
    parser.add_argument(
        "--test-modes",
        nargs="*",
        default=["MJPG:640x480:30", "MJPG:1280x720:30"],
        help=(
            "Modes to benchmark as FOURCC:WxH:FPS. "
            "Example: MJPG:640x480:30 YUYV:640x480:30"
        ),
    )
    args = parser.parse_args()

    print(f"== v4l2-ctl formats for {args.device} ==")
    raw = run_v4l2_list_formats(args.device)
    if not raw:
        print("No v4l2-ctl output; skipping format list.")
    else:
        print(raw)
        parsed = parse_formats(raw)
        if parsed:
            print("\n== Parsed formats (FOURCC WxH [fps...]) ==")
            for fourcc, w, h, fps_list in parsed:
                fps_str = ", ".join(f"{f:.1f}" for f in fps_list) if fps_list else "?"
                print(f"  {fourcc} {w}x{h} [{fps_str}]")

    print("\n== Benchmarking capture time with OpenCV ==")
    for mode_str in args.test_modes:
        try:
            fourcc_s, wh_s, fps_s = mode_str.split(":")
            w_s, h_s = wh_s.lower().split("x")
            w = int(w_s)
            h = int(h_s)
            fps = float(fps_s)
        except ValueError:
            print(f"[WARN] Cannot parse test mode '{mode_str}', expected FOURCC:WxH:FPS")
            continue

        benchmark_mode(args.index, fourcc_s, w, h, fps)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nInterrupted by user.", file=sys.stderr)