# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
"""
Level 2: Run all existing tuned configs through --run_config to verify
the production operator works with every shape in the config CSVs.

For each tuner family, discovers all tuned CSVs (default + model_configs),
merges them via pathsep, and runs --run_config to benchmark every shape.
Any shape that errors (us=-1 or exception) is reported as a test failure.

Run:
    python3 -m unittest op_tests.tuning_tests.test_run_config -v
"""

import csv
import os
import re
import sys
import subprocess
import unittest

AITER_ROOT = os.path.dirname(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
CONFIGS_DIR = os.path.join(AITER_ROOT, "aiter", "configs")
MODEL_CONFIGS_DIR = os.path.join(CONFIGS_DIR, "model_configs")

# Override: specify a tuner family and config CSV to test directly.
#   TUNE_TEST_FAMILY=a8w8_blockscale TUNE_TEST_CONFIG=/path/to/tuned.csv \
#     python3 -m unittest op_tests.tuning_tests.test_run_config.TestRunConfigCustom -v
#
# TUNE_TEST_CONFIG supports pathsep (:) for merging multiple CSVs, e.g.:
#   TUNE_TEST_CONFIG="configs/a8w8_blockscale_tuned_gemm.csv:model_configs/xxx.csv"
TUNE_TEST_FAMILY = os.environ.get("TUNE_TEST_FAMILY")
TUNE_TEST_CONFIG = os.environ.get("TUNE_TEST_CONFIG")


def _gpu_available():
    try:
        import torch

        return torch.cuda.is_available() and torch.cuda.device_count() > 0
    except ImportError:
        return False


def _find_tuned_csvs(pattern):
    """Find all tuned CSVs matching pattern in configs/ and model_configs/."""
    found = []
    for d in (CONFIGS_DIR, MODEL_CONFIGS_DIR):
        if not os.path.isdir(d):
            continue
        for f in sorted(os.listdir(d)):
            if (
                pattern in f
                and "tuned" in f
                and "untuned" not in f
                and f.endswith(".csv")
            ):
                found.append(os.path.join(d, f))
    return found


def _resolve_config_via_aiter(config_property):
    """Resolve config file through AITER_CONFIGS (same path as production).
    Returns the resolved file path, or None if unavailable."""
    try:
        from aiter.jit.core import AITER_CONFIGS

        config_file = getattr(AITER_CONFIGS, config_property, None)
        if config_file and os.path.exists(config_file):
            return config_file
    except Exception:
        pass
    return None


def _merge_config_paths(csv_list):
    """Merge multiple CSV paths with os.pathsep (like AITER_CONFIG_* env)."""
    return os.pathsep.join(csv_list)


def _run_config(script, config_csv, timeout=600, extra_args=None):
    """Run tuner with --run_config <tuned_csv> and return result."""
    cmd = [
        sys.executable,
        os.path.join(AITER_ROOT, script),
        "--run_config",
        config_csv,
        "--warmup",
        "2",
        "--iters",
        "5",
    ]
    if extra_args:
        cmd.extend(extra_args)
    env = os.environ.copy()
    script_dir = os.path.dirname(os.path.join(AITER_ROOT, script))
    env["PYTHONPATH"] = script_dir + ":" + env.get("PYTHONPATH", "")
    try:
        return subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout,
            cwd=AITER_ROOT,
            env=env,
        )
    except subprocess.TimeoutExpired as e:
        raise AssertionError(
            f"run_config timed out after {timeout}s\n"
            f"  cmd: {' '.join(cmd)}\n"
            f"  stdout (last 500): {(e.stdout or b'')[-500:]}\n"
            f"  stderr (last 500): {(e.stderr or b'')[-500:]}"
        ) from None


def _parse_benchmark_results(lines):
    """Parse benchmark output lines. Returns (errors, mismatches, ok_count, skip_count).
    ERROR/MISMATCH entries include the reason from the following line if present."""
    error_shapes = []
    mismatch_shapes = []
    ok_count = 0
    skip_count = 0
    for i, line in enumerate(lines):
        stripped = line.strip()
        if "| " not in stripped:
            continue
        reason = ""
        for j in range(i + 1, min(i + 4, len(lines))):
            if lines[j].strip().startswith("reason:"):
                reason = " | " + lines[j].strip()
                break
        if stripped.endswith("ERROR"):
            error_shapes.append(stripped + reason)
        elif stripped.endswith("MISMATCH"):
            mismatch_shapes.append(stripped + reason)
        elif stripped.endswith("OK"):
            ok_count += 1
        elif stripped.endswith("SKIP"):
            skip_count += 1
    return error_shapes, mismatch_shapes, ok_count, skip_count


