"""Generate GPT-5 CoT beliefs for benchmark/v1 train/val/test ticks (parallel).

Reads:  benchmark/v1/data/{split}.parquet (tick-level records with frame_indices)
Writes: data/cot_corpus_v3/v1_{split}_perframe.jsonl

Schema (one record per tick, matches SFT trainer expectations):
  {
    "id": "v1_{split}_{i:06d}",
    "video_id": str,
    "video_path": str,
    "source": str,
    "category": str,
    "frame_indices": List[int][8],
    "actions_per_frame": List[str][8],   # SILENT/OBSERVE/ALERT
    "beliefs_per_frame": List[str][8],   # GPT-5 generated, ≤25 words each
    "danger_per_frame": List[float][8],  # derived from action label
    "tta_per_frame": List[float][8],
    "tick_action": str,
    "tick_tta_raw": float,
    "source_kind": "video_file" | "frame_folder",
    "hazard_category": str,              # GPT-5 generated
    "one_sentence_rationale": str,       # GPT-5 generated
    "gpt5_model": str,
    "in_tokens": int, "out_tokens": int, "cost_usd": float,
  }

Cost cap enforced via shared ledger. Resume-on-failure via output-file scan.

Usage:
  python tools/run_v1_gpt5_cot.py --split val --parallel 16 --max_cost_usd 200
  python tools/run_v1_gpt5_cot.py --split train --parallel 16 --max_cost_usd 1500
"""
from __future__ import annotations
import argparse
import base64
import hashlib
import io
import json
import logging
import os
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Dict, List, Optional

import cv2
import numpy as np
from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("v1_gpt5_cot")

KEY_PATH = Path("~/Desktop/openai_api_key.txt")

# Per-million-token costs (same as tools/vlalert_x_distill.py)
COSTS = {
    "gpt-5.5": (5.00, 15.00),
    "gpt-5.4": (3.00,  9.00),
    "gpt-5":   (2.50,  7.50),
    "gpt-4o":  (2.50, 10.00),
}

PROMPT = """You are a safety analyst labelling an 8-frame dashcam montage
(2 rows x 4 cols, left-to-right then top-to-bottom = frame 0..7, last
frame is most recent). Output a strict JSON record.

Output schema (no extras, no missing keys):
{
  "hazard_category": one of [pedestrian, vrurider, vehicle_cross,
        vehicle_oncoming, vehicle_lead, weather, infrastructure, none],
  "per_frame_belief": [
      {"frame": 0, "belief": "<=25-word phrase describing the scene
            and threat status visible in this frame"},
      ... (exactly 8 entries, frames 0..7)
  ],
  "one_sentence_rationale": "<=25-word summary of the risk evolution"
}

Rules:
- The clip's outcome is unknown -- judge from visual evidence only.
- Each `belief` must be a *phrase*, not a full sentence with a period.
- Use simple physical descriptors (vehicle position, motion cue,
  conflict sign), avoid temporal claims like "will collide".
- If the scene is benign, use `hazard_category: none` and briefly note
  the dominant safe-driving cue per frame.
"""


# ─────────────── shared cost ledger (thread-safe) ───────────────
class CostLedger:
    def __init__(self, path: Path, model: str, max_cost_usd: float):
        self.path = path
        self.model = model
        self.max_cost_usd = max_cost_usd
        self.lock = threading.Lock()
        if path.exists():
            d = json.loads(path.read_text())
            self.n_calls = d.get("n_calls", 0)
            self.cost_usd = d.get("cost_usd", 0.0)
            self.in_tokens = d.get("in_tokens", 0)
            self.out_tokens = d.get("out_tokens", 0)
        else:
            self.n_calls = 0; self.cost_usd = 0.0
            self.in_tokens = 0; self.out_tokens = 0

    def can_spend(self, projected_usd: float) -> bool:
        with self.lock:
            return self.cost_usd + projected_usd <= self.max_cost_usd

    def add(self, in_tok: int, out_tok: int):
        cin, cout = COSTS.get(self.model, (5.0, 15.0))
        cost = (in_tok / 1e6) * cin + (out_tok / 1e6) * cout
        with self.lock:
            self.n_calls += 1
            self.in_tokens += in_tok
            self.out_tokens += out_tok
            self.cost_usd += cost
            self.path.parent.mkdir(parents=True, exist_ok=True)
            self.path.write_text(json.dumps({
                "primary_model": self.model,
                "n_calls": self.n_calls,
                "in_tokens": self.in_tokens,
                "out_tokens": self.out_tokens,
                "cost_usd": self.cost_usd,
            }, indent=2))
        return cost


