"""
Build the 5K SFT/OPD dataset:
  - 1K from TemporalBench (long_qa)
  - 4K from LLaVA-Video-178K (cap / oe / mc across academic_v0_1, nextqa, activitynetqa, perceptiontest)

Filters: only entries whose video is in the corresponding manifest_jsonl (≥32 decodable frames).

Output:
  - {OUT_DIR}/sft_5k.jsonl       — chat-template messages, multi-turn preserved
  - {OUT_DIR}/sft_5k.parquet     — verl-ready (prompt[list], videos[list], response, data_source, ability, extra_info)
  - {OUT_DIR}/eval_100.jsonl     — held-out subset for sanity checks (not in sft_5k)

Each row has:
  prompt:   list[ChatMessage]   (only user-side messages; for OPD the student generates the response)
  response: str                 (the chosen ground-truth response for SFT; ignored by OPD)
  videos:   list[dict]          (one entry per <video> placeholder)
  ability:  str
  data_source: str
  extra_info: dict

For SFT scripts that want full multi-turn supervision, see jsonl `messages` field.
"""
from __future__ import annotations

import argparse
import json
import os
import random
import sys
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

import pandas as pd

ROOT = Path("/mnt/local-fast/opd_zt")
RAW = ROOT / "data" / "raw"
VIDEOS = ROOT / "data" / "videos"
FILTERED = ROOT / "data" / "filtered"
OUT_DIR = ROOT / "data"

# Where each source's videos actually live (relative to VIDEOS / src).
# Manifest rel_paths are relative to VIDEOS/<src>.
TB_VID_ROOT = VIDEOS / "tb"            # contains long_video/{ActivityNet,COIN,...}
LV_VID_ROOT = VIDEOS / "lv178k"        # contains {split}/academic_source/... or similar


# ---------------------------------------------------------------------------
# helpers


def load_manifest(src: str) -> dict[str, dict]:
    p = FILTERED / f"{src}_manifest.jsonl"
    if not p.exists():
        print(f"FATAL: manifest missing: {p}", file=sys.stderr)
        sys.exit(2)
    out: dict[str, dict] = {}
    with p.open() as f:
        for line in f:
            r = json.loads(line)
            out[r["rel_path"]] = r
    return out


def make_video_dict(abs_path: str, manifest_entry: dict, frames: int) -> dict:
    """Format the video dict consumed by Qwen2.5-VL processor.

    Pin both min and max frames so qwen-vl-utils samples exactly N frames.
    """
    return {
        "video": f"file://{abs_path}",
        "max_frames": frames,
        "min_frames": frames,
        # explicit nframes also accepted by qwen-vl-utils
        "nframes": frames,
        # Qwen2.5-VL: keep pixel budget moderate per frame to control vmemory.
        "max_pixels": 360 * 420,
    }


def conv_to_messages(conv: list[dict]) -> list[dict]:
    """LLaVA conversation -> chat messages. Replace <image> with <video> placeholder.

    The convention: the first human turn has <image>; we replace it with <video>
    so verl's rl_dataset._build_messages routes the entry through the video path.
    """
    msgs: list[dict] = []
    for turn in conv:
        role = "user" if turn["from"] == "human" else "assistant"
        text = turn["value"]
        # llava uses <image> as the video placeholder for video data
        text = text.replace("<image>", "<video>").strip()
        msgs.append({"role": role, "content": text})
    return msgs


# ---------------------------------------------------------------------------
# TemporalBench → records


def load_tb_records(tb_manifest: dict[str, dict]) -> list[dict]:
    """Use long_qa.json (multi-choice). 1 turn per record."""
    p = RAW / "tb" / "temporalbench_long_qa.json"
    if not p.exists():
        print(f"FATAL: TB long_qa missing: {p}", file=sys.stderr)
        sys.exit(2)

    raw = json.loads(p.read_text())
    out: list[dict] = []
    skipped = 0
    for entry in raw:
        rel = entry["video_name"]   # 'long_video/ActivityNet/abc.mp4'
        if rel not in tb_manifest:
            skipped += 1
            continue
        abs_path = TB_VID_ROOT / rel
        question = entry["question"]
        gt_letter = entry["GT"]
        # answer string: the option letter expected (paper format)
        answer = gt_letter
        msgs = [
            {"role": "user", "content": f"<video>\n{question}"},
            {"role": "assistant", "content": answer},
        ]
        rec = {
            "data_source": "tb_long_qa",
            "ability": "video_mc_qa",
            "video_rel": rel,
            "video_abs": str(abs_path),
            "messages": msgs,
            "extra_info": {
                "idx": entry["idx"],
                "dataset": entry["dataset"],
                "manifest": tb_manifest[rel],
            },
        }
        out.append(rec)
    print(f"[tb] loaded {len(out)} records from long_qa (skipped {skipped} missing)",
          flush=True)
    return out


