#!/usr/bin/env python3
"""Prepare one EasyCom session for causal Qwen2.5-Omni extraction.

The released EasyCom session consists of retained, approximately one-minute
MP4 clips.  EgoSpeak removes source-timeline gaps by concatenating those clips
and crops the left participant-ID banner.  This script reproduces that dense
timeline and writes fixed-global-rate 2 FPS and/or 5 FPS AV files.

It intentionally never fills gaps inferred from clip filenames.  The number
of 5 Hz task steps is derived from ``target_perframe`` (preferred) or from the
sum of the released 20 FPS video frames.
"""

from __future__ import annotations

import argparse
import json
import math
import shlex
import subprocess
from dataclasses import asdict, dataclass
from fractions import Fraction
from pathlib import Path

from easycom_manifest import parse_timecode


SOURCE_FPS = 20
TASK_FPS = 5
KEEP_WIDTH = 1920
KEEP_HEIGHT = 1080


@dataclass(frozen=True)
class ClipInfo:
    path: str
    timecode_seconds: float
    width: int
    height: int
    fps: float
    frames: int


def _run_json(command: list[str]) -> dict:
    try:
        completed = subprocess.run(
            command, check=True, capture_output=True, text=True
        )
    except subprocess.CalledProcessError as exc:
        stderr = (exc.stderr or "").strip()
        detail = f"\n{stderr}" if stderr else ""
        raise RuntimeError(
            f"Command failed ({exc.returncode}): {shlex.join(command)}{detail}"
        ) from exc
    try:
        return json.loads(completed.stdout)
    except json.JSONDecodeError as exc:
        raise RuntimeError(
            f"Command returned invalid JSON: {shlex.join(command)}"
        ) from exc


def probe_clip(path: Path) -> ClipInfo:
    payload = _run_json(
        [
            "ffprobe",
            "-v",
            "error",
            "-select_streams",
            "v:0",
            "-show_entries",
            "stream=width,height,avg_frame_rate,nb_frames",
            "-of",
            "json",
            str(path),
        ]
    )
    streams = payload.get("streams", [])
    if not streams:
        raise RuntimeError(f"No video stream in {path}")
    stream = streams[0]
    if stream.get("nb_frames") in (None, "N/A"):
        raise RuntimeError(f"Exact nb_frames is required for {path}")
    fps = float(Fraction(stream["avg_frame_rate"]))
    return ClipInfo(
        path=str(path.resolve()),
        timecode_seconds=parse_timecode(path.stem),
        width=int(stream["width"]),
        height=int(stream["height"]),
        fps=fps,
        frames=int(stream["nb_frames"]),
    )


def load_target_steps(path: Path) -> int:
    try:
        import numpy as np
    except ImportError as exc:
        raise RuntimeError("numpy is required when --target is used") from exc
    target = np.load(path, mmap_mode="r")
    if target.ndim != 2 or target.shape[1] != 3:
        raise ValueError(f"Expected target shape [T,3], got {target.shape}")
    return int(target.shape[0])


def ffconcat_quote(path: str) -> str:
    # ffconcat uses single-quoted paths; escape an embedded quote as '\''.
    return "'" + path.replace("'", "'\\''") + "'"


def validate_inventory(clips: list[ClipInfo], target_steps: int | None) -> int:
    if not clips:
        raise ValueError("No clips found")
    for clip in clips:
        if not abs(clip.fps - SOURCE_FPS) < 1e-6:
            raise ValueError(f"Expected 20 FPS, got {clip.fps} in {clip.path}")
        if clip.width < KEEP_WIDTH or clip.height != KEEP_HEIGHT:
            raise ValueError(
                f"Cannot take rightmost {KEEP_WIDTH}x{KEEP_HEIGHT} from "
                f"{clip.width}x{clip.height}: {clip.path}"
            )
    total_frames = sum(clip.frames for clip in clips)
    frames_per_step = SOURCE_FPS // TASK_FPS
    media_steps = total_frames // frames_per_step
    if target_steps is not None and media_steps != target_steps:
        raise ValueError(
            f"Dense media gives {media_steps} complete 5 Hz steps from "
            f"{total_frames} frames, but target has {target_steps}"
        )
    # A session may end with 1--3 source frames, which cannot form another
    # complete 5 Hz target interval. The target length is authoritative; the
    # final video/audio trim below discards only this sub-step tail.
    return target_steps if target_steps is not None else media_steps


