#!/usr/bin/env python3
from __future__ import annotations

import argparse
import base64
import csv
import json
import re
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any

import requests


SYSTEM = "You are a strict visual memory evaluator. Return only valid JSON."

USER_PROMPT = """Evaluate whether the revisit tail frames preserve the same object and scene as the first frame.

The most important target is the toy bear if it appears in the first frame. Judge its appearance consistency strongly:
- bear identity, color, shape, face, pose, clothing/accessories if visible
- whether the final revisit still depicts the same bear after camera motion

Scene/layout consistency is secondary. Give lower weight to background unless it contradicts the target.

Return JSON with:
{
  "bear_appearance_score": 0-5,
  "bear_presence_score": 0-5,
  "scene_consistency_score": 0-5,
  "view_revisit_score": 0-5,
  "overall_score": 0-100,
  "verdict": "pass|partial|fail",
  "reason": "short evidence"
}

Weights for overall_score: bear_appearance 45%, bear_presence 25%, view_revisit 20%, scene_consistency 10%.
"""


def load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def encode_image(path: str | Path) -> str:
    return base64.b64encode(Path(path).read_bytes()).decode("ascii")


def parse_json(text: str) -> dict[str, Any]:
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip().strip("`")
    match = re.search(r"\{.*\}", text, flags=re.DOTALL)
    if not match:
        raise ValueError(f"no JSON in VLM response: {text[:300]}")
    return json.loads(match.group(0))


def wait_ready(api_base: str, timeout: int) -> None:
    url = f"{api_base.rstrip('/')}/models"
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            if requests.get(url, timeout=5).status_code == 200:
                return
        except Exception:
            pass
        time.sleep(2)
    raise RuntimeError(f"VLM server not ready at {api_base}")


def score_one(metrics_path: Path, api_base: str, model: str, timeout: int, force: bool) -> Path:
    out_path = metrics_path.parent / "vlm_score.json"
    if out_path.is_file() and not force:
        return out_path

    metrics = load_json(metrics_path)
    images = list(metrics.get("first_frame_paths") or []) + list(metrics.get("revisit_tail_paths") or [])
    content: list[dict[str, Any]] = [{"type": "text", "text": USER_PROMPT}]
    for image_path in images:
        content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{encode_image(image_path)}"}})

    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": [{"type": "text", "text": SYSTEM}]},
            {"role": "user", "content": content},
        ],
        "temperature": 0.0,
        "max_tokens": 800,
        "chat_template_kwargs": {"enable_thinking": False},
    }
    response = requests.post(f"{api_base.rstrip('/')}/chat/completions", json=payload, timeout=timeout)
    response.raise_for_status()
    raw = response.json()["choices"][0]["message"]["content"]
    vlm = parse_json(raw)

    out = {
        "metrics_path": str(metrics_path),
        "run_id": metrics.get("run_id"),
        "domain": (metrics.get("sample") or {}).get("domain"),
        "sample_id": (metrics.get("sample") or {}).get("sample_id"),
        "mode": metrics.get("mode"),
        "traditional": {
            "closure_psnr": metrics.get("closure_psnr"),
            "closure_ssim": metrics.get("closure_ssim"),
            "closure_mse": metrics.get("closure_mse"),
        },
        "vlm": vlm,
    }
    out_path.write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8")
    return out_path


def maybe_score_vlm(metrics_paths: list[Path], args: argparse.Namespace) -> None:
    if not args.score_vlm:
        return
    wait_ready(args.vlm_api_base, args.startup_wait_sec)
    todo = [p for p in metrics_paths if args.force_vlm or not (p.parent / "vlm_score.json").is_file()]
    print(f"[vlm] scoring {len(todo)} / {len(metrics_paths)} cases with workers={args.vlm_workers}")
    if not todo:
        return
    with ThreadPoolExecutor(max_workers=args.vlm_workers) as pool:
        futures = [pool.submit(score_one, p, args.vlm_api_base, args.vlm_model, args.vlm_timeout, args.force_vlm) for p in todo]
        for i, fut in enumerate(as_completed(futures), 1):
            path = fut.result()
            print(f"[vlm] {i}/{len(futures)} {path}")


