import argparse
import av
import cv2
import platform
import subprocess
import sys
import time

TARGET_FPS = 30
WIDTH = 1920
HEIGHT = 1080


def list_windows_dshow_devices():
    """
    Print DirectShow devices using ffmpeg, if available.
    """
    cmd = ["ffmpeg", "-hide_banner", "-list_devices", "true", "-f", "dshow", "-i", "dummy"]
    print("Listing Windows DirectShow devices via ffmpeg...")
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=False)
    except FileNotFoundError as exc:
        raise RuntimeError("ffmpeg was not found in PATH; cannot list DirectShow devices.") from exc

    output = (result.stderr or "") + (result.stdout or "")
    print(output.strip())


def open_pyav_capture(system_name, device, pixel_format):
    """
    Open a camera device via PyAV on Linux or Windows.
    """
    if system_name == "linux":
        options = {
            "video_size": f"{WIDTH}x{HEIGHT}",
            "framerate": str(TARGET_FPS),
            "input_format": pixel_format or "mjpeg",
        }
        container = av.open(device or "/dev/video4", format="v4l2", options=options)
    elif system_name == "windows":
        if not device:
            raise RuntimeError(
                "On Windows, PyAV/dshow requires a camera name, for example "
                '--device "USB Camera". Use --list-devices first.'
            )
        options = {
            "video_size": f"{WIDTH}x{HEIGHT}",
            "framerate": str(TARGET_FPS),
        }
        if pixel_format:
            options["pixel_format"] = pixel_format
        container = av.open(f"video={device}", format="dshow", options=options)
    else:
        raise RuntimeError(
            f"Unsupported OS: {platform.system()}. Only Linux and Windows are supported."
        )

    stream = container.streams.video[0]
    return container, stream


def main():
    parser = argparse.ArgumentParser(
        description="Capture from camera via PyAV on Linux or Windows, optionally save to MP4."
    )
    parser.add_argument(
        "-d",
        "--device",
        type=str,
        default="",
        help=(
            "Linux: V4L2 device path like /dev/video4. "
            "Windows: DirectShow camera name like 'USB Camera'. "
            "If empty, Linux uses /dev/video4; Windows requires --device."
        ),
    )
    parser.add_argument(
        "--pixel-format",
        type=str,
        default="",
        help=(
            "Optional input pixel format / format hint. "
            "Examples: Linux 'mjpeg'; Windows may accept values such as 'mjpeg' if supported."
        ),
    )
    parser.add_argument(
        "--list-devices",
        action="store_true",
        help="On Windows, list DirectShow device names via ffmpeg and exit.",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        default="",
        help="Output MP4 path (e.g. capture.mp4). If empty, only preview, no file.",
    )
    args = parser.parse_args()

    system_name = platform.system().lower()
    if args.list_devices:
        if system_name != "windows":
            raise RuntimeError("--list-devices is currently only implemented for Windows.")
        list_windows_dshow_devices()
        return

    container, stream = open_pyav_capture(system_name, args.device, args.pixel_format)
    print(
        f"Using PyAV backend on {platform.system()} "
        f"with device={args.device or '/dev/video4'}"
    )

    frame_count = 0
    t_start = time.perf_counter()

    writer = None
    if args.output:
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        writer = cv2.VideoWriter(
            args.output, fourcc, float(TARGET_FPS), (WIDTH, HEIGHT)
        )
        if not writer.isOpened():
            container.close()
            raise RuntimeError(f"Could not open VideoWriter for {args.output}")
        print(f"Recording to {args.output} @ {TARGET_FPS} fps, {WIDTH}x{HEIGHT}")

    try:
        for frame in container.decode(stream):
            t0 = time.perf_counter()
            img = frame.to_ndarray(format="bgr24")  # OpenCV-compatible
            t1 = time.perf_counter()

            if writer is not None:
                writer.write(img)

            cv2.imshow("pyav_cam", img)
            key = cv2.waitKey(1) & 0xFF
            t2 = time.perf_counter()

            frame_count += 1
            if frame_count % 30 == 0:
                elapsed = t2 - t_start
                print(
                    f"Frame {frame_count}: convert={1000*(t1-t0):.1f}ms, "
                    f"imshow+waitKey={1000*(t2-t1):.1f}ms, avg_fps={frame_count/elapsed:.1f}"
                )

            if key == ord("q"):
                break
    finally:
        container.close()
        if writer is not None:
            writer.release()
            print(f"Saved {frame_count} frames to {args.output}")
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()