#!/usr/bin/env python3
import argparse
import subprocess
import tempfile
import shutil
import statistics
from pathlib import Path

# This script compares the performance of two commits in the utf8cpp library
# by running benchmarks for different scenarios and functions.

# Example:
# ./bench/compare_commits.py HEAD~1 HEAD

SCENARIOS = ["ascii", "cyrillic", "mixed"]
FUNCTIONS = ["utf8::next", "utf8::unchecked::next", "utf8::find_invalid"]
BENCHMARK_SOURCE = Path(__file__).with_name("benchmark.cpp")

def run(cmd, cwd=None):
    result = subprocess.run(cmd, cwd=cwd, stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"Command failed: {' '.join(cmd)}\n{result.stderr}")
    return result.stdout

def checkout_commit(commit_ref):
    tempdir = tempfile.mkdtemp(prefix="utfcpp_")
    archive = subprocess.Popen(["git", "archive", commit_ref],
                               stdout=subprocess.PIPE)
    subprocess.run(["tar", "-xC", tempdir], stdin=archive.stdout)
    shutil.copy2(BENCHMARK_SOURCE, Path(tempdir) / "bench" / "benchmark.cpp")
    return tempdir

def build_benchmark(source_dir):
    build_dir = Path(source_dir) / "build"
    build_dir.mkdir(exist_ok=True)
    run(["cmake", "..", "-DCMAKE_BUILD_TYPE=Release",
         "-DUTF8CPP_ENABLE_BENCHMARKS=ON"], cwd=build_dir)
    run(["cmake", "--build", ".", "--config", "Release", "--target", "benchmark"],
        cwd=build_dir)
    exe = build_dir / "bench" / "benchmark"
    if not exe.exists():
        raise RuntimeError("Benchmark executable not found")
    return exe

def parse_csv_output(text):
    results = {}
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue

        parts = [p.strip() for p in line.split(",")]
        if parts[0] == "Function":
            continue
        if len(parts) < 4:
            continue

        func = parts[0]
        try:
            time_us = float(parts[1])
            mbps = float(parts[3])
        except ValueError:
            continue

        results[func] = {"time_us": time_us, "mbps": mbps}
    return results

def new_scenario_results():
    return {s: {f: {"time_us": [], "mbps": []} for f in FUNCTIONS}
            for s in SCENARIOS}

def add_benchmark_result(scenario_results, scenario, exe):
    parsed = parse_csv_output(run([str(exe), scenario]))
    missing_functions = set(FUNCTIONS) - parsed.keys()
    if missing_functions:
        missing = ", ".join(sorted(missing_functions))
        raise RuntimeError(f"Benchmark output is missing: {missing}")

    for func in FUNCTIONS:
        scenario_results[scenario][func]["time_us"].append(parsed[func]["time_us"])
        scenario_results[scenario][func]["mbps"].append(parsed[func]["mbps"])

def median_results(scenario_results):
    medians = {}
    for scenario in SCENARIOS:
        medians[scenario] = {}
        for func in FUNCTIONS:
            vals = scenario_results[scenario][func]
            if vals["time_us"]:
                medians[scenario][func] = {
                    "time_us": statistics.median(vals["time_us"]),
                    "mbps": statistics.median(vals["mbps"])
                }
    return medians

def run_benchmarks(exe1, exe2, runs):
    results1 = new_scenario_results()
    results2 = new_scenario_results()

    for scenario in SCENARIOS:
        print(f"  Running scenario: {scenario}")
        for run_index in range(runs):
            executions = [(results1, exe1), (results2, exe2)]
            if run_index % 2:
                executions.reverse()
            for results, exe in executions:
                add_benchmark_result(results, scenario, exe)

    return median_results(results1), median_results(results2)

def pct_change(old, new):
    if old == 0:
        return 0.0
    return (new - old) / old * 100.0

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("commit1")
    parser.add_argument("commit2")
    parser.add_argument("--runs", type=int, default=100,
                        help="Number of runs per scenario")
    args = parser.parse_args()
    if args.runs < 1:
        parser.error("--runs must be at least 1")

    commit1_sha = run(["git", "rev-parse", args.commit1]).strip()[:7]
    commit2_sha = run(["git", "rev-parse", args.commit2]).strip()[:7]

    print(f"Comparing commits:")
    print(f"  Commit1: {commit1_sha} ({args.commit1})")
    print(f"  Commit2: {commit2_sha} ({args.commit2})")
    print("")

    tempdirs = []
    try:
        print(f"Building {commit1_sha}...")
        src1 = checkout_commit(args.commit1)
        tempdirs.append(src1)
        exe1 = build_benchmark(src1)

        print(f"Building {commit2_sha}...")
        src2 = checkout_commit(args.commit2)
        tempdirs.append(src2)
        exe2 = build_benchmark(src2)

        print("Running interleaved benchmark samples...")
        results1, results2 = run_benchmarks(exe1, exe2, args.runs)

        print("=" * 80)
        print(f"Benchmark Comparison (median of {args.runs} runs): {commit1_sha} vs {commit2_sha}")
        print("=" * 80)

        for scenario in SCENARIOS:
            print(f"\nScenario: {scenario.upper()}")
            print("-" * 80)
            print(f"{'Function':<30} {'Commit1 MB/s':>15} {'Commit2 MB/s':>15} {'Change (%)':>12}  Result")
            print("-" * 80)

            for func in FUNCTIONS:
                old = results1[scenario][func]["mbps"]
                new = results2[scenario][func]["mbps"]
                change = pct_change(old, new)

                if change > 2:
                    verdict = "Commit2 is faster"
                elif change < -2:
                    verdict = "Commit2 is slower"
                else:
                    verdict = "Similar performance"

                print(f"{func:<30} {old:>15.2f} {new:>15.2f} {change:>12.2f}  {verdict}")

        print("\n" + "=" * 80)
        print("Interpretation:")
        print("  Change > +2%    Commit2 is faster")
        print("  Change < -2%    Commit2 is slower")
        print("  Otherwise       Similar performance")
        print("=" * 80)
    finally:
        for d in tempdirs:
            shutil.rmtree(d, ignore_errors=True)

if __name__ == "__main__":
    main()
