# SPDX-License-Identifier: Apache-2.0
"""Config-driven inference performance tests.

Benchmark configs live in .buildkite/performance-benchmarks/tests/*.json.
Each JSON file defines model params, generation kwargs, run config, and
per-device thresholds.  This test module auto-discovers all configs and
parametrizes a single test function over them.
"""
import glob
import json
import os
import time
from collections.abc import Mapping
from datetime import datetime, timezone

import torch
import pytest

from fastvideo import VideoGenerator
from fastvideo.logger import init_logger
from fastvideo.worker.multiproc_executor import MultiprocExecutor

logger = init_logger(__name__)

STAGE_METRIC_MAP: dict[str, str] = {
    "TextEncodingStage": "text_encoder_time_s",
    "DenoisingStage": "dit_time_s",
    "DmdDenoisingStage": "dit_time_s",
    "DecodingStage": "vae_decode_time_s",
}
V2_CONFIG_SCHEMA_VERSION = 2
V2_REQUIRED_IDENTITY_FIELDS = (
    "workload_id",
    "variant_id",
    "benchmark_version",
)
V2_OPTIONAL_METADATA_FIELDS = (
    "recipe",
    "metric_threshold_policy",
    "quality_metadata",
)

# -- Config discovery -------------------------------------------------------

_BENCHMARKS_DIR = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    "..",
    "..",
    "..",
    ".buildkite",
    "performance-benchmarks",
    "tests",
)


def _has_v2_fields(cfg):
    v2_fields = V2_REQUIRED_IDENTITY_FIELDS + V2_OPTIONAL_METADATA_FIELDS
    return any(field in cfg for field in v2_fields)


def _is_v2_config(cfg):
    return cfg.get("config_schema_version") == V2_CONFIG_SCHEMA_VERSION


def _validate_non_empty_string(value, field, path):
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{path}: v2 identity field {field!r} must be a non-empty string")


def _validate_integer(value, field, path):
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError(f"{path}: v2 identity field {field!r} must be an integer")


def _validate_benchmark_config(cfg, path="<memory>"):
    missing_common = [field for field in ("benchmark_id",) if field not in cfg]
    if missing_common:
        raise ValueError(f"{path}: missing required benchmark config fields: {', '.join(missing_common)}")

    schema_version = cfg.get("config_schema_version")
    if schema_version is None:
        if _has_v2_fields(cfg):
            raise ValueError(f"{path}: v2 benchmark identity fields require config_schema_version=2")
        return

    if schema_version != V2_CONFIG_SCHEMA_VERSION:
        raise ValueError(f"{path}: unsupported benchmark config_schema_version={schema_version!r}")

    missing_v2 = [field for field in V2_REQUIRED_IDENTITY_FIELDS if field not in cfg]
    if missing_v2:
        raise ValueError(f"{path}: missing required v2 identity fields: {', '.join(missing_v2)}")

    _validate_non_empty_string(cfg["workload_id"], "workload_id", path)
    _validate_non_empty_string(cfg["variant_id"], "variant_id", path)
    _validate_integer(cfg["benchmark_version"], "benchmark_version", path)

    for field in V2_OPTIONAL_METADATA_FIELDS:
        if field in cfg and not isinstance(cfg[field], Mapping):
            raise ValueError(f"{path}: optional v2 metadata field {field!r} must be an object")


def _config_identity_metadata(cfg):
    if not _is_v2_config(cfg):
        return {}
    metadata = {
        "config_schema_version": cfg["config_schema_version"],
        "workload_id": cfg["workload_id"],
        "variant_id": cfg["variant_id"],
        "benchmark_version": cfg["benchmark_version"],
    }
    for field in V2_OPTIONAL_METADATA_FIELDS:
        if field in cfg:
            metadata[field] = cfg[field]
    return metadata


def _benchmark_display_id(cfg):
    return cfg["benchmark_id"]


def _discover_benchmarks():
    """Glob benchmark JSON configs and return list of (id, config) tuples."""
    pattern = os.path.join(_BENCHMARKS_DIR, "*.json")
    configs = []
    for path in sorted(glob.glob(pattern)):
        with open(path) as f:
            cfg = json.load(f)
        _validate_benchmark_config(cfg, path)
        configs.append(cfg)
    return configs


_BENCHMARK_CONFIGS = _discover_benchmarks()

# -- Helpers ----------------------------------------------------------------


