#!/usr/bin/env python3
"""Build fixed-frame SegmentQA clips and ms-swift JSONL manifests."""

from __future__ import annotations

import argparse
import json
import os
import re
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import datetime, timezone
from fractions import Fraction
from pathlib import Path
from typing import Any

import av
import pandas as pd
from PIL import Image, ImageOps


DATASET_NAMES = ("heico-focus-vqa", "lapchole-focus-vqa")
STATUS_NAME = "segment_status.jsonl"
BASELINE_SYSTEM_PROMPT = (
    "You are a surgical video analysis assistant. You are shown a short clip from a laparoscopic procedure and asked a single question about it.\n\n"
    "The clip is given as a sequence of frames in chronological order. Each frame is preceded by its timestamp in the source video, written as "
    "<SECONDS seconds> -- a frame marked <1234.5 seconds> was taken 1234.5 seconds (00:20:34) after the video began. Consecutive frames may be several seconds "
    "apart, so an event can happen between two of them.\n\n"
    "A foreign object (FO) is any object fully introduced into the patient's body\n"
    "cavity during surgery that must be retrieved or accounted for. Importantly,\n"
    "standard surgical instruments that remain connected to the external environment\n"
    "(e.g., graspers, scissors, trocars, staplers, cameras) are not considered foreign\n"
    "objects. Furthermore, we exclude detachable parts of surgical instruments,\n"
    "particularly anvil components of staplers.\n\n"
    "The foreign object classes are exactly: Sponge, Clip, Specimen Bag, Silicone Loop, External Drain, Needle, Gallstone, Specimen, Mesh, Absorbable Hemostatic Agent.\n\n"
    "Reply with the answer and nothing else -- no reasoning, no preamble, no\n"
    "explanation, no restating the question. A single short line.\n\n"
    "Rules for the answer:\n"
    "- Write the value only. No sentence, no explanation, no units, no trailing\n"
    "  period, and never repeat the question.\n"
    "- Asks yes or no -> write exactly: yes   or   no\n"
    "- Asks how many / for a count -> write digits only, e.g. 0 or 1 or 2.\n"
    "- Asks which foreign object class(es) -> write class names exactly as spelled\n"
    "  in the list above, comma-separated (e.g. Clip, Sponge), or exactly: none\n"
    "  Never answer with a generic description such as \"surgical instrument\".\n"
    "- Asks when something happens -> write hh:mm:ss, read off the frame timestamps.\n"
    "  Those timestamps count from the START OF THE VIDEO, not from the start of this\n"
    "  clip, so <1234.5 seconds> is 00:20:34. Convert carefully: 3600 seconds is one\n"
    "  hour, 60 seconds is one minute. If the event falls between two frames, give\n"
    "  your best estimate between their timestamps rather than snapping to one.\n"
    "  For several time points, separate them with commas: 00:09:19, 00:12:44\n"
    "- Asks for a percentage or a share -> write the number only, e.g. 40\n"
    "- Lists options to choose from -> copy exactly one of those options, verbatim.\n"
    "- Anything else -> a short phrase, at most a few words.\n\n"
    "If you are unsure, still commit to your single best answer in the required\n"
    "form. An empty, hedged, or explanatory answer is scored as wrong.\n"
)


@dataclass(frozen=True)
class Segment:
    dataset: str
    split: str
    video: str
    start: float
    end: float
    relpath: str
    source_path: str

    @property
    def key(self) -> tuple[str, str, str, float, float]:
        return self.dataset, self.split, self.video, self.start, self.end


def parse_time(value: Any) -> float:
    if isinstance(value, (int, float)):
        return float(value)
    parts = str(value).split(":")
    if len(parts) != 3:
        raise ValueError(f"expected HH:MM:SS timestamp, got {value!r}")
    hours, minutes, seconds = parts
    return int(hours) * 3600 + int(minutes) * 60 + float(seconds)


def jsonable(value: Any) -> Any:
    if hasattr(value, "tolist"):
        return value.tolist()
    if not isinstance(value, (list, tuple, dict)) and pd.isna(value):
        return None
    return value


def slug(value: str) -> str:
    return re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("_")


def seconds_token(value: float) -> str:
    return f"{value:.3f}".rstrip("0").rstrip(".").replace(".", "p")