# ─────────────── frame extraction + montage ───────────────
def _load_frames(video_path: str, frame_indices: List[int],
                 size: int = 256) -> Optional[List[Image.Image]]:
    """Load 8 frames as resized PIL images."""
    p = Path(video_path)
    if p.suffix.lower() == ".mp4" and p.exists():
        cap = cv2.VideoCapture(str(p))
        frames = []
        for fi in frame_indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(fi))
            ok, fr = cap.read()
            if not ok: return None
            fr = cv2.cvtColor(fr, cv2.COLOR_BGR2RGB)
            fr = cv2.resize(fr, (size, size), interpolation=cv2.INTER_AREA)
            frames.append(Image.fromarray(fr))
        cap.release()
        return frames if len(frames) == len(frame_indices) else None
    elif p.is_dir():
        frames = []
        for fi in frame_indices:
            for w in (3, 4, 5, 6):
                fp = p / f"{int(fi):0{w}d}.jpg"
                if fp.exists():
                    img = Image.open(fp).convert("RGB")
                    img.thumbnail((size, size))
                    frames.append(img); break
            else:
                fp = p / "images" / f"{int(fi):06d}.jpg"
                if fp.exists():
                    img = Image.open(fp).convert("RGB")
                    img.thumbnail((size, size))
                    frames.append(img)
                else:
                    return None
        return frames if len(frames) == len(frame_indices) else None
    return None


def _build_montage(frames: List[Image.Image], cell: int = 224) -> Image.Image:
    """2 rows x 4 cols, return PIL."""
    canvas = Image.new("RGB", (cell * 4, cell * 2), (0, 0, 0))
    for i, im in enumerate(frames):
        r, c = i // 4, i % 4
        im_r = im.resize((cell, cell))
        canvas.paste(im_r, (c * cell, r * cell))
    return canvas


def _pil_to_data_url(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/jpeg;base64,{b64}"


# ─────────────── GPT-5 call ───────────────
def _call_gpt5(client, montage: Image.Image, model: str,
               max_retries: int = 3) -> Optional[Dict]:
    url = _pil_to_data_url(montage)
    last_err = None
    for attempt in range(max_retries):
        try:
            resp = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": PROMPT},
                    {"role": "user", "content": [
                        {"type": "image_url",
                         "image_url": {"url": url, "detail": "low"}},
                        {"type": "text", "text":
                         "Analyze the 8-frame montage and output the JSON."},
                    ]},
                ],
                max_completion_tokens=3000,
                response_format={"type": "json_object"},
            )
            text = resp.choices[0].message.content
            if not text or not text.strip():
                last_err = f"empty response (finish_reason={resp.choices[0].finish_reason})"
                if attempt < max_retries - 1:
                    time.sleep(1.0)
                continue
            data = json.loads(text)
            return {
                "data": data,
                "in_tokens": resp.usage.prompt_tokens,
                "out_tokens": resp.usage.completion_tokens,
                "model": resp.model,
            }
        except Exception as e:
            last_err = str(e)
            if attempt < max_retries - 1:
                time.sleep(2.0 * (attempt + 1))
    logger.warning(f"GPT call failed after {max_retries} retries: {last_err}")
    return None


# ─────────────── per-tick worker ───────────────
def _process_tick(rec: Dict, client, model: str, ledger: CostLedger) -> Optional[Dict]:
    if not ledger.can_spend(0.02):
        return {"skip_reason": "budget_cap"}
    frames = _load_frames(rec["video_path"], rec["frame_indices"])
    if frames is None or len(frames) != 8:
        return {"skip_reason": "frame_load_failed"}
    montage = _build_montage(frames)
    result = _call_gpt5(client, montage, model)
    if result is None:
        return {"skip_reason": "gpt_failed"}
    ledger.add(result["in_tokens"], result["out_tokens"])
    # Extract per-frame beliefs (defensive: handle multiple formats)
    pf = result["data"].get("per_frame_belief", [])
    beliefs = [""] * 8
    for i, entry in enumerate(pf):
        if isinstance(entry, dict):
            f = entry.get("frame", i)
            b = entry.get("belief", "")
        elif isinstance(entry, str):
            f, b = i, entry  # GPT returned plain string array
        else:
            continue
        try:
            f = int(f)
        except (TypeError, ValueError):
            f = i
        if 0 <= f < 8:
            beliefs[f] = str(b) if b else ""
    # Output record (compatible with SFT trainer)
    out = {
        "id":               rec["id"],
        "video_id":         rec["video_id"],
        "video_path":       rec["video_path"],
        "source":           rec["source"],
        "category":         rec["category"],
        "frame_indices":    rec["frame_indices"],
        "actions_per_frame": rec["actions_per_frame"],
        "beliefs_per_frame": beliefs,
        "danger_per_frame":  rec["danger_per_frame"],
        "tta_per_frame":     rec["tta_per_frame"],
        "tick_action":      rec["tick_action"],
        "tick_tta_raw":     rec["tick_tta_raw"],
        "source_kind":      rec["source_kind"],
        "hazard_category":  result["data"].get("hazard_category", "none"),
        "one_sentence_rationale": result["data"].get("one_sentence_rationale", ""),
        "gpt5_model":       result["model"],
        "in_tokens":        result["in_tokens"],
        "out_tokens":       result["out_tokens"],
    }
    return out