def _get_thresholds(cfg):
    """Return thresholds dict for the current GPU from config."""
    device_name = torch.cuda.get_device_name()
    thresholds = cfg.get("thresholds", {})
    for gpu_key, thresh in thresholds.items():
        if gpu_key in device_name:
            logger.info("Using thresholds for %s: %s", gpu_key, thresh)
            return thresh
    default = thresholds.get("default", {
        "max_generation_time_s": 120.0,
        "max_peak_memory_mb": 30000.0,
    })
    logger.warning("No thresholds for device '%s', using defaults", device_name)
    return default


def _shutdown_executor(generator):
    if generator is None:
        return
    if isinstance(generator.executor, MultiprocExecutor):
        generator.executor.shutdown()


def _extract_component_times(result: dict) -> dict[str, float | None]:
    component_times: dict[str, float | None] = {
        "text_encoder_time_s": None,
        "dit_time_s": None,
        "vae_decode_time_s": None,
    }
    logging_info = result.get("logging_info")
    if logging_info is None:
        return component_times
    if isinstance(logging_info, Mapping):
        stages: dict = logging_info.get("stages", {}) or {}
    else:
        stages: dict = getattr(logging_info, "stages", {}) or {}
    if not stages:
        return component_times
    logger.info("Discovered pipeline stages: %s", list(stages.keys()))
    for stage_name, stage_data in stages.items():
        if not isinstance(stage_data, Mapping):
            logger.debug("Skipping malformed stage '%s' data: %r", stage_name, stage_data)
            continue
        stage_class = stage_data.get("stage_class", stage_name)
        component_metric = stage_data.get("component_metric")
        if isinstance(component_metric, str) and component_metric in component_times:
            metric_key = component_metric
        else:
            metric_key = STAGE_METRIC_MAP.get(stage_class)
        if metric_key is None:
            logger.debug("Unmapped stage '%s' class '%s' (%.3fs)",
                         stage_name,
                         stage_class,
                         stage_data.get("execution_time", 0))
            continue
        elapsed = stage_data.get("execution_time")
        if elapsed is None:
            continue
        existing = component_times[metric_key]
        component_times[metric_key] = (
            elapsed if existing is None else existing + elapsed)
    return component_times


def _avg_component(all_component_times: list[dict], key: str) -> float | None:
    vals = [r[key] for r in all_component_times if r.get(key) is not None]
    return round(sum(vals) / len(vals), 3) if vals else None


def _run_generation(generator, prompt, generation_kwargs):
    """Run a single generation, return (elapsed_s, peak_memory_mb, component_times)."""
    torch.cuda.synchronize()
    start = time.perf_counter()
    result = generator.generate_video(prompt, **generation_kwargs)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start
    peak_memory_mb = result.get("peak_memory_mb", 0.0) or 0.0
    component_times = _extract_component_times(result)
    return elapsed, peak_memory_mb, component_times

