#!/usr/bin/env python3
"""Extract EgoSpeak's Ego4D clips while loading Qwen only once.

Manifest sample indices are one-based and refer to data rows, excluding the
header.  For example, ``--start-index 31 --end-index 60`` selects exactly the
same clips as ``sed -n '32,61p' ego4d_egospeak.tsv``.

Only immutable model weights and the processor are reused.  Every clip invokes
the direct extractor independently, which creates a new media reader, rolling
KV cache, TMRoPE origin, and readout sequence.  Consequently no temporal state
is carried between dataset samples.
"""

from __future__ import annotations

import argparse
import csv
import gc
import math
import os
import sys
import time
import traceback
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence

import torch

import streaming_turn_model as core
import streaming_turn_model_directly as direct


TARGET_HZ = 5
DEFAULT_OUTPUT_NAME = "open_stream_2fps_20s_decord"
EXPECTED_MANIFEST_COLUMNS = (
    "split",
    "uid",
    "steps",
    "duration",
    "video_frames_2fps",
)


@dataclass(frozen=True)
class ClipRecord:
    index: int
    split: str
    uid: str
    steps: int
    duration: float
    video_frames_2fps: int


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--data-root",
        type=Path,
        default=Path("/data/wangx/Datasets/ego4d"),
        help="Ego4D experiment root (default: %(default)s)",
    )
    parser.add_argument(
        "--manifest",
        type=Path,
        help="Defaults to DATA_ROOT/ego4d_egospeak.tsv",
    )
    parser.add_argument(
        "--video-root",
        type=Path,
        help="Defaults to DATA_ROOT/v1/clips",
    )
    parser.add_argument(
        "--target-root",
        type=Path,
        help="Defaults to DATA_ROOT/ego4d_target_frame",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        help=(
            "Defaults to DATA_ROOT/omni_features/"
            f"{DEFAULT_OUTPUT_NAME}"
        ),
    )
    parser.add_argument(
        "--model-id",
        default="/data/wangx/Models/Qwen2.5-Omni-3B",
    )
    parser.add_argument(
        "--start-index",
        type=int,
        default=1,
        help="First one-based manifest data row to process (default: 1)",
    )
    parser.add_argument(
        "--end-index",
        type=int,
        help="Last one-based manifest data row to process, inclusive",
    )
    parser.add_argument(
        "--split",
        choices=("all", "train", "test"),
        default="all",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Replace completed outputs instead of resuming around them",
    )
    parser.add_argument(
        "--continue-on-error",
        action="store_true",
        help="Record a failed clip and continue with the next selected clip",
    )
    parser.add_argument("--memory-seconds", type=float, default=20.0)
    parser.add_argument("--step-seconds", type=float, default=0.2)
    parser.add_argument("--native-block-seconds", type=float, default=2.0)
    parser.add_argument("--video-fps", type=float, default=2.0)
    parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
    parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
    parser.add_argument("--frame-cache-size", type=int, default=64)
    parser.add_argument(
        "--decoder",
        choices=("decord", "pynvcodec"),
        default="decord",
    )
    parser.add_argument("--decoder-gpu-id", type=int, default=0)
    parser.add_argument("--decoder-threads", type=int, default=1)
    parser.add_argument("--max-open-readers", type=int, default=1)
    parser.add_argument("--verify-every", type=int, default=0)
    parser.add_argument("--verify-continuous-every", type=int, default=0)
    parser.add_argument("--no-progress", action="store_true")
    parser.add_argument(
        "--dtype",
        choices=("auto", "bfloat16", "float16", "float32"),
        default="bfloat16",
    )
    parser.add_argument("--flash-attention-2", action="store_true")
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument(
        "--empty-cache-every",
        type=int,
        default=0,
        help="Call torch.cuda.empty_cache every N processed clips; 0 disables",
    )
    args = parser.parse_args(argv)

    if args.start_index <= 0:
        parser.error("--start-index must be positive")
    if args.end_index is not None and args.end_index < args.start_index:
        parser.error("--end-index must be >= --start-index")
    if args.empty_cache_every < 0:
        parser.error("--empty-cache-every must be >= 0")
    if args.step_seconds <= 0:
        parser.error("--step-seconds must be positive")
    return args


def resolve_paths(args: argparse.Namespace) -> None:
    args.data_root = args.data_root.expanduser().resolve()
    args.manifest = (
        args.manifest or args.data_root / "ego4d_egospeak.tsv"
    ).expanduser().resolve()
    args.video_root = (
        args.video_root or args.data_root / "v1" / "clips"
    ).expanduser().resolve()
    args.target_root = (
        args.target_root or args.data_root / "ego4d_target_frame"
    ).expanduser().resolve()
    args.output_root = (
        args.output_root
        or args.data_root / "omni_features" / DEFAULT_OUTPUT_NAME
    ).expanduser().resolve()