# ─────────────── main ───────────────
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--split", required=True, choices=["train", "val", "test", "all"])
    ap.add_argument("--parallel", type=int, default=16)
    ap.add_argument("--max_cost_usd", type=float, default=200.0)
    ap.add_argument("--model", default="gpt-4o")
    ap.add_argument("--max_priority", type=int, default=7,
                    help="Skip ticks with priority > this (1=highest, 99=skip)")
    ap.add_argument("--limit", type=int, default=0,
                    help="cap n samples (for smoke test)")
    args = ap.parse_args()

    # Prefer priority-sorted manifest from cot_corpus_v3; fallback to v2.
    pri_jsonl = ROOT / f"data/cot_corpus_v3/v1_{args.split}_priority.jsonl"
    base_jsonl = ROOT / f"data/cot_corpus_v2/v1_{args.split}_perframe.jsonl"
    src_jsonl = pri_jsonl if pri_jsonl.exists() else base_jsonl
    if not src_jsonl.exists():
        logger.error(f"Source jsonl not found: {src_jsonl}")
        return
    records = []
    n_skip_pri = 0
    with src_jsonl.open() as f:
        for line in f:
            line = line.strip()
            if not line: continue
            r = json.loads(line)
            pri = r.get("priority", 99)
            if pri > args.max_priority:
                n_skip_pri += 1
                continue
            records.append(r)
    if args.limit:
        records = records[:args.limit]
    logger.info(f"[load] {len(records):,} records from {src_jsonl}, "
                f"skipped {n_skip_pri} (priority > {args.max_priority})")

    out_path = ROOT / f"data/cot_corpus_v3/v1_{args.split}_perframe.jsonl"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    ledger_path = ROOT / f"eval_results/openai_teacher/v1_gpt5_{args.split}_ledger.json"

    # Resume: scan existing output for completed IDs
    seen = set()
    if out_path.exists():
        with out_path.open() as f:
            for line in f:
                try:
                    d = json.loads(line)
                    if "id" in d:
                        seen.add(d["id"])
                except Exception:
                    pass
        logger.info(f"[resume] skipping {len(seen):,} already-done")
    todo = [r for r in records if r["id"] not in seen]
    logger.info(f"[plan] {len(todo):,} ticks to generate, "
                f"max cost ${args.max_cost_usd}, parallel={args.parallel}")

    # Init OpenAI
    os.environ["OPENAI_API_KEY"] = KEY_PATH.read_text().strip()
    from openai import OpenAI
    client = OpenAI()

    ledger = CostLedger(ledger_path, args.model, args.max_cost_usd)
    logger.info(f"[ledger] start cost=${ledger.cost_usd:.3f} "
                f"of ${args.max_cost_usd}")

    n_done, n_failed, n_skipped_budget = 0, 0, 0
    t0 = time.time()
    out_lock = threading.Lock()

    with ThreadPoolExecutor(max_workers=args.parallel) as ex:
        futures = {ex.submit(_process_tick, rec, client, args.model, ledger): rec
                   for rec in todo}
        for fut in as_completed(futures):
            try:
                res = fut.result()
            except Exception as e:
                logger.warning(f"  [worker crash] {e}")
                n_failed += 1
                continue
            if res is None:
                n_failed += 1
                continue
            if "skip_reason" in res:
                if res["skip_reason"] == "budget_cap":
                    n_skipped_budget += 1
                    # Cancel remaining
                    for f in futures: f.cancel()
                    break
                else:
                    n_failed += 1
                    continue
            with out_lock:
                with out_path.open("a") as f:
                    f.write(json.dumps(res) + "\n")
            n_done += 1
            if n_done % 50 == 0:
                el = time.time() - t0
                rate = n_done / max(el, 1e-9)
                logger.info(f"  done={n_done}, failed={n_failed}, "
                             f"cost=${ledger.cost_usd:.2f}, "
                             f"rate={rate:.1f}/s, "
                             f"eta={(len(todo) - n_done) / max(rate, 1e-9) / 60:.0f}min")

    logger.info(f"\n[final] done={n_done}, failed={n_failed}, "
                 f"skipped_budget={n_skipped_budget}, cost=${ledger.cost_usd:.2f}")


if __name__ == "__main__":
    main()