def segment_relpath(dataset: str, split: str, video: str, start: float, end: float) -> str:
    stem = slug(Path(video).stem)
    name = f"{stem}__{seconds_token(start)}_{seconds_token(end)}.mp4"
    return str(Path("videos") / split / dataset / name)


def target_times(start: float, end: float, frames: int) -> list[float]:
    duration = end - start
    if duration <= 0:
        raise ValueError(f"non-positive segment duration: start={start} end={end}")
    return [start + (i + 0.5) * duration / frames for i in range(frames)]


def build_datasets(data_root: Path, splits: list[str]) -> dict[str, dict[str, Path]]:
    return {
        name: {
            **{split: data_root / name / f"data/segment/{split}.parquet" for split in splits},
            "videos": data_root / name / "videos",
        }
        for name in DATASET_NAMES
    }


def load_rows(datasets: dict[str, dict[str, Path]], splits: list[str]) -> tuple[list[dict[str, Any]], list[Segment]]:
    rows: list[dict[str, Any]] = []
    by_key: dict[tuple[str, str, str, float, float], Segment] = {}
    for split in splits:
        for dataset, cfg in datasets.items():
            df = pd.read_parquet(cfg[split])
            for _, row in df.iterrows():
                raw = {k: jsonable(v) for k, v in row.to_dict().items()}
                start = parse_time(raw["timestamp_start"])
                end = parse_time(raw["timestamp_end"])
                video = str(raw["video"])
                source = cfg["videos"] / video
                rel = segment_relpath(dataset, split, video, start, end)
                segment = Segment(dataset, split, video, start, end, rel, str(source))
                by_key.setdefault(segment.key, segment)
                raw.update(
                    {
                        "dataset": dataset,
                        "split": split,
                        "source_video_path": str(source),
                        "segment_start_seconds": start,
                        "segment_end_seconds": end,
                        "derived_video": rel,
                    }
                )
                rows.append(raw)
    return rows, list(by_key.values())


def apply_limit(segments: list[Segment], limit: int | None) -> list[Segment]:
    ordered = sorted(segments, key=lambda s: s.key)
    if limit is None:
        return ordered
    buckets = {name: [s for s in ordered if s.dataset == name] for name in DATASET_NAMES}
    picked: list[Segment] = []
    index = 0
    while len(picked) < limit:
        before = len(picked)
        for dataset in DATASET_NAMES:
            if index < len(buckets[dataset]) and len(picked) < limit:
                picked.append(buckets[dataset][index])
        if len(picked) == before:
            break
        index += 1
    return picked


def frame_seconds(frame: av.VideoFrame, stream: av.video.stream.VideoStream) -> float | None:
    if frame.pts is None:
        return None
    base = frame.time_base or stream.time_base
    return float(frame.pts * base)


def letterbox(frame: av.VideoFrame, width: int, height: int) -> av.VideoFrame:
    image = frame.to_image().convert("RGB")
    padded = ImageOps.pad(image, (width, height), method=Image.Resampling.BICUBIC, color=(0, 0, 0), centering=(0.5, 0.5))
    return av.VideoFrame.from_image(padded).reformat(format="yuv420p")


def select_frames(segment: Segment, frames: int, width: int, height: int) -> tuple[list[av.VideoFrame], list[float]]:
    targets = target_times(segment.start, segment.end, frames)
    selected: list[av.VideoFrame | None] = [None] * frames
    selected_pts: list[float | None] = [None] * frames
    last: av.VideoFrame | None = None
    last_pts: float | None = None
    next_idx = 0

    with av.open(segment.source_path) as container:
        stream = container.streams.video[0]
        seek_offset = int(max(segment.start - 2.0, 0.0) / float(stream.time_base))
        container.seek(seek_offset, stream=stream, backward=True, any_frame=False)
        stop_decode = False
        for packet in container.demux(stream):
            for frame in packet.decode():
                pts = frame_seconds(frame, stream)
                if pts is None:
                    continue
                if pts < segment.start:
                    continue
                if pts >= segment.end:
                    stop_decode = True
                    break
                while last is not None and next_idx < frames and targets[next_idx] <= pts:
                    if abs(last_pts - targets[next_idx]) <= abs(pts - targets[next_idx]):
                        chosen, chosen_pts = last, last_pts
                    else:
                        chosen, chosen_pts = frame, pts
                    selected[next_idx] = letterbox(chosen, width, height)
                    selected_pts[next_idx] = chosen_pts
                    next_idx += 1
                if last is None and next_idx < frames and targets[next_idx] <= pts:
                    selected[next_idx] = letterbox(frame, width, height)
                    selected_pts[next_idx] = pts
                    next_idx += 1
                last, last_pts = frame, pts
            if stop_decode or next_idx >= frames:
                break

    if last is None or last_pts is None:
        raise RuntimeError("decoded zero frames inside segment")
    for index, frame in enumerate(selected):
        if frame is None:
            selected[index] = letterbox(last, width, height)
            selected_pts[index] = last_pts
    return [f for f in selected if f is not None], [float(x) for x in selected_pts if x is not None]