def flat_row(metrics_path: Path, eval_root: Path) -> dict[str, Any]:
    metrics = load_json(metrics_path)
    vlm_path = metrics_path.parent / "vlm_score.json"
    vlm = load_json(vlm_path).get("vlm", {}) if vlm_path.is_file() else {}
    sample = metrics.get("sample") or {}
    align = metrics.get("alignment") or {}
    init_ctx = align.get("initial_context") or {}
    key_ev = align.get("ckpt_key_evidence") or {}

    return {
        "eval_root": eval_root.name,
        "run_id": metrics.get("run_id"),
        "ckpt": metrics.get("ckpt"),
        "domain": sample.get("domain"),
        "sample_id": sample.get("sample_id"),
        "video_name": sample.get("video_name"),
        "start_frame": sample.get("start_frame"),
        "mode": metrics.get("mode"),
        "num_chunks": metrics.get("num_chunks"),
        "context_frames": metrics.get("context_frames"),
        "closure_psnr": metrics.get("closure_psnr"),
        "closure_ssim": metrics.get("closure_ssim"),
        "closure_mse": metrics.get("closure_mse"),
        "vlm_overall_score": vlm.get("overall_score"),
        "vlm_bear_appearance_score": vlm.get("bear_appearance_score"),
        "vlm_bear_presence_score": vlm.get("bear_presence_score"),
        "vlm_scene_consistency_score": vlm.get("scene_consistency_score"),
        "vlm_view_revisit_score": vlm.get("view_revisit_score"),
        "vlm_verdict": vlm.get("verdict"),
        "vlm_reason": vlm.get("reason"),
        "memory_profile": align.get("memory_profile"),
        "memory_profile_matched": align.get("memory_profile_matched"),
        "action_injection_impl": align.get("action_injection_impl"),
        "camera_inject_mode_effective": align.get("camera_inject_mode_effective"),
        "rt_encoding_effective": align.get("rt_encoding_effective"),
        "use_framepack_memory": align.get("use_framepack_memory"),
        "use_framepack_length_compress": align.get("use_framepack_length_compress"),
        "framepack_ratio": align.get("framepack_ratio"),
        "use_spatial_memory": align.get("use_spatial_memory"),
        "use_spatial_memory_legacy": align.get("use_spatial_memory_legacy"),
        "spatial_memory_tokens": align.get("spatial_memory_tokens"),
        "spatial_memory_inject_mode": align.get("spatial_memory_inject_mode"),
        "ckpt_has_camera_encoder_keys": key_ev.get("has_camera_encoder_keys"),
        "ckpt_has_spatial_memory_module_keys": key_ev.get("has_spatial_memory_module_keys"),
        "ckpt_has_block_wise_ssm_keys": key_ev.get("has_block_wise_ssm_keys"),
        "ckpt_has_videossm_hybrid_keys": key_ev.get("has_videossm_hybrid_keys"),
        "ckpt_has_ssm_keys": key_ev.get("has_ssm_keys"),
        "initial_context_source": init_ctx.get("training_memory_source"),
        "initial_context_detail": init_ctx.get("context_source_detail"),
        "initial_context_frame_count": init_ctx.get("context_frame_count"),
        "stage1_metrics_path": str(metrics_path),
        "vlm_score_path": str(vlm_path) if vlm_path.is_file() else "",
        "video_path": metrics.get("video_path"),
    }


def mean(values: list[Any]) -> float | None:
    xs = []
    for value in values:
        try:
            if value is not None and value != "":
                xs.append(float(value))
        except (TypeError, ValueError):
            pass
    return sum(xs) / len(xs) if xs else None


