"""Zero-shot scoring on benchmark/v1/val using Gemini Flash Lite.

Uses the cheapest production multimodal model (gemini-2.0-flash-lite) to emit a
3-class action label (SILENT/OBSERVE/ALERT) + a [0,1] danger score for each
8-frame tick. Output matches the per_tick PT schema produced by
tools/score_v1_val_baselines.py so the existing aggregators auto-include it.

Cost (val only, ~11,220 ticks):
  per tick ≈ 8 images @ ~258 image tokens + ~120 prompt + ~30 output
           ≈ 2.2k input tokens + 30 output tokens
           ≈ $0.00025  (Flash-Lite: $0.075/1M input + $0.30/1M output)
  full split ≈ $2.80   (hard cap $5.00 — exits early if exceeded)

Usage:
    GEMINI_API_KEY=$(cat ~/Desktop/GEMINI_API.txt) \
        python tools/score_v1_val_gemini.py [--max_ticks N] [--workers 10]

Resumable: re-running skips ticks already in the sha256 cache.
"""
from __future__ import annotations
import argparse
import base64
import hashlib
import io
import json
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import cv2
import numpy as np
import torch
from PIL import Image
from tqdm import tqdm

from google import genai
from google.genai import types as genai_types

ROOT = Path(__file__).resolve().parents[1]
MANIFEST   = ROOT / "eval_results/benchmark_v1_val/val_manifest.json"
OUT_PT     = ROOT / "eval_results/benchmark_v1_val/per_tick/gemini_zeroshot.pt"
CACHE_DIR  = ROOT / "eval_results/benchmark_v1_val/gemini_cache"
COST_FILE  = ROOT / "eval_results/benchmark_v1_val/gemini_cost.json"
LOG_FILE   = ROOT / "logs/v4/gemini_score.log"

MODEL_NAME = "gemini-2.5-flash-lite"   # cheapest current production model
FRAME_SIZE = 256                       # downscale frames for token efficiency
PRICE_IN   = 0.10 / 1_000_000          # USD/token  (2.5-flash-lite)
PRICE_OUT  = 0.40 / 1_000_000
HARD_CAP   = 5.00                      # USD

ACTION_MAP = {"SILENT": 0, "OBSERVE": 1, "ALERT": 2}

PROMPT = (
    "You are a driving-safety system. You see 8 consecutive dashcam frames in "
    "temporal order from an ego vehicle. Decide whether a collision or critical "
    "hazard is about to occur within the next ~2 seconds. "
    "Output STRICT JSON ONLY in this exact schema:\n"
    '{"action": "SILENT" | "OBSERVE" | "ALERT", "danger": <float 0..1>}\n'
    "Definitions: SILENT = normal driving, no hazard. OBSERVE = potential "
    "hazard developing (2-4 s out). ALERT = imminent collision risk (< 2 s). "
    "Return ONLY the JSON, no prose."
)


# ──────────────────────── frame loading ──────────────────────────