def fixed_grid_frame_count(*, fps: float, dense_duration: float) -> int:
    """Number of samples on the grid ``n / fps < dense_duration``."""
    return int(math.ceil(dense_duration * fps - 1e-7))


def build_ffmpeg_command(
    clips: list[ClipInfo],
    output: Path,
    *,
    fps: float,
    dense_duration: float,
) -> list[str]:
    expected_frames = fixed_grid_frame_count(
        fps=fps, dense_duration=dense_duration
    )
    # Trim every decoded stream independently before concatenation. Released
    # AAC clips can contain roughly one codec-padding frame beyond their video
    # duration; concatenating them first would accumulate AV drift.
    command = [
        "ffmpeg",
        "-hide_banner",
        "-loglevel",
        "warning",
        "-y",
    ]
    for clip in clips:
        command.extend(["-i", clip.path])

    filters: list[str] = []
    concat_inputs: list[str] = []
    for index, clip in enumerate(clips):
        duration = clip.frames / SOURCE_FPS
        filters.append(
            f"[{index}:v:0]trim=start=0:end={duration:.9f},"
            f"setpts=PTS-STARTPTS,"
            f"crop={KEEP_WIDTH}:{KEEP_HEIGHT}:iw-{KEEP_WIDTH}:0[v{index}]"
        )
        filters.append(
            f"[{index}:a:0]atrim=start=0:end={duration:.9f},"
            f"asetpts=PTS-STARTPTS[a{index}]"
        )
        concat_inputs.extend([f"[v{index}]", f"[a{index}]"])
    filters.append(
        "".join(concat_inputs)
        + f"concat=n={len(clips)}:v=1:a=1[vcat][acat]"
    )
    # Establish one global visual lattice after the clips have become a dense
    # timeline. Some FFmpeg versions omit the last grid point when the input
    # ends exactly on an integer-second boundary. Add two disposable clones,
    # then trim by frame index, so EOF timestamp rounding cannot change the
    # number of samples. The clones are retained only if a real final grid
    # sample was missing.
    filters.append(
        f"[vcat]trim=start=0:end={dense_duration:.9f},"
        f"setpts=PTS-STARTPTS,"
        f"fps=fps={fps}:start_time=0:round=down,"
        f"tpad=stop_mode=clone:stop=2,"
        f"trim=start_frame=0:end_frame={expected_frames},"
        f"setpts=N/({fps}*TB),format=yuv420p[vout]"
    )
    filters.append(
        f"[acat]atrim=start=0:end={dense_duration:.9f},"
        "asetpts=PTS-STARTPTS[aout]"
    )

    command.extend(
        [
        "-filter_complex",
        ";".join(filters),
        "-map",
        "[vout]",
        "-map",
        "[aout]",
        "-c:v",
        "libx264",
        "-preset",
        "medium",
        "-crf",
        "18",
        "-c:a",
        "aac",
        "-b:a",
        "128k",
        "-movflags",
        "+faststart",
        str(output),
        ]
    )
    return command


