"""
Generate structured CoT labels for Nexar train clips using GPT-4o as teacher.

Schema (enforced by response_format json_object):
{
  "scene": "short description of the driving scene",
  "critical_objects": ["list", "of", "hazardous agents"],
  "threat_analysis": "short reasoning about what could collide and when",
  "verdict": "yes" | "no",
  "confidence": integer 0-100
}

Usage:
  export OPENAI_API_KEY=$(cat ~/Desktop/openai_api_key.txt)
  python -m training.VLA.build_cot_labels \
      --train_csv nexar-collision-prediction/train.csv \
      --video_dir nexar-collision-prediction/train \
      --out data/vla_cot/train_cot.jsonl \
      --n_clips 30 --n_frames 8
"""
from __future__ import annotations

import argparse
import json
import os
import random
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import pandas as pd
from tqdm import tqdm

try:
    from openai import OpenAI
except ImportError as e:
    raise SystemExit("pip install openai") from e

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from training.VLA.frame_utils import pil_to_data_url, sample_frames_from_mp4

SYSTEM_PROMPT = (
    "You are a senior driving-safety analyst reviewing dashcam footage. "
    "You will see 8 uniformly-sampled frames from a short clip and a ground-truth label "
    "indicating whether the clip ends in a collision or near-collision. "
    "Produce a concise chain-of-thought in strict JSON with this exact schema:\n"
    '{\n'
    '  "scene": "<=25-word scene description (road type, weather, lighting, traffic)",\n'
    '  "critical_objects": ["each item is an agent/object that matters, <=6 words, max 4 items"],\n'
    '  "threat_analysis": "<=40-word reasoning on kinematics and likely collision path",\n'
    '  "verdict": "yes" or "no",\n'
    '  "confidence": integer 0-100\n'
    '}\n'
    "Rules:\n"
    "- verdict MUST match the ground-truth label.\n"
    "- Be specific and grounded — reference colors, positions, and motions actually visible.\n"
    "- NEVER say \"based on the label\"; write as if you inferred yourself.\n"
    "- Output JSON only, no prose."
)

USER_TEMPLATE = (
    "Ground-truth label for this clip: collision = {label}.\n"
    "Analyze the 8 frames (earliest → latest, left-to-right) and output the JSON."
)


def build_messages(label: int, frames, detail: str = "low"):
    content = []
    for img in frames:
        content.append(
            {"type": "image_url", "image_url": {"url": pil_to_data_url(img), "detail": detail}}
        )
    label_word = "YES" if label == 1 else "NO"
    content.append({"type": "text", "text": USER_TEMPLATE.format(label=label_word)})
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": content},
    ]


def call_gpt4o(client, clip_id, label, frames, model: str, max_retries: int = 3, detail: str = "low"):
    messages = build_messages(label, frames, detail=detail)
    last_err = None
    for attempt in range(max_retries):
        try:
            resp = client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=0.2,
                max_tokens=350,
                response_format={"type": "json_object"},
                timeout=60,
            )
            raw = resp.choices[0].message.content
            parsed = json.loads(raw)
            # minimal schema validation
            assert parsed.get("verdict") in ("yes", "no"), "bad verdict"
            assert isinstance(parsed.get("confidence"), (int, float)), "bad confidence"
            assert (parsed["verdict"] == "yes") == (label == 1), "verdict/label mismatch"
            return {"id": clip_id, "label": int(label), "cot": parsed, "usage": resp.usage.model_dump() if resp.usage else None}
        except (json.JSONDecodeError, AssertionError, Exception) as e:  # noqa: BLE001
            last_err = e
            if attempt + 1 < max_retries:
                time.sleep(2 * (attempt + 1))
    return {"id": clip_id, "label": int(label), "cot": None, "error": str(last_err)}


def load_done(out_path: Path):
    done = set()
    if out_path.exists():
        with out_path.open() as f:
            for line in f:
                try:
                    rec = json.loads(line)
                    if rec.get("cot") is not None:
                        done.add(str(rec["id"]))
                except Exception:
                    continue
    return done


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--train_csv", required=True)
    ap.add_argument("--video_dir", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--n_clips", type=int, default=30, help="total clips (balanced pos/neg)")
    ap.add_argument("--n_frames", type=int, default=8)
    ap.add_argument("--resize_short", type=int, default=336)
    ap.add_argument("--model", default="gpt-4o")
    ap.add_argument("--detail", default="low", choices=["low", "high", "auto"])
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--skip_ids", default=None, help="comma-separated IDs to exclude (e.g. eval split)")
    args = ap.parse_args()

    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        raise SystemExit("Set OPENAI_API_KEY (source the file into env)")

    df = pd.read_csv(args.train_csv, dtype={"id": str})
    df["id"] = df["id"].astype(str).str.zfill(5)

    skip = set()
    if args.skip_ids:
        skip = set(s.strip().zfill(5) for s in args.skip_ids.split(",") if s.strip())
    df = df[~df["id"].isin(skip)]

    rng = random.Random(args.seed)
    pos = df[df["target"] == 1]["id"].tolist()
    neg = df[df["target"] == 0]["id"].tolist()
    rng.shuffle(pos); rng.shuffle(neg)
    half = args.n_clips // 2
    picked = pos[:half] + neg[:args.n_clips - half]
    rng.shuffle(picked)

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    done = load_done(out_path)
    todo = [pid for pid in picked if pid not in done]
    print(f"[cot] picked={len(picked)} already_done={len(done)} todo={len(todo)}")

    video_dir = Path(args.video_dir)
    client = OpenAI(api_key=api_key)

    def worker(pid):
        label = int(df[df["id"] == pid]["target"].iloc[0])
        video_path = video_dir / f"{pid}.mp4"
        if not video_path.exists():
            return {"id": pid, "label": label, "cot": None, "error": "missing_mp4"}
        try:
            frames = sample_frames_from_mp4(video_path, n_frames=args.n_frames, resize_short=args.resize_short)
        except Exception as e:  # noqa: BLE001
            return {"id": pid, "label": label, "cot": None, "error": f"frame_err:{e}"}
        return call_gpt4o(client, pid, label, frames, model=args.model, detail=args.detail)

    total_tokens_in = 0
    total_tokens_out = 0
    n_ok = 0
    n_err = 0
    with out_path.open("a") as fout, ThreadPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(worker, pid): pid for pid in todo}
        for fut in tqdm(as_completed(futs), total=len(futs), desc="cot"):
            rec = fut.result()
            fout.write(json.dumps(rec) + "\n")
            fout.flush()
            if rec.get("cot") is not None:
                n_ok += 1
                u = rec.get("usage") or {}
                total_tokens_in += u.get("prompt_tokens", 0)
                total_tokens_out += u.get("completion_tokens", 0)
            else:
                n_err += 1

    print(f"[cot] ok={n_ok} err={n_err} prompt_tokens={total_tokens_in} compl_tokens={total_tokens_out}")
    # gpt-4o pricing (2026-04): $2.5/M in, $10/M out
    est_usd = total_tokens_in * 2.5 / 1e6 + total_tokens_out * 10 / 1e6
    print(f"[cot] est cost: ${est_usd:.4f}")


if __name__ == "__main__":
    main()