def load_frames_for_sample(sample: dict) -> list[bytes]:
    """Return a list of 8 JPEG-encoded frame bytes."""
    src = sample.get("source_dir", "")
    fis = sample.get("frame_indices", [])[:8]
    p = Path(src)

    frames_rgb = []
    if p.suffix.lower() in (".mp4", ".avi") and p.exists():
        cap = cv2.VideoCapture(str(p))
        for fi in fis:
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(fi))
            ok, fr = cap.read()
            if ok:
                frames_rgb.append(cv2.cvtColor(fr, cv2.COLOR_BGR2RGB))
            else:
                frames_rgb.append(np.zeros((FRAME_SIZE, FRAME_SIZE, 3), np.uint8))
        cap.release()
    else:
        # Frame-folder path (DOTA / DAD / DADA)
        search_dirs = [p, p / "images"]
        for fi in fis:
            arr = None
            for d in search_dirs:
                if not d.is_dir():
                    continue
                for w in (3, 4, 5, 6):
                    fp = d / f"{int(fi):0{w}d}.jpg"
                    if fp.exists():
                        arr = np.array(Image.open(fp).convert("RGB"))
                        break
                if arr is not None:
                    break
            frames_rgb.append(arr if arr is not None
                              else np.zeros((FRAME_SIZE, FRAME_SIZE, 3), np.uint8))

    # Resize + JPEG encode
    out = []
    for fr in frames_rgb:
        h, w = fr.shape[:2]
        s = min(h, w)
        sq = fr[(h - s) // 2:(h - s) // 2 + s, (w - s) // 2:(w - s) // 2 + s]
        sq = cv2.resize(sq, (FRAME_SIZE, FRAME_SIZE), interpolation=cv2.INTER_AREA)
        buf = io.BytesIO()
        Image.fromarray(sq).save(buf, format="JPEG", quality=80)
        out.append(buf.getvalue())
    return out


# ──────────────────────── Gemini call ────────────────────────────

def parse_response(text: str) -> tuple[str, float, bool]:
    """Return (action_str, danger_float, ok)."""
    t = text.strip()
    if t.startswith("```"):
        t = t.strip("`").lstrip("json").strip()
    # Try strict JSON first
    try:
        d = json.loads(t)
        a = str(d.get("action", "")).upper().strip()
        if a not in ACTION_MAP:
            # Try keyword match
            for k in ACTION_MAP:
                if k in a:
                    a = k; break
        if a not in ACTION_MAP:
            return "SILENT", 0.05, False
        dv = float(d.get("danger", 0.05))
        dv = max(0.0, min(1.0, dv))
        return a, dv, True
    except Exception:
        # Fallback keyword scan
        T = t.upper()
        for k in ("ALERT", "OBSERVE", "SILENT"):
            if k in T:
                # Crude danger inference
                dv = {"ALERT": 0.9, "OBSERVE": 0.5, "SILENT": 0.05}[k]
                return k, dv, False
        return "SILENT", 0.05, False


def call_gemini(client, frame_bytes: list[bytes],
                 max_retries: int = 5) -> tuple[str, str, float, bool, dict]:
    """Return (raw_text, action, danger, ok, usage)."""
    parts = [genai_types.Part.from_text(text=PROMPT)]
    for jpg in frame_bytes:
        parts.append(genai_types.Part.from_bytes(data=jpg, mime_type="image/jpeg"))
    contents = [genai_types.Content(role="user", parts=parts)]

    for attempt in range(max_retries):
        try:
            resp = client.models.generate_content(
                model=MODEL_NAME,
                contents=contents,
                config=genai_types.GenerateContentConfig(
                    temperature=0.0,
                    max_output_tokens=80,
                    response_mime_type="application/json",
                ),
            )
            text = resp.text or ""
            action, danger, ok = parse_response(text)
            um = resp.usage_metadata
            usage = {
                "input": int(um.prompt_token_count or 0),
                "output": int(um.candidates_token_count or 0),
            }
            return text, action, danger, ok, usage
        except Exception as e:
            msg = str(e).lower()
            if "429" in msg or "quota" in msg or "rate" in msg or "503" in msg:
                time.sleep(2 ** attempt)
                continue
            return f"ERROR: {e}", "SILENT", 0.05, False, {"input": 0, "output": 0}
    return "ERROR: max retries", "SILENT", 0.05, False, {"input": 0, "output": 0}


# ──────────────────────── orchestrator ───────────────────────────

def score_one(client, sample: dict, idx: int) -> dict:
    sid = sample.get("video_id", f"sample_{idx}")
    tick = sample.get("tick_idx", 0)
    cache_key = hashlib.sha256(
        f"{sid}_{tick}_{sample['frame_indices'][0]}_{MODEL_NAME}".encode()
    ).hexdigest()[:24]
    cache_fp = CACHE_DIR / f"{cache_key}.json"
    if cache_fp.exists():
        try:
            cached = json.loads(cache_fp.read_text())
            cached["from_cache"] = True
            return cached
        except Exception:
            pass

    frames = load_frames_for_sample(sample)
    text, action, danger, ok, usage = call_gemini(client, frames)
    result = {
        "idx": idx, "sid": sid, "tick_idx": tick,
        "raw_text": text, "action_str": action, "danger": danger,
        "ok": ok, "usage": usage, "from_cache": False,
    }
    cache_fp.write_text(json.dumps(result))
    return result


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--max_ticks", type=int, default=0,
                    help="0 = all ticks; >0 = smoke test with this many")
    ap.add_argument("--workers",   type=int, default=10)
    ap.add_argument("--api_key_file", type=Path,
                    default=Path("~/Desktop/GEMINI_API.txt"))
    ap.add_argument("--cost_cap", type=float, default=HARD_CAP)
    args = ap.parse_args()

    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    LOG_FILE.parent.mkdir(parents=True, exist_ok=True)

    api_key = args.api_key_file.read_text().strip()
    client = genai.Client(api_key=api_key)
    print(f"[init] model={MODEL_NAME}  workers={args.workers}  cap=${args.cost_cap}")

    samples = json.loads(MANIFEST.read_text())["samples"]
    if args.max_ticks > 0:
        samples = samples[:args.max_ticks]
    N = len(samples)
    print(f"[load] {N} ticks from {MANIFEST.name}")

    results = [None] * N
    total_in = total_out = 0
    cost_so_far = 0.0
    n_done = n_ok = n_cache = 0
    stop_flag = {"v": False}

    def worker(i):
        if stop_flag["v"]:
            return None
        return score_one(client, samples[i], i)

    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(worker, i): i for i in range(N)}
        pbar = tqdm(total=N, ncols=100, desc="gemini")
        for fut in as_completed(futs):
            i = futs[fut]
            try:
                r = fut.result()
            except Exception as e:
                print(f"[err] tick {i}: {e}")
                r = None
            if r is None:
                pbar.update(1); continue
            results[i] = r
            n_done += 1
            if r["ok"]:
                n_ok += 1
            if r.get("from_cache"):
                n_cache += 1
            u = r.get("usage", {})
            total_in  += u.get("input", 0)
            total_out += u.get("output", 0)
            cost_so_far = total_in * PRICE_IN + total_out * PRICE_OUT
            if not r.get("from_cache") and cost_so_far > args.cost_cap:
                print(f"\n[STOP] cost cap reached: ${cost_so_far:.3f} > ${args.cost_cap}")
                stop_flag["v"] = True
            pbar.set_postfix({
                "ok": f"{n_ok}/{n_done}",
                "cache": n_cache,
                "$": f"{cost_so_far:.3f}",
            })
            pbar.update(1)
        pbar.close()

    # ────────── persist cost ──────────
    COST_FILE.write_text(json.dumps({
        "model": MODEL_NAME,
        "input_tokens": total_in,
        "output_tokens": total_out,
        "cost_usd": cost_so_far,
        "n_ticks": N, "n_done": n_done, "n_ok": n_ok,
    }, indent=2))
    print(f"\n[cost] ${cost_so_far:.4f}  in={total_in:,}  out={total_out:,}")
    print(f"[done] {n_done}/{N} ticks  ({n_ok} parsed OK, {n_cache} from cache)")

    # ────────── build per_tick PT in baseline schema ──────────
    raw_logits = torch.zeros(N, 3, dtype=torch.float32)
    scores3    = torch.zeros(N, 3, dtype=torch.float32)
    scores_bin = torch.zeros(N, dtype=torch.float32)
    actions_str = []
    raw_texts  = []
    tick_labels = torch.zeros(N, dtype=torch.long)
    tta_raw     = torch.zeros(N, dtype=torch.float32)
    frame_indices = torch.zeros(N, 8, dtype=torch.long)
    fps_tensor    = torch.zeros(N, dtype=torch.float32)
    ids, sources, categories, raw_categories, tick_idxs = [], [], [], [], []

    for i, s in enumerate(samples):
        ids.append(s.get("video_id", ""))
        sources.append(s.get("source", ""))
        categories.append(s.get("category", ""))
        raw_categories.append(s.get("raw_category", ""))
        tick_idxs.append(s.get("tick_idx", 0))
        tick_labels[i] = int(s.get("action_label", 0))
        tta_raw[i] = float(s.get("tta_raw", -1.0))
        fis = s.get("frame_indices", [])[:8]
        if len(fis) < 8: fis = fis + [fis[-1] if fis else 0] * (8 - len(fis))
        frame_indices[i] = torch.tensor(fis, dtype=torch.long)
        fps_tensor[i] = float(s.get("fps", 30.0))

        r = results[i]
        if r is None:
            actions_str.append("SILENT")
            raw_texts.append("MISSING")
            scores3[i]    = torch.tensor([0.85, 0.10, 0.05])
            scores_bin[i] = 0.05
            raw_logits[i] = torch.tensor([0.85, 0.10, 0.05]).log()
            continue
        a = r["action_str"]
        d = float(r["danger"])
        actions_str.append(a)
        raw_texts.append(r["raw_text"][:200])
        # 3-class soft: put 0.85 on chosen class, split rest based on danger
        soft = torch.full((3,), (1 - 0.85) / 2)
        soft[ACTION_MAP[a]] = 0.85
        # Blend danger into ALERT prob: scores_3class[i,2] gets danger value
        soft[2] = max(soft[2].item(), d * 0.9)
        soft = soft / soft.sum()
        scores3[i] = soft
        scores_bin[i] = d
        raw_logits[i] = soft.log()

    out = {
        "method": "gemini_flash_lite_zeroshot",
        "model": MODEL_NAME,
        "manifest": str(MANIFEST),
        "n_ticks": N,
        "ids": ids, "source": sources,
        "category": categories, "raw_category": raw_categories,
        "frame_indices": frame_indices, "tta_raw": tta_raw,
        "fps": fps_tensor, "n_frames": torch.full((N,), 8, dtype=torch.long),
        "tick_idx": torch.tensor(tick_idxs, dtype=torch.long),
        "tick_label": tick_labels,
        "raw_logits": raw_logits,
        "scores_3class": scores3,
        "scores_binary": scores_bin,
        # extras for case-study debugging
        "gemini_raw_text": raw_texts,
        "gemini_action_str": actions_str,
        "cost_usd": cost_so_far,
    }
    OUT_PT.parent.mkdir(parents=True, exist_ok=True)
    torch.save(out, OUT_PT)
    print(f"[save] {OUT_PT}")


if __name__ == "__main__":
    sys.exit(main())