# SPDX-License-Identifier: Apache-2.0
"""Generic AISBench VBench evaluation for pre-generated videos.

This module deliberately has no FastVideo or algorithm dependency.
"""

from __future__ import annotations

import importlib
import os
import re
import subprocess
import sys
from pathlib import Path

from flashgen.benchmark.datasets import (
    DEFAULT_MINI_RATIO,
    DEFAULT_SAMPLING,
    MiniRatio,
    SamplingMethod,
    load_mini_dataset,
    validate_videos,
)
from flashgen.benchmark.runtime import require_npu


def _require_aisbench_cli() -> None:
    try:
        importlib.import_module("ais_bench.benchmark.cli.main")
    except ImportError as exc:
        raise RuntimeError("The evaluate command requires AISBench Benchmark with VBench support installed.") from exc


def _safe_abbr(value: str) -> str:
    sanitized = re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("_.-")
    return sanitized or "vbench_input"


def _write_aisbench_config(
    output_root: Path,
    video_root: Path,
    metadata_path: Path,
    dimensions: tuple[str, ...],
    model_abbr: str,
    cache_dir: str,
) -> Path:
    """Write the same eval-only configuration accepted by the AISBench CLI."""
    config_path = output_root / "aisbench_vbench_config.py"
    config_text = f'''# Auto-generated by FlashGen benchmark evaluate.
from ais_bench.benchmark.datasets import VBenchDataset
from ais_bench.benchmark.partitioners import NaivePartitioner
from ais_bench.benchmark.runners import LocalRunner
from ais_bench.benchmark.summarizers import VBenchSummarizer
from ais_bench.benchmark.tasks import VBenchEvalTask

DATA_PATH = {str(video_root)!r}
FULL_JSON_DIR = {str(metadata_path)!r}
VBENCH_CACHE_DIR = {cache_dir!r}
VBENCH_DIMENSIONS = {list(dimensions)!r}

models = [
    dict(
        attr="local",
        type="VBenchEvalPlaceholder",
        abbr={model_abbr!r},
    )
]

datasets = [
    dict(
        abbr=f"vbench_{{dimension}}",
        type=VBenchDataset,
        path=DATA_PATH,
        full_json_dir=FULL_JSON_DIR,
        eval_cfg=dict(
            dimension_list=[dimension],
            load_ckpt_from_local=True,
        ),
    )
    for dimension in VBENCH_DIMENSIONS
]

eval = dict(
    partitioner=dict(type=NaivePartitioner),
    runner=dict(
        type=LocalRunner,
        task=dict(type=VBenchEvalTask),
    ),
)

summarizer = dict(
    attr="accuracy",
    type=VBenchSummarizer,
)
'''
    config_path.write_text(config_text, encoding="utf-8")
    return config_path


def evaluate_videos(
    videos_dir: str | os.PathLike[str],
    work_dir: str | os.PathLike[str],
    *,
    mini_ratio: MiniRatio = DEFAULT_MINI_RATIO,
    sampling: SamplingMethod = DEFAULT_SAMPLING,
    dataset_root: str | os.PathLike[str] | None = None,
    name: str | None = None,
    vbench_cache_dir: str | os.PathLike[str] | None = None,
    max_num_workers: int = 1,
    max_workers_per_gpu: int = 1,
) -> Path:
    """Validate and evaluate videos through AISBench's native VBench workflow."""
    require_npu()
    if max_num_workers < 1:
        raise ValueError("max_num_workers must be at least 1")
    if max_workers_per_gpu < 1:
        raise ValueError("max_workers_per_gpu must be at least 1")
    dataset = load_mini_dataset(mini_ratio, sampling, dataset_root)
    video_root = validate_videos(dataset, videos_dir)
    output_root = Path(work_dir).expanduser().resolve()
    output_root.mkdir(parents=True, exist_ok=True)

    _require_aisbench_cli()
    cache = Path(vbench_cache_dir).expanduser().resolve() if vbench_cache_dir is not None else None
    if cache is not None:
        os.environ["VBENCH_CACHE_DIR"] = str(cache)

    model_abbr = _safe_abbr(name or video_root.name)
    config_path = _write_aisbench_config(
        output_root=output_root,
        video_root=video_root,
        metadata_path=dataset.metadata_path,
        dimensions=dataset.dimensions,
        model_abbr=model_abbr,
        cache_dir=str(cache) if cache is not None else os.environ.get("VBENCH_CACHE_DIR", ""),
    )
    command = [
        sys.executable,
        "-m",
        "ais_bench.benchmark.cli.main",
        str(config_path),
        "--mode",
        "eval",
        "--work-dir",
        str(output_root.parent),
        "--reuse",
        output_root.name,
        "--max-num-workers",
        str(max_num_workers),
        "--max-workers-per-gpu",
        str(max_workers_per_gpu),
        "--dump-eval-details",
    ]
    env = os.environ.copy()
    env.setdefault("PYTHONUNBUFFERED", "1")
    try:
        subprocess.run(command, cwd=output_root, env=env, check=True)
    except subprocess.CalledProcessError as exc:
        raise RuntimeError(f"AISBench VBench CLI failed with exit code {exc.returncode}.") from exc
    return output_root