#!/usr/bin/env python3
"""Score generated naive baselines with VBench-8 and build final summaries.

This script does not generate videos. Run
``generate_vbench8_extended_naive_baselines.py`` first, then use this entry
point to score FFFF plus any selected naive strategies. Strategies are scored
sequentially on the requested GPU; separate processes may be used on different
GPUs by passing disjoint ``--strategy`` sets and ``--no-summarize``.
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from scripts.naive_vbench_policies import (
    ALL_STRATEGY_NAMES,
    EVALUATION_STRATEGY_NAMES,
)


OUTPUT_DEFAULT = REPO_ROOT / "evaluation_runs/vbench8_extended_naive_baselines"
MAPPING_DEFAULT = REPO_ROOT / "assets/vbench8_extended_subset_mapping.json"
EXTENDED_PROMPTS_DEFAULT = REPO_ROOT / "prompts/vbench/all_dimension_extended.txt"
VBENCH_SITE_DEFAULT = REPO_ROOT / ".evaluation_env/vbench_site2"
VBENCH_INFO_DEFAULT = VBENCH_SITE_DEFAULT / "vbench/VBench_full_info.json"
VBENCH_CACHE_DEFAULT = Path("/data3/chenzhuo/.cache/vbench")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gpu", default="0")
    parser.add_argument("--output-root", type=Path, default=OUTPUT_DEFAULT)
    parser.add_argument("--mapping", type=Path, default=MAPPING_DEFAULT)
    parser.add_argument(
        "--extended-prompts", type=Path, default=EXTENDED_PROMPTS_DEFAULT
    )
    parser.add_argument("--vbench-site", type=Path, default=VBENCH_SITE_DEFAULT)
    parser.add_argument("--vbench-info", type=Path, default=VBENCH_INFO_DEFAULT)
    parser.add_argument("--vbench-cache", type=Path, default=VBENCH_CACHE_DEFAULT)
    parser.add_argument(
        "--strategy",
        action="append",
        choices=EVALUATION_STRATEGY_NAMES,
        default=None,
        help="Naive strategy to score; repeat as needed. Default: all six.",
    )
    parser.add_argument(
        "--score-ffff",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Also score the matched FFFF videos.",
    )
    parser.add_argument(
        "--skip-existing", action=argparse.BooleanOptionalAction, default=True
    )
    parser.add_argument(
        "--summarize",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Build pixel and final summaries for this exact strategy set.",
    )
    return parser.parse_args()


def checked_file(path: Path, label: str) -> Path:
    resolved = path.resolve()
    if not resolved.is_file():
        raise FileNotFoundError(f"Missing {label}: {resolved}")
    return resolved


def run(command: list[str], env: dict[str, str]) -> None:
    print("[run] " + " ".join(command), flush=True)
    subprocess.run(command, cwd=REPO_ROOT, env=env, check=True)


def main() -> None:
    args = parse_args()
    selected = args.strategy or list(EVALUATION_STRATEGY_NAMES)
    if len(set(selected)) != len(selected):
        raise ValueError("--strategy values must be unique")
    strategies = (["ffff"] if args.score_ffff else []) + selected
    if args.summarize and "ffff" not in strategies:
        raise ValueError("Final summaries require --score-ffff")
    unknown = set(strategies) - set(ALL_STRATEGY_NAMES)
    if unknown:
        raise ValueError(f"Unknown strategies: {sorted(unknown)}")

    output_root = args.output_root.resolve()
    mapping = checked_file(args.mapping, "mapping")
    extended_prompts = checked_file(args.extended_prompts, "extended prompts")
    vbench_site = args.vbench_site.resolve()
    if not vbench_site.is_dir():
        raise FileNotFoundError(f"Missing VBench site directory: {vbench_site}")
    vbench_info = checked_file(args.vbench_info, "VBench full info")
    videos_root = output_root / "generated_videos"
    if not videos_root.is_dir():
        raise FileNotFoundError(f"Missing generated videos: {videos_root}")

    env = os.environ.copy()
    existing_pythonpath = env.get("PYTHONPATH")
    pythonpath_parts = [str(vbench_site), str(REPO_ROOT)]
    if existing_pythonpath:
        pythonpath_parts.append(existing_pythonpath)
    env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
    env["VBENCH_CACHE_DIR"] = str(args.vbench_cache.resolve())
    env["VBENCH_BERT_MODEL_DIR"] = str(
        (args.vbench_cache / "bert-base-uncased").resolve()
    )

    if args.summarize:
        run(
            [
                sys.executable,
                str(REPO_ROOT / "scripts/summarize_vbench8_generation.py"),
                "--output-root",
                str(output_root),
                "--strategies",
                *strategies,
            ],
            env,
        )

    for strategy in strategies:
        command = [
            sys.executable,
            str(REPO_ROOT / "scripts/eval_vbench8_extended_subset.py"),
            "--gpu",
            str(args.gpu),
            "--strategy",
            strategy,
            "--mapping",
            str(mapping),
            "--videos-root",
            str(videos_root),
            "--output-root",
            str(output_root),
            "--vbench-info",
            str(vbench_info),
        ]
        if args.skip_existing:
            command.append("--skip-existing")
        run(command, env)

    if args.summarize:
        run(
            [
                sys.executable,
                str(REPO_ROOT / "scripts/summarize_vbench8_extended.py"),
                "--output-root",
                str(output_root),
                "--mapping",
                str(mapping),
                "--extended-prompts",
                str(extended_prompts),
                "--vbench-info",
                str(vbench_info),
                "--strategies",
                *strategies,
            ],
            env,
        )
    print(f"[complete] strategies={','.join(strategies)}", flush=True)


if __name__ == "__main__":
    main()