def _extract_repro_and_reasons(lines):
    """Extract Repro CSV block from tuner output."""
    parts = []
    repro = []
    in_repro = False
    for line in lines:
        if "Repro CSV" in line:
            in_repro = True
        if in_repro:
            repro.append(line)
    if repro:
        parts.append("\n" + "\n".join(repro))
    return "\n".join(parts)


def _parse_all_benchmark_results(lines):
    """Parse all benchmark result lines, returning a list of dicts with timing data.
    Each dict: {shape, e2e_us, status, reason}."""
    all_results = []
    for i, line in enumerate(lines):
        stripped = line.strip()
        if "| " not in stripped:
            continue
        if stripped.startswith("Shape") or stripped.startswith("-"):
            continue
        parts = [p.strip() for p in stripped.split("|")]
        if len(parts) < 3:
            continue
        status = parts[-1]
        if status not in ("OK", "ERROR", "MISMATCH", "SKIP"):
            continue
        shape = parts[0]
        e2e_us_str = parts[-2]
        try:
            e2e_us = float(e2e_us_str)
        except ValueError:
            e2e_us = None
        reason = ""
        for j in range(i + 1, min(i + 4, len(lines))):
            if lines[j].strip().startswith("reason:"):
                reason = lines[j].strip()[len("reason:") :].strip()
                break
        result = {"shape": shape, "e2e_us": e2e_us, "status": status}
        if reason:
            result["reason"] = reason
        # Parse shape tuple into individual columns (M, N, K, ...)
        m = re.match(r"\((.+)\)", shape)
        if m:
            vals = [v.strip() for v in m.group(1).split(",")]
            for idx, col in enumerate(["M", "N", "K"]):
                if idx < len(vals):
                    try:
                        result[col] = int(vals[idx])
                    except ValueError:
                        result[col] = vals[idx]
        all_results.append(result)
    return all_results