def load_manifest(path: Path) -> list[ClipRecord]:
    if not path.is_file():
        raise FileNotFoundError(f"Missing EgoSpeak manifest: {path}")
    records: list[ClipRecord] = []
    seen: set[str] = set()
    with path.open(newline="") as file:
        reader = csv.DictReader(file, delimiter="\t")
        if reader.fieldnames != list(EXPECTED_MANIFEST_COLUMNS):
            raise RuntimeError(
                f"Unexpected manifest columns in {path}: {reader.fieldnames}"
            )
        for index, row in enumerate(reader, start=1):
            split = row["split"].strip()
            uid = row["uid"].strip()
            if split not in {"train", "test"}:
                raise RuntimeError(f"Invalid split at manifest row {index}: {split}")
            if not uid or uid in seen:
                raise RuntimeError(
                    f"Missing or duplicate UID at manifest row {index}: {uid!r}"
                )
            seen.add(uid)
            steps = int(row["steps"])
            duration = float(row["duration"])
            video_frames = int(row["video_frames_2fps"])
            if steps <= 0:
                raise RuntimeError(f"Non-positive steps for {uid}: {steps}")
            expected_duration = steps / TARGET_HZ
            if not math.isclose(
                duration, expected_duration, rel_tol=0.0, abs_tol=1e-6
            ):
                raise RuntimeError(
                    f"Manifest duration mismatch for {uid}: {duration} vs "
                    f"{expected_duration} from {steps} target steps"
                )
            records.append(
                ClipRecord(
                    index=index,
                    split=split,
                    uid=uid,
                    steps=steps,
                    duration=duration,
                    video_frames_2fps=video_frames,
                )
            )
    if not records:
        raise RuntimeError(f"No clips in {path}")
    return records


def select_records(
    records: list[ClipRecord], args: argparse.Namespace
) -> list[ClipRecord]:
    end_index = args.end_index if args.end_index is not None else len(records)
    selected = [
        record
        for record in records
        if args.start_index <= record.index <= end_index
        and (args.split == "all" or record.split == args.split)
    ]
    if not selected:
        raise RuntimeError(
            f"No manifest rows selected by {args.start_index=} "
            f"{args.end_index=} {args.split=}"
        )
    return selected


def preflight(selected: list[ClipRecord], args: argparse.Namespace) -> None:
    if not args.video_root.is_dir():
        raise FileNotFoundError(f"Missing video root: {args.video_root}")
    if not args.target_root.is_dir():
        raise FileNotFoundError(f"Missing target root: {args.target_root}")
    model_path = Path(args.model_id).expanduser()
    if model_path.is_absolute() and not model_path.is_dir():
        raise FileNotFoundError(f"Missing local model: {model_path}")

    problems: list[str] = []
    for record in selected:
        video = args.video_root / f"{record.uid}.mp4"
        target = args.target_root / f"{record.uid}.npy"
        if not video.is_file() or video.stat().st_size == 0:
            problems.append(f"missing video: {video}")
        if not target.is_file() or target.stat().st_size == 0:
            problems.append(f"missing target: {target}")
            continue
        try:
            target_steps = core.target_num_steps(target)
        except Exception as error:  # report all bad inputs before model loading
            problems.append(f"invalid target {target}: {error}")
            continue
        if target_steps != record.steps:
            problems.append(
                f"target length mismatch {record.uid}: manifest={record.steps}, "
                f"npy={target_steps}"
            )
    if problems:
        preview = "\n".join(f"- {problem}" for problem in problems[:20])
        suffix = "" if len(problems) <= 20 else f"\n... and {len(problems)-20} more"
        raise RuntimeError(f"Ego4D preflight failed:\n{preview}{suffix}")


def clip_argv(
    record: ClipRecord,
    args: argparse.Namespace,
    output: Path,
) -> list[str]:
    values = [
        "--model-id",
        str(args.model_id),
        "--video",
        str(args.video_root / f"{record.uid}.mp4"),
        "--target",
        str(args.target_root / f"{record.uid}.npy"),
        "--mode",
        "open-stream-replay",
        "--step-seconds",
        str(args.step_seconds),
        "--native-block-seconds",
        str(args.native_block_seconds),
        "--memory-seconds",
        str(args.memory_seconds),
        "--video-fps",
        str(args.video_fps),
        "--video-min-pixels",
        str(args.video_min_pixels),
        "--video-max-pixels",
        str(args.video_max_pixels),
        "--frame-cache-size",
        str(args.frame_cache_size),
        "--decoder",
        args.decoder,
        "--decoder-gpu-id",
        str(args.decoder_gpu_id),
        "--decoder-threads",
        str(args.decoder_threads),
        "--max-open-readers",
        str(args.max_open_readers),
        "--verify-every",
        str(args.verify_every),
        "--verify-continuous-every",
        str(args.verify_continuous_every),
        "--dtype",
        args.dtype,
        "--device",
        args.device,
        "--output",
        str(output),
    ]
    if args.flash_attention_2:
        values.append("--flash-attention-2")
    if args.no_progress:
        values.append("--no-progress")
    return values