def validate_output(path: Path, *, fps: float, dense_duration: float) -> dict:
    payload = _run_json(
        [
            "ffprobe",
            "-v",
            "error",
            "-show_entries",
            "stream=index,codec_type,avg_frame_rate,nb_frames,duration,sample_rate,channels",
            "-of",
            "json",
            str(path),
        ]
    )
    video = next(
        (stream for stream in payload.get("streams", []) if stream["codec_type"] == "video"),
        None,
    )
    audio = next(
        (stream for stream in payload.get("streams", []) if stream["codec_type"] == "audio"),
        None,
    )
    if video is None or audio is None:
        raise RuntimeError(f"Prepared file must contain video and audio: {path}")
    actual_fps = float(Fraction(video["avg_frame_rate"]))
    actual_frames = int(video["nb_frames"])
    expected_frames = fixed_grid_frame_count(
        fps=fps, dense_duration=dense_duration
    )
    if not math.isclose(actual_fps, fps, rel_tol=0.0, abs_tol=1e-6):
        raise RuntimeError(f"Prepared FPS is {actual_fps}, expected {fps}: {path}")
    if actual_frames != expected_frames:
        raise RuntimeError(
            f"Prepared frame count is {actual_frames}, expected "
            f"{expected_frames}: {path}"
        )
    audio_duration = float(audio["duration"])
    # AAC frame padding can shift the container duration by a few milliseconds.
    if abs(audio_duration - dense_duration) > 0.05:
        raise RuntimeError(
            f"Prepared audio duration {audio_duration:.6f}s differs from dense "
            f"timeline {dense_duration:.6f}s: {path}"
        )
    return {
        "video_fps": actual_fps,
        "video_frames": actual_frames,
        "expected_video_frames": expected_frames,
        "video_duration_seconds": float(video["duration"]),
        "audio_duration_seconds": audio_duration,
        "audio_sample_rate": int(audio["sample_rate"]),
        "audio_channels": int(audio["channels"]),
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--session-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--target", type=Path)
    parser.add_argument(
        "--fps",
        type=float,
        nargs="+",
        default=[2.0, 5.0],
        help="Fixed global visual rates to generate (default: 2 5).",
    )
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument(
        "--reuse-existing",
        action="store_true",
        help="Validate and reuse an existing output MP4 instead of re-encoding it.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    clip_paths = sorted(
        args.session_dir.glob("*.mp4"), key=lambda path: parse_timecode(path.stem)
    )
    clips = [probe_clip(path) for path in clip_paths]
    target_steps = load_target_steps(args.target) if args.target else None
    num_steps = validate_inventory(clips, target_steps)
    session = args.session_dir.name
    duration = num_steps / TASK_FPS

    args.output_dir.mkdir(parents=True, exist_ok=True)
    concat_path = args.output_dir / f"{session}.ffconcat"
    concat_text = "ffconcat version 1.0\n" + "".join(
        f"file {ffconcat_quote(clip.path)}\n" for clip in clips
    )
    concat_path.write_text(concat_text)

    outputs: dict[str, str] = {}
    output_media: dict[str, dict] = {}
    commands: list[list[str]] = []
    for fps in args.fps:
        if fps <= 0:
            raise ValueError("Every --fps value must be positive")
        fps_label = f"{fps:g}".replace(".", "p")
        output = args.output_dir / f"{session}_{fps_label}fps.mp4"
        outputs[f"{fps:g}"] = str(output.resolve())
        command = build_ffmpeg_command(
            clips,
            output,
            fps=fps,
            dense_duration=duration,
        )
        commands.append(command)
        if args.dry_run:
            print(shlex.join(command))
        elif args.reuse_existing and output.is_file():
            print(f"Reusing {output} ...", flush=True)
            output_media[f"{fps:g}"] = validate_output(
                output, fps=fps, dense_duration=duration
            )
        else:
            print(f"Preparing {output} ...", flush=True)
            subprocess.run(command, check=True)
            output_media[f"{fps:g}"] = validate_output(
                output, fps=fps, dense_duration=duration
            )

    source_frames = sum(clip.frames for clip in clips)
    used_source_frames = num_steps * (SOURCE_FPS // TASK_FPS)
    metadata = {
        "session": session,
        "source_session_dir": str(args.session_dir.resolve()),
        "source_clip_order": [asdict(clip) for clip in clips],
        "source_num_clips": len(clips),
        "source_frames_20hz": source_frames,
        "used_source_frames_20hz": used_source_frames,
        "discarded_tail_frames_20hz": source_frames - used_source_frames,
        "task_fps": TASK_FPS,
        "num_steps": num_steps,
        "dense_duration_seconds": duration,
        "crop": {
            "policy": "rightmost_1920x1080",
            "width": KEEP_WIDTH,
            "height": KEEP_HEIGHT,
        },
        "visual_sampling": "fixed_session_global_grid",
        "outputs": outputs,
        "output_media": output_media,
        "target": str(args.target.resolve()) if args.target else None,
        "dry_run": bool(args.dry_run),
        "reuse_existing": bool(args.reuse_existing),
        "commands": [shlex.join(command) for command in commands],
    }
    metadata_path = args.output_dir / f"{session}.json"
    metadata_path.write_text(json.dumps(metadata, indent=2) + "\n")
    print(
        f"{session}: {len(clips)} clips, {num_steps} steps, "
        f"{duration:.3f}s dense timeline"
    )
    print(f"Saved metadata to {metadata_path}")


if __name__ == "__main__":
    main()