def count_video_frames(path: Path) -> tuple[int, int, int]:
    with av.open(str(path)) as container:
        stream = container.streams.video[0]
        count = 0
        width = height = 0
        for frame in container.decode(stream):
            count += 1
            width, height = frame.width, frame.height
        return count, width, height


def metadata_from_pts(source_pts: list[float], width: int, height: int) -> dict[str, Any]:
    frame_indices = [int(round(value * 1000)) for value in source_pts]
    total_num_frames = max(frame_indices) + 1
    return {
        "fps": 1000,
        "frames_indices": frame_indices,
        "total_num_frames": total_num_frames,
        "width": width,
        "height": height,
        "duration": total_num_frames / 1000,
        "video_backend": "precomputed_absolute_pts",
    }


def metadata_contract_matches(item: dict[str, Any], frames: int, width: int, height: int) -> bool:
    metadata = item.get("video_metadata")
    source_pts = item.get("selected_source_pts")
    if not isinstance(metadata, dict) or not isinstance(source_pts, list) or len(source_pts) != frames:
        return False
    frame_indices = [int(round(float(value) * 1000)) for value in source_pts]
    return (
        metadata.get("fps") == 1000
        and metadata.get("frames_indices") == frame_indices
        and metadata.get("total_num_frames") == max(frame_indices) + 1
        and metadata.get("width") == width
        and metadata.get("height") == height
        and metadata.get("duration") == metadata.get("total_num_frames") / 1000
        and metadata.get("video_backend") == "precomputed_absolute_pts"
    )


def verified_success_item(item: dict[str, Any], output_dir: Path, frames: int, width: int, height: int) -> bool:
    path = output_dir / item.get("relpath", "")
    if item.get("status") != "success" or item.get("verified_frame_count") != frames:
        return False
    if item.get("width") != width or item.get("height") != height or not path.exists():
        return False
    try:
        actual = count_video_frames(path)
    except Exception:
        return False
    return actual == (frames, width, height) and metadata_contract_matches(item, frames, width, height)


def encode_segment(segment: Segment, output_dir: Path, frames: int, width: int, height: int) -> dict[str, Any]:
    out_path = output_dir / segment.relpath
    partial = out_path.with_suffix(out_path.suffix + ".partial")
    selected, source_pts = select_frames(segment, frames, width, height)
    if len(selected) != frames:
        raise RuntimeError(f"selected {len(selected)} frames, expected {frames}")

    out_path.parent.mkdir(parents=True, exist_ok=True)
    try:
        with av.open(str(partial), "w", format="mp4") as container:
            stream = container.add_stream("libx264", rate=Fraction(20, 1))
            stream.width = width
            stream.height = height
            stream.pix_fmt = "yuv420p"
            stream.options = {"crf": "18", "preset": "veryfast", "an": "1"}
            for index, frame in enumerate(selected):
                frame.pts = index
                frame.time_base = Fraction(1, 20)
                for packet in stream.encode(frame):
                    container.mux(packet)
            for packet in stream.encode():
                container.mux(packet)
        verified_count, verified_width, verified_height = count_video_frames(partial)
        if (verified_count, verified_width, verified_height) != (frames, width, height):
            raise RuntimeError(
                f"encoded {(verified_count, verified_width, verified_height)}, expected {(frames, width, height)}")
        os.replace(partial, out_path)
    except Exception:
        try:
            partial.unlink()
        except FileNotFoundError:
            pass
        raise
    metadata = metadata_from_pts(source_pts, width, height)
    return {
        "dataset": segment.dataset,
        "split": segment.split,
        "video": segment.video,
        "source_video_path": segment.source_path,
        "start_seconds": segment.start,
        "end_seconds": segment.end,
        "relpath": segment.relpath,
        "frame_count": frames,
        "verified_frame_count": verified_count,
        "width": verified_width,
        "height": verified_height,
        "selected_source_pts": source_pts,
        "video_metadata": metadata,
        "sampling_rule": f"uniform_pts_bin_centers_[start,end)_{frames}",
        "resize": f"letterbox_{width}x{height}",
        "status": "success",
    }