def aggregate(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    groups: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        groups[(row["run_id"], row["domain"], row["mode"])].append(row)

    out = []
    for (run_id, domain, mode), group in sorted(groups.items()):
        out.append(
            {
                "run_id": run_id,
                "domain": domain,
                "mode": mode,
                "num_cases": len(group),
                "mean_closure_psnr": mean([r["closure_psnr"] for r in group]),
                "mean_closure_ssim": mean([r["closure_ssim"] for r in group]),
                "mean_closure_mse": mean([r["closure_mse"] for r in group]),
                "mean_vlm_overall_score": mean([r["vlm_overall_score"] for r in group]),
                "mean_vlm_bear_appearance_score": mean([r["vlm_bear_appearance_score"] for r in group]),
                "mean_vlm_bear_presence_score": mean([r["vlm_bear_presence_score"] for r in group]),
                "mean_vlm_scene_consistency_score": mean([r["vlm_scene_consistency_score"] for r in group]),
                "mean_vlm_view_revisit_score": mean([r["vlm_view_revisit_score"] for r in group]),
                "memory_profile": group[0]["memory_profile"],
                "context_frames": group[0]["context_frames"],
                "use_framepack_memory": group[0]["use_framepack_memory"],
                "use_framepack_length_compress": group[0]["use_framepack_length_compress"],
                "framepack_ratio": group[0]["framepack_ratio"],
                "use_spatial_memory": group[0]["use_spatial_memory"],
                "use_spatial_memory_legacy": group[0]["use_spatial_memory_legacy"],
                "spatial_memory_tokens": group[0]["spatial_memory_tokens"],
                "spatial_memory_inject_mode": group[0]["spatial_memory_inject_mode"],
                "eval_roots": "|".join(sorted({str(r["eval_root"]) for r in group})),
            }
        )
    return out


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if not rows:
        path.write_text("", encoding="utf-8")
        return
    fieldnames = list(rows[0].keys())
    with path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--eval-root", action="append", required=True)
    ap.add_argument("--out-dir", default="revisit_materials")
    ap.add_argument("--prefix", default="ep0_revisit")
    ap.add_argument("--score-vlm", action="store_true")
    ap.add_argument("--force-vlm", action="store_true")
    ap.add_argument("--vlm-workers", type=int, default=8)
    ap.add_argument("--vlm-api-base", default="http://127.0.0.1:8000/v1")
    ap.add_argument("--vlm-model", default="Qwen/Qwen2.5-VL-72B-Instruct")
    ap.add_argument("--vlm-timeout", type=int, default=180)
    ap.add_argument("--startup-wait-sec", type=int, default=900)
    args = ap.parse_args()

    eval_roots = [Path(p) for p in args.eval_root]
    metrics_paths: list[Path] = []
    for root in eval_roots:
        stage1 = root / "stage1"
        found = sorted(stage1.rglob("stage1_metrics.json"))
        print(f"[collect] {root}: {len(found)} stage1 metrics")
        metrics_paths.extend(found)

    maybe_score_vlm(metrics_paths, args)

    rows = []
    for root in eval_roots:
        for path in sorted((root / "stage1").rglob("stage1_metrics.json")):
            rows.append(flat_row(path, root))
    agg_rows = aggregate(rows)

    out_dir = Path(args.out_dir)
    cases_csv = out_dir / f"{args.prefix}_cases.csv"
    aggregate_csv = out_dir / f"{args.prefix}_aggregate_by_model_domain_mode.csv"
    write_csv(cases_csv, rows)
    write_csv(aggregate_csv, agg_rows)
    (out_dir / f"{args.prefix}_cases.json").write_text(json.dumps(rows, ensure_ascii=False, indent=2), encoding="utf-8")
    (out_dir / f"{args.prefix}_aggregate_by_model_domain_mode.json").write_text(json.dumps(agg_rows, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"[export] cases={len(rows)} -> {cases_csv}")
    print(f"[export] aggregates={len(agg_rows)} -> {aggregate_csv}")


if __name__ == "__main__":
    main()