# ---------------------------------------------------------------------------
# LLaVA-Video-178K → records


def lv_video_abs(split: str, rel_in_split: str) -> Path:
    """Convert LLaVA-Video-178K 'video' field to abs path under VIDEOS/lv178k/<split>/.

    The 'video' field is like 'academic_source/Charades/Q04US.mp4'.
    Videos live under VIDEOS/lv178k/<split>/<...>.
    """
    return LV_VID_ROOT / split / rel_in_split


def lv_rel_for_manifest(split: str, rel_in_split: str) -> str:
    """The rel_path key the manifest uses (relative to VIDEOS/lv178k/)."""
    return f"{split}/{rel_in_split}"


def collect_lv_entries() -> list[dict]:
    """Walk JSON files; return raw entries with split metadata.

    We use cap (single-turn caption) + oe (open-ended QA) + mc (multi-choice QA)
    from the splits we downloaded.
    """
    base = RAW / "lv178k"
    files: list[Path] = sorted(base.glob("*/*_processed.json"))
    out: list[dict] = []
    for p in files:
        split = p.parent.name  # e.g. '30_60_s_academic_v0_1'
        name = p.name
        if "cap_processed" in name:
            qtype = "cap"
        elif "oe_v0_1_qa_processed" in name or "oe_qa" in name:
            qtype = "oe"
        elif "mc_v0_1_qa_processed" in name or "mc_qa" in name:
            qtype = "mc"
        else:
            continue
        try:
            data = json.loads(p.read_text())
        except Exception as e:
            print(f"[lv] skip {p}: {e}", flush=True)
            continue
        for entry in data:
            if not isinstance(entry, dict) or "conversations" not in entry:
                continue
            entry["_split"] = split
            entry["_qtype"] = qtype
            out.append(entry)
    print(f"[lv] collected {len(out)} raw entries across {len(files)} files", flush=True)
    return out


def load_lv_records(lv_manifest: dict[str, dict]) -> list[dict]:
    raw = collect_lv_entries()
    out: list[dict] = []
    skipped_missing = 0
    skipped_youtube = 0
    for entry in raw:
        split = entry["_split"]
        qtype = entry["_qtype"]
        # Skip youtube splits — we didn't download them in this pass.
        if "youtube" in split:
            skipped_youtube += 1
            continue
        rel_in_split = entry.get("video")
        if not rel_in_split:
            continue
        manifest_key = lv_rel_for_manifest(split, rel_in_split)
        if manifest_key not in lv_manifest:
            skipped_missing += 1
            continue
        abs_path = lv_video_abs(split, rel_in_split)
        conv = entry["conversations"]
        if len(conv) < 2 or conv[0]["from"] != "human":
            continue
        # First user/assistant pair only — simpler for OPD rollouts.
        msgs = [
            {"role": "user", "content": conv[0]["value"].replace("<image>", "<video>").strip()},
            {"role": "assistant", "content": conv[1]["value"].strip()},
        ]
        rec = {
            "data_source": f"lv178k_{qtype}",
            "ability": {"cap": "video_caption", "oe": "video_oe_qa", "mc": "video_mc_qa"}[qtype],
            "video_rel": manifest_key,
            "video_abs": str(abs_path),
            "messages": msgs,
            "extra_info": {
                "id": entry.get("id"),
                "split": split,
                "qtype": qtype,
                "manifest": lv_manifest[manifest_key],
            },
        }
        out.append(rec)
    print(f"[lv] kept {len(out)}  skipped_missing={skipped_missing}  skipped_youtube={skipped_youtube}",
          flush=True)
    return out


# ---------------------------------------------------------------------------
# Sampling + writing