def done_segments(status_path: Path, output_dir: Path, frames: int, width: int, height: int) -> set[str]:
    done: set[str] = set()
    if not status_path.exists():
        return done
    for line in status_path.read_text().splitlines():
        if not line.strip():
            continue
        item = json.loads(line)
        if verified_success_item(item, output_dir, frames, width, height):
            done.add(item["relpath"])
    return done


def append_jsonl(path: Path, item: dict[str, Any]) -> None:
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(item, ensure_ascii=False) + "\n")
        handle.flush()


def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")


def ms_swift_record(row: dict[str, Any], segment_info: dict[str, Any], output_dir: Path, include_answer: bool) -> dict[str, Any]:
    messages = [
        {"role": "system", "content": BASELINE_SYSTEM_PROMPT},
        {"role": "user", "content": f"<video>\n{row['question']}"},
    ]
    if include_answer:
        messages.append({"role": "assistant", "content": str(row["answer"])})
    return {
        "messages": messages,
        "videos": [str((output_dir / row["derived_video"]).resolve())],
        "chat_template_kwargs": {
            "nframes": segment_info["frame_count"],
            "enable_thinking": False,
            "video_metadata": segment_info["video_metadata"],
        },
        "segment_metadata": {
            k: row.get(k)
            for k in [
                "id",
                "dataset",
                "split",
                "video",
                "timestamp_start",
                "timestamp_end",
                "answer_format",
                "segment_start_seconds",
                "segment_end_seconds",
                "derived_video",
            ]
        },
    }


def write_outputs(output_dir: Path, rows: list[dict[str, Any]], manifest_by_rel: dict[str, dict[str, Any]], splits: list[str]) -> None:
    selected = [row for row in rows if row["derived_video"] in manifest_by_rel]
    if "train" in splits:
        write_jsonl(
            output_dir / "train.jsonl",
            [ms_swift_record(row, manifest_by_rel[row["derived_video"]], output_dir, True) for row in selected if row["split"] == "train"],
        )
    if "test" in splits:
        test_rows = [row for row in selected if row["split"] == "test"]
        write_jsonl(
            output_dir / "test_requests.jsonl",
            [ms_swift_record(row, manifest_by_rel[row["derived_video"]], output_dir, False) for row in test_rows],
        )
        write_jsonl(
            output_dir / "test_references.jsonl",
            [{"id": row["id"], "dataset": row["dataset"], "answer": row["answer"], "answer_format": row["answer_format"]} for row in test_rows],
        )


def write_docs(output_dir: Path, args: argparse.Namespace, datasets: dict[str, dict[str, Path]], questions: int, segments: int, failures: int) -> None:
    now = datetime.now(timezone.utc).isoformat()
    provenance = {
        "version": "fixed_frames_v1",
        "created_at_utc": now,
        "command": " ".join(os.sys.argv),
        "inputs": {name: {k: str(v) for k, v in cfg.items()} for name, cfg in datasets.items()},
        "splits": args.splits,
        "question_count": questions,
        "unique_segment_count": segments,
        "failed_count": failures,
        "frame_count": args.frames,
        "sampling_rule": "uniform PTS bin centers in [start,end), nearest real source PTS",
        "resize": f"proportional letterbox {args.width}x{args.height}",
        "encoding": {"container": "mp4", "codec": "H.264/libx264", "audio": "none", "pix_fmt": "yuv420p"},
        "metadata_contract": {
            "fps": 1000,
            "frames_indices": "round(abs_pts*1000)",
            "total_num_frames": "max(frames_indices)+1",
            "width": args.width,
            "height": args.height,
            "duration": "total_num_frames/fps",
            "video_backend": "precomputed_absolute_pts",
        },
        "system_prompt_source": {
            "author_repo_commit": "3fc0850f9f76261c529e75c78ad2272558ff9e4e",
            "source": "https://github.com/AzJ24/Orena/blob/3fc0850f9f76261c529e75c78ad2272558ff9e4e/orena_sft/prompts.py",
            "function": "build_system_prompt",
            "parameters": {"include_definitions": False, "style": "direct", "track": "segment"},
            "foreign_objects_source": "orena-focus 0.3.4 wheel: focus/foreign_objects.py",
        },
        "system_prompt": BASELINE_SYSTEM_PROMPT,
        "output_dir": str(output_dir),
        "resume": bool(args.resume),
        "workers": args.workers,
        "limit": args.limit,
        "seed": args.seed,
    }
    (output_dir / "provenance.json").write_text(json.dumps(provenance, indent=2, ensure_ascii=False) + "\n")
    report = {
        "status": "pass" if failures == 0 and args.success_count == segments else "fail",
        "expected_frame_count": args.frames,
        "expected_width": args.width,
        "expected_height": args.height,
        "segments": segments,
        "success": args.success_count,
        "missing": segments - args.success_count,
        "failures": failures,
        "checked_at_utc": now,
    }
    (output_dir / "validation_report.json").write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n")