def format_eta(seconds: float) -> str:
    if not math.isfinite(seconds) or seconds < 0:
        return "unknown"
    hours, remainder = divmod(int(seconds), 3600)
    minutes, secs = divmod(remainder, 60)
    return f"{hours:02d}:{minutes:02d}:{secs:02d}"


def main(argv: Sequence[str] | None = None) -> None:
    args = parse_args(argv)
    resolve_paths(args)
    core._configure_qwen_logging()

    records = load_manifest(args.manifest)
    selected = select_records(records, args)
    preflight(selected, args)
    (args.output_root / "train").mkdir(parents=True, exist_ok=True)
    (args.output_root / "test").mkdir(parents=True, exist_ok=True)

    pending = [
        record
        for record in selected
        if args.overwrite
        or not (args.output_root / record.split / f"{record.uid}.pt").is_file()
        or (args.output_root / record.split / f"{record.uid}.pt").stat().st_size == 0
    ]
    skipped = len(selected) - len(pending)
    print(
        f"Selected manifest samples {selected[0].index}-{selected[-1].index}: "
        f"{len(selected)} clip(s), pending={len(pending)}, skipped={skipped}",
        flush=True,
    )
    print(f"Output root: {args.output_root}", flush=True)
    if not pending:
        print("All selected clips already have non-empty outputs.")
        return

    print(f"Loading Qwen once from {args.model_id} ...", flush=True)
    load_started = time.perf_counter()
    model, processor = core.load_qwen_omni(
        args.model_id,
        dtype=args.dtype,
        device_map={"": args.device},
        use_flash_attention_2=args.flash_attention_2,
    )
    model.requires_grad_(False)
    print(
        f"Model loaded once in {time.perf_counter() - load_started:.1f}s; "
        "per-clip KV/media state will reset.",
        flush=True,
    )

    run_started = time.perf_counter()
    completed = 0
    failures: list[tuple[ClipRecord, str]] = []
    for pending_index, record in enumerate(pending, start=1):
        output = args.output_root / record.split / f"{record.uid}.pt"
        temporary = output.with_name(f"{output.name}.partial.{os.getpid()}")
        print(
            f"\n[{pending_index}/{len(pending)} | manifest {record.index}] "
            f"{record.split}/{record.uid}: {record.steps} states, "
            f"{record.duration:.1f}s",
            flush=True,
        )
        clip_started = time.perf_counter()
        try:
            direct.main(
                clip_argv(record, args, temporary),
                model=model,
                processor=processor,
            )
            os.replace(temporary, output)
            completed += 1
            elapsed = time.perf_counter() - clip_started
            average = (time.perf_counter() - run_started) / completed
            remaining = len(pending) - pending_index
            print(
                f"Completed {output} in {elapsed:.1f}s; "
                f"ETA {format_eta(average * remaining)}",
                flush=True,
            )
        except KeyboardInterrupt:
            print(f"Interrupted; partial path: {temporary}", file=sys.stderr)
            raise
        except Exception as error:
            failures.append((record, str(error)))
            print(
                f"FAILED manifest sample {record.index} ({record.uid}): {error}\n"
                f"Partial path, if created: {temporary}",
                file=sys.stderr,
                flush=True,
            )
            traceback.print_exc()
            if not args.continue_on_error:
                raise
        finally:
            # The model and processor stay referenced above.  Collection here
            # releases the completed clip's reader, PIL frames, branch caches,
            # and Python cache wrappers before opening the next sample.
            gc.collect()
            processed = completed + len(failures)
            if (
                args.empty_cache_every > 0
                and processed > 0
                and processed % args.empty_cache_every == 0
                and torch.cuda.is_available()
            ):
                torch.cuda.empty_cache()

    total = time.perf_counter() - run_started
    print(
        f"\nBatch finished in {format_eta(total)}: completed={completed}, "
        f"skipped={skipped}, failed={len(failures)}",
        flush=True,
    )
    if failures:
        for record, message in failures:
            print(f"- manifest {record.index} {record.uid}: {message}", file=sys.stderr)
        raise SystemExit(1)


if __name__ == "__main__":
    main()