def _save_results_csv(name, all_results):
    """Save all benchmark results to CSV for later comparison."""
    if not all_results:
        return None
    os.makedirs(REPORT_DIR, exist_ok=True)
    csv_file = os.path.join(REPORT_DIR, f"{name}_run_config_results.csv")
    fieldnames = ["M", "N", "K", "e2e_us", "status", "reason"]
    with open(csv_file, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(all_results)
    return csv_file


REPORT_DIR = "/tmp/tuning_test_reports"


def _format_failures(
    name, error_shapes, mismatch_shapes, skip_count=0, output_lines=None
):
    """Write detailed report to file, return short summary for AssertionError."""
    os.makedirs(REPORT_DIR, exist_ok=True)
    report_file = os.path.join(REPORT_DIR, f"{name}_run_config_report.txt")

    with open(report_file, "w") as f:
        f.write(f"=== {name} run_config report ===\n\n")
        if error_shapes:
            f.write(f"Errors ({len(error_shapes)} shapes):\n")
            f.write("\n".join(error_shapes) + "\n\n")
        if mismatch_shapes:
            f.write(f"Accuracy mismatches ({len(mismatch_shapes)} shapes):\n")
            f.write("\n".join(mismatch_shapes) + "\n\n")
        if output_lines:
            repro = _extract_repro_and_reasons(output_lines)
            if repro:
                f.write(repro + "\n")

    parts = []
    if error_shapes:
        parts.append(f"Errors: {len(error_shapes)} shapes")
    if mismatch_shapes:
        parts.append(f"Mismatches: {len(mismatch_shapes)} shapes")
    if skip_count > 0:
        parts.append(f"Skipped: {skip_count} shapes")
    if error_shapes or mismatch_shapes:
        parts.append(f"\nFull report: {report_file}")
    return parts


TUNER_FAMILIES = {
    "a8w8": {
        "script": "csrc/ck_gemm_a8w8/gemm_a8w8_tune.py",
        "csv_pattern": "a8w8_tuned_gemm",
        "exclude_patterns": ["bpreshuffle", "blockscale", "batched"],
        "config_property": "AITER_CONFIG_GEMM_A8W8_FILE",
    },
    "a8w8_bpreshuffle": {
        "script": "csrc/ck_gemm_a8w8_bpreshuffle/gemm_a8w8_bpreshuffle_tune.py",
        "csv_pattern": "a8w8_bpreshuffle_tuned_gemm",
        "exclude_patterns": ["blockscale"],
        "config_property": "AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE",
    },
    "a8w8_blockscale": {
        "script": "csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_tune.py",
        "csv_pattern": "a8w8_blockscale_tuned_gemm",
        "exclude_patterns": ["bpreshuffle", "fmoe"],
        "config_property": "AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_FILE",
    },
    "a8w8_blockscale_bpreshuffle": {
        "script": "csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_tune.py",
        "csv_pattern": "a8w8_blockscale_bpreshuffle_tuned_gemm",
        "exclude_patterns": ["fmoe"],
        "extra_args": ["--preshuffle"],
        "config_property": "AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_FILE",
    },
    "a4w4_blockscale": {
        "script": "csrc/ck_gemm_a4w4_blockscale/gemm_a4w4_blockscale_tune.py",
        "csv_pattern": "a4w4_blockscale_tuned_gemm",
        "exclude_patterns": [],
        "config_property": "AITER_CONFIG_GEMM_A4W4_FILE",
    },
    "batched_a8w8": {
        "script": "csrc/ck_batched_gemm_a8w8/batched_gemm_a8w8_tune.py",
        "csv_pattern": "a8w8_tuned_batched_gemm",
        "exclude_patterns": [],
        "config_property": "AITER_CONFIG_A8W8_BATCHED_GEMM_FILE",
    },
    "batched_bf16": {
        "script": "csrc/ck_batched_gemm_bf16/batched_gemm_bf16_tune.py",
        "csv_pattern": "bf16_tuned_batched_gemm",
        "exclude_patterns": [],
        "config_property": "AITER_CONFIG_BF16_BATCHED_GEMM_FILE",
    },
    "fmoe": {
        "script": "csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py",
        "csv_pattern": "tuned_fmoe",
        "exclude_patterns": ["untuned", "profile"],
        "timeout": 1200,
        "config_property": "AITER_CONFIG_FMOE_FILE",
    },
    "gradlib_bf16": {
        "script": "gradlib/gradlib/gemm_tuner.py",
        "csv_pattern": "bf16_tuned_gemm",
        "exclude_patterns": ["batched"],
        "config_property": "AITER_CONFIG_GEMM_BF16_FILE",
    },
}


@unittest.skipUnless(_gpu_available(), "No GPU available")
class TestRunConfig(unittest.TestCase):
    """Run --run_config on all existing tuned CSVs to verify production ops."""

    def _test_family(self, name):
        cfg = TUNER_FAMILIES[name]
        timeout = cfg.get("timeout", 600)
        extra_args = cfg.get("extra_args", None)

        config_prop = cfg.get("config_property")
        merged = _resolve_config_via_aiter(config_prop) if config_prop else None

        if merged:
            csv_names = [f"{config_prop} -> {merged}"]
        else:
            pattern = cfg["csv_pattern"]
            excludes = cfg.get("exclude_patterns", [])
            csvs = _find_tuned_csvs(pattern)
            csvs = [
                c for c in csvs if not any(ex in os.path.basename(c) for ex in excludes)
            ]
            if not csvs:
                self.skipTest(f"No tuned CSVs found for {name} (pattern={pattern})")
            merged = _merge_config_paths(csvs)
            csv_names = [os.path.basename(c) for c in csvs]

        result = _run_config(
            cfg["script"], merged, timeout=timeout, extra_args=extra_args
        )

        output = result.stdout + result.stderr
        if result.returncode != 0:
            print(f"\n=== {name} run_config FAILED ===")
            print(f"CSVs: {csv_names}")
            print(f"STDOUT (last 2000):\n{result.stdout[-2000:]}")
            print(f"STDERR (last 2000):\n{result.stderr[-2000:]}")

        self.assertEqual(
            result.returncode, 0, f"{name} run_config failed (csvs={csv_names})"
        )

        lines = output.split("\n")
        error_shapes, mismatch_shapes, ok_count, skip_count = _parse_benchmark_results(
            lines
        )

        all_results = _parse_all_benchmark_results(lines)
        csv_file = _save_results_csv(name, all_results)
        if csv_file:
            print(f"\n  Results CSV: {csv_file} ({len(all_results)} shapes)")

        failures = _format_failures(
            name, error_shapes, mismatch_shapes, skip_count, lines
        )
        self.assertEqual(
            len(failures),
            0,
            f"{name} run_config: {ok_count} OK, {skip_count} SKIP, "
            f"{len(error_shapes)} ERROR, {len(mismatch_shapes)} MISMATCH\n"
            + "\n".join(failures),
        )

    def test_a8w8(self):
        self._test_family("a8w8")

    def test_a8w8_bpreshuffle(self):
        self._test_family("a8w8_bpreshuffle")

    def test_a8w8_blockscale(self):
        self._test_family("a8w8_blockscale")

    def test_a8w8_blockscale_bpreshuffle(self):
        self._test_family("a8w8_blockscale_bpreshuffle")

    def test_a4w4_blockscale(self):
        self._test_family("a4w4_blockscale")

    def test_batched_a8w8(self):
        self._test_family("batched_a8w8")

    def test_batched_bf16(self):
        self._test_family("batched_bf16")

    def test_fmoe(self):
        self._test_family("fmoe")

    def test_gradlib_bf16(self):
        self._test_family("gradlib_bf16")


@unittest.skipUnless(_gpu_available(), "No GPU available")
@unittest.skipUnless(
    TUNE_TEST_FAMILY,
    "Set TUNE_TEST_FAMILY (and optionally TUNE_TEST_CONFIG) to run",
)
class TestRunConfigCustom(unittest.TestCase):
    """Run --run_config with user-specified family and config CSV.

    Usage:
        # Use AITER_CONFIGS resolution (same as production):
        TUNE_TEST_FAMILY=a8w8_blockscale \
        python3 -m unittest op_tests.tuning_tests.test_run_config.TestRunConfigCustom -v

        # Explicit config CSV:
        TUNE_TEST_FAMILY=a8w8_blockscale \
        TUNE_TEST_CONFIG="aiter/configs/a8w8_blockscale_tuned_gemm.csv" \
        python3 -m unittest op_tests.tuning_tests.test_run_config.TestRunConfigCustom -v

        # Multiple configs (merged):
        TUNE_TEST_FAMILY=a8w8_blockscale \
        TUNE_TEST_CONFIG="aiter/configs/a8w8_blockscale_tuned_gemm.csv:aiter/configs/model_configs/a8w8_blockscale_tuned_gemm_ds_v3.csv" \
        python3 -m unittest op_tests.tuning_tests.test_run_config.TestRunConfigCustom -v
    """

    def test_custom(self):
        family = TUNE_TEST_FAMILY
        config = TUNE_TEST_CONFIG
        self.assertIn(
            family,
            TUNER_FAMILIES,
            f"Unknown family '{family}'. Available: {list(TUNER_FAMILIES.keys())}",
        )
        cfg = TUNER_FAMILIES[family]
        timeout = cfg.get("timeout", 600)
        extra_args = cfg.get("extra_args", None)

        if config:
            # Explicit config: resolve relative paths against AITER_ROOT
            resolved = []
            for p in config.split(os.pathsep):
                p = p.strip()
                if not p:
                    continue
                if not os.path.isabs(p):
                    p = os.path.join(AITER_ROOT, p)
                self.assertTrue(os.path.exists(p), f"Config not found: {p}")
                resolved.append(p)
            merged = os.pathsep.join(resolved)
            csv_names = [os.path.basename(p) for p in resolved]
        else:
            # No config specified: resolve via AITER_CONFIGS (production path)
            config_prop = cfg.get("config_property")
            self.assertIsNotNone(
                config_prop, f"No config_property defined for family '{family}'"
            )
            merged = _resolve_config_via_aiter(config_prop)
            self.assertIsNotNone(
                merged,
                f"Could not resolve config for {family} via AITER_CONFIGS.{config_prop}",
            )
            csv_names = [f"{config_prop} -> {merged}"]

        print(f"\nRunning {family} --run_config with: {csv_names}")
        result = _run_config(
            cfg["script"], merged, timeout=timeout, extra_args=extra_args
        )

        output = result.stdout + result.stderr
        if result.returncode != 0:
            print(f"STDOUT:\n{result.stdout[-3000:]}")
            print(f"STDERR:\n{result.stderr[-3000:]}")
        self.assertEqual(result.returncode, 0, f"{family} run_config failed")

        lines = output.split("\n")
        error_shapes, mismatch_shapes, ok_count, skip_count = _parse_benchmark_results(
            lines
        )

        all_results = _parse_all_benchmark_results(lines)
        csv_file = _save_results_csv(family, all_results)
        if csv_file:
            print(f"\n  Results CSV: {csv_file} ({len(all_results)} shapes)")

        failures = _format_failures(
            family, error_shapes, mismatch_shapes, skip_count, lines
        )
        self.assertEqual(
            len(failures),
            0,
            f"{family} run_config: {ok_count} OK, {skip_count} SKIP, "
            f"{len(error_shapes)} ERROR, {len(mismatch_shapes)} MISMATCH\n"
            + "\n".join(failures),
        )


if __name__ == "__main__":
    unittest.main(verbosity=2)