def run(args: argparse.Namespace) -> int:
    data_root = Path(args.data_root)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    datasets = build_datasets(data_root, args.splits)
    rows, all_segments = load_rows(datasets, args.splits)
    segments = apply_limit(all_segments, args.limit)
    selected_relpaths = {segment.relpath for segment in segments}
    selected_rows = [row for row in rows if row["derived_video"] in selected_relpaths]
    status_path = output_dir / STATUS_NAME
    existing = done_segments(status_path, output_dir, args.frames, args.width, args.height) if args.resume else set()
    todo = [s for s in segments if s.relpath not in existing]
    failures = 0
    manifest_by_rel: dict[str, dict[str, Any]] = {}
    print(json.dumps({"event": "start", "output": str(output_dir), "splits": args.splits, "segments": len(segments), "todo": len(todo), "python": os.sys.executable}, ensure_ascii=False), flush=True)

    def work(segment: Segment) -> dict[str, Any]:
        started = time.time()
        try:
            item = encode_segment(segment, output_dir, args.frames, args.width, args.height)
        except Exception as exc:
            return {
                "dataset": segment.dataset,
                "split": segment.split,
                "video": segment.video,
                "source_video_path": segment.source_path,
                "start_seconds": segment.start,
                "end_seconds": segment.end,
                "relpath": segment.relpath,
                "status": "failed",
                "error": str(exc),
                "elapsed_seconds": round(time.time() - started, 3),
            }
        item["elapsed_seconds"] = round(time.time() - started, 3)
        return item

    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = [pool.submit(work, segment) for segment in todo]
        for future in as_completed(futures):
            item = future.result()
            append_jsonl(status_path, item)
            if item["status"] == "success":
                manifest_by_rel[item["relpath"]] = item
            else:
                failures += 1
            print(f"relpath={item['relpath']} status={item['status']} elapsed_seconds={item.get('elapsed_seconds')}", flush=True)

    for line in status_path.read_text().splitlines() if status_path.exists() else []:
        if not line.strip():
            continue
        item = json.loads(line)
        if item.get("relpath") in selected_relpaths and verified_success_item(
            item, output_dir, args.frames, args.width, args.height
        ):
            manifest_by_rel[item["relpath"]] = item
    manifest = [manifest_by_rel[s.relpath] for s in segments if s.relpath in manifest_by_rel]
    write_jsonl(output_dir / "segment_manifest.jsonl", manifest)
    write_outputs(output_dir, selected_rows, manifest_by_rel, args.splits)
    args.success_count = len(manifest)
    write_docs(output_dir, args, datasets, len(selected_rows), len(segments), failures)
    print(json.dumps({"questions": len(selected_rows), "segments": len(segments), "success": len(manifest), "failed": failures}, ensure_ascii=False))
    return 0 if failures == 0 and len(manifest) == len(segments) else 1


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-root", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--splits", nargs="+", choices=["train", "test"], required=True)
    parser.add_argument("--frames", type=int, default=80)
    parser.add_argument("--width", type=int, default=420)
    parser.add_argument("--height", type=int, default=224)
    parser.add_argument("--workers", type=int, default=1)
    parser.add_argument("--limit", type=int)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()
    if args.workers < 1:
        raise SystemExit("--workers must be >= 1")
    if args.frames < 1:
        raise SystemExit("--frames must be >= 1")
    if args.width < 1 or args.height < 1:
        raise SystemExit("--width and --height must be >= 1")
    return run(args)


if __name__ == "__main__":
    raise SystemExit(main())