def _write_results(results):
    """Write JSON results to the results directory."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    results_dir = os.path.join(script_dir, "results")
    os.makedirs(results_dir, exist_ok=True)

    bid = results.get("benchmark_id", "unknown")
    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    filename = f"perf_{bid}_{ts}.json"
    filepath = os.path.join(results_dir, filename)

    with open(filepath, "w") as f:
        json.dump(results, f, indent=2)
    logger.info("Performance results written to %s", filepath)


# -- Test -------------------------------------------------------------------

def _run_benchmark(cfg):
    run_config = cfg.get("run_config", {})
    required_gpus = run_config.get("required_gpus", 1)
    available = torch.cuda.device_count()
    if available < required_gpus:
        pytest.skip(f"Need {required_gpus} GPUs, only {available} available")

    model_info = cfg["model"]
    init_kwargs = dict(cfg.get("init_kwargs", {}))
    gen_kwargs = dict(cfg.get("generation_kwargs", {}))
    prompts = cfg.get("test_prompts", ["A cinematic video."])
    prompt = prompts[0]

    num_warmup = run_config.get("num_warmup_runs", 1)
    num_measure = run_config.get("num_measurement_runs", 3)
    thresholds = _get_thresholds(cfg)

    # Remap JSON keys to VideoGenerator kwargs
    text_enc_prec = init_kwargs.pop("text_encoder_precisions", None)
    if text_enc_prec is not None:
        init_kwargs["text_encoder_precisions"] = tuple(text_enc_prec)

    # Output directory for generated videos
    script_dir = os.path.dirname(os.path.abspath(__file__))
    output_dir = os.path.join(script_dir, "generated_videos",
                              cfg["benchmark_id"])
    os.makedirs(output_dir, exist_ok=True)
    gen_kwargs["output_path"] = output_dir

    generator = None
    try:
        generator = VideoGenerator.from_pretrained(
            model_path=model_info["model_path"],
            **init_kwargs,
        )

        for i in range(num_warmup):
            logger.info("Warmup run %d/%d", i + 1, num_warmup)
            _run_generation(generator, prompt, gen_kwargs)

        times = []
        peak_memories = []
        all_component_times = []
        for i in range(num_measure):
            logger.info("Measurement run %d/%d", i + 1, num_measure)
            elapsed, peak_mb, component_times = _run_generation(
                generator,
                prompt,
                gen_kwargs,
            )
            logger.info("  Time: %.2fs, Peak memory: %.0fMB", elapsed, peak_mb)
            times.append(elapsed)
            peak_memories.append(peak_mb)
            all_component_times.append(component_times)
    finally:
        _shutdown_executor(generator)

    avg_time = sum(times) / len(times)
    max_peak_memory = max(peak_memories)
    device_name = torch.cuda.get_device_name()
    num_frames = gen_kwargs.get("num_frames")
    throughput_fps = (1.0 / avg_time) if avg_time > 0 else None
    if isinstance(num_frames, (int, float)) and avg_time > 0:
        throughput_fps = num_frames / avg_time

    results = {
        "benchmark_id": cfg["benchmark_id"],
        **_config_identity_metadata(cfg),
        "model_short_name": model_info.get("model_short_name", ""),
        "device": device_name,
        "num_gpus": init_kwargs.get("num_gpus", 1),
        "num_warmup_runs": num_warmup,
        "num_measurement_runs": num_measure,
        "avg_generation_time_s": round(avg_time, 3),
        "individual_times_s": [round(t, 3) for t in times],
        "throughput_fps": round(throughput_fps, 3)
        if throughput_fps is not None else None,
        "max_peak_memory_mb": round(max_peak_memory, 1),
        "individual_peak_memories_mb": [round(m, 1) for m in peak_memories],
        "thresholds": thresholds,
        "regression_thresholds": cfg.get("regression_thresholds", {}),
        "commit": os.environ.get("BUILDKITE_COMMIT", ""),
        "pr_number": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "text_encoder_time_s": _avg_component(all_component_times,
                                              "text_encoder_time_s"),
        "dit_time_s": _avg_component(all_component_times, "dit_time_s"),
        "vae_decode_time_s": _avg_component(all_component_times,
                                            "vae_decode_time_s"),
    }

    logger.info(
        "Performance results: avg_time=%.2fs, "
        "max_peak_memory=%.0fMB", avg_time, max_peak_memory)
    _write_results(results)

    max_time = thresholds["max_generation_time_s"]
    max_mem = thresholds["max_peak_memory_mb"]

    assert avg_time <= max_time, (
        f"Average generation time {avg_time:.2f}s exceeds "
        f"threshold {max_time:.1f}s for {device_name}")

    assert max_peak_memory <= max_mem, (
        f"Peak memory {max_peak_memory:.0f}MB exceeds "
        f"threshold {max_mem:.0f}MB for {device_name}")

    component_thresholds = {
        "text_encoder_time_s": thresholds.get("max_text_encoder_time_s"),
        "dit_time_s": thresholds.get("max_dit_time_s"),
        "vae_decode_time_s": thresholds.get("max_vae_decode_time_s"),
    }
    for metric, max_val in component_thresholds.items():
        if max_val is None:
            continue
        actual = results[metric]
        if actual is not None:
            assert actual <= max_val, (
                f"{metric} {actual:.3f}s exceeds threshold {max_val:.3f}s "
                f"for {device_name}")


@pytest.mark.parametrize(
    "cfg",
    _BENCHMARK_CONFIGS,
    ids=[_benchmark_display_id(c) for c in _BENCHMARK_CONFIGS],
)
def test_inference_performance(cfg):
    """Measure generation latency, peak GPU memory, and component-level timings
    (text encoder, DiT, VAE decode). Assert each against device-aware thresholds.
    """

    original_env = os.environ.get("FASTVIDEO_STAGE_LOGGING")
    os.environ["FASTVIDEO_STAGE_LOGGING"] = "1"
    try:
        _run_benchmark(cfg)
    finally:
        if original_env is None:
            os.environ.pop("FASTVIDEO_STAGE_LOGGING", None)
        else:
            os.environ["FASTVIDEO_STAGE_LOGGING"] = original_env