def sample_balanced(records: list[dict], n: int, group_key, rng: random.Random) -> list[dict]:
    groups: dict[Any, list[dict]] = defaultdict(list)
    for r in records:
        groups[group_key(r)].append(r)
    # Round-robin sample to keep ability/category balance.
    for g in groups.values():
        rng.shuffle(g)
    keys = list(groups.keys())
    rng.shuffle(keys)
    out: list[dict] = []
    i = 0
    while len(out) < n:
        any_added = False
        for k in keys:
            if groups[k]:
                out.append(groups[k].pop())
                any_added = True
                if len(out) >= n:
                    break
        if not any_added:
            break
        i += 1
    return out


def to_verl_row(rec: dict, frames: int) -> dict:
    """Return a row in verl parquet schema."""
    msgs = rec["messages"]
    assert msgs[0]["role"] == "user" and msgs[-1]["role"] == "assistant"
    prompt_msgs = [m for m in msgs if m["role"] in ("system", "user")]
    response = msgs[-1]["content"]
    video_dict = make_video_dict(rec["video_abs"], rec["extra_info"]["manifest"], frames)
    return {
        "prompt": prompt_msgs,
        "response": response,
        "videos": [video_dict],
        "data_source": rec["data_source"],
        "ability": rec["ability"],
        "extra_info": {**rec["extra_info"], "video_rel": rec["video_rel"]},
    }


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--n_tb", type=int, default=1000)
    p.add_argument("--n_lv", type=int, default=4000)
    p.add_argument("--n_eval", type=int, default=100)
    p.add_argument("--frames", type=int, default=32)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--out_dir", default=str(OUT_DIR))
    args = p.parse_args()

    rng = random.Random(args.seed)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    print("[load] reading manifests...", flush=True)
    tb_manifest = load_manifest("tb")
    lv_manifest = load_manifest("lv178k")
    print(f"[load] tb manifest: {len(tb_manifest)}", flush=True)
    print(f"[load] lv manifest: {len(lv_manifest)}", flush=True)

    tb_records = load_tb_records(tb_manifest)
    lv_records = load_lv_records(lv_manifest)

    # 1k TB balanced across datasets.
    tb_sample = sample_balanced(
        tb_records, args.n_tb,
        group_key=lambda r: r["extra_info"]["dataset"], rng=rng,
    )
    # 4k LV178K balanced across (split, qtype).
    lv_sample = sample_balanced(
        lv_records, args.n_lv,
        group_key=lambda r: (r["extra_info"]["split"], r["extra_info"]["qtype"]),
        rng=rng,
    )
    print(f"[sample] tb={len(tb_sample)}  lv={len(lv_sample)}", flush=True)

    all_records = tb_sample + lv_sample
    rng.shuffle(all_records)

    # Held-out eval split (drawn from leftover, not from selected sample).
    used_ids = {(r["data_source"], r["video_rel"], r["messages"][0]["content"][:80])
                for r in all_records}
    leftover_pool = [
        r for r in (tb_records + lv_records)
        if (r["data_source"], r["video_rel"], r["messages"][0]["content"][:80]) not in used_ids
    ]
    rng.shuffle(leftover_pool)
    eval_records = leftover_pool[: args.n_eval]
    print(f"[sample] eval held-out: {len(eval_records)}", flush=True)

    # Write jsonl (with full messages including assistant) for SFT.
    jsonl_path = out_dir / "sft_5k.jsonl"
    with jsonl_path.open("w") as f:
        for r in all_records:
            f.write(json.dumps(r) + "\n")
    print(f"[write] {jsonl_path}  ({jsonl_path.stat().st_size / 1e6:.1f} MB)",
          flush=True)

    eval_path = out_dir / "eval_100.jsonl"
    with eval_path.open("w") as f:
        for r in eval_records:
            f.write(json.dumps(r) + "\n")
    print(f"[write] {eval_path}", flush=True)

    # Write parquet for verl OPD.
    rows = [to_verl_row(r, args.frames) for r in all_records]
    df = pd.DataFrame(rows)
    pq_path = out_dir / "sft_5k.parquet"
    df.to_parquet(pq_path, index=False)
    print(f"[write] {pq_path}  ({pq_path.stat().st_size / 1e6:.1f} MB)",
          flush=True)

    eval_rows = [to_verl_row(r, args.frames) for r in eval_records]
    pd.DataFrame(eval_rows).to_parquet(out_dir / "eval_100.parquet", index=False)

    # Print breakdown.
    by_src = Counter(r["data_source"] for r in all_records)
    print(f"[summary] sft_5k by data_source: {dict(by_src)}", flush=True)


if __name__ == "__main__":
    main()