#!/usr/bin/env python3
"""
Benchmark script for wasot CLI.

Generates test files of various sizes and benchmarks wasot CLI operations.
Results are stored in a JSON file.

Usage:
    ./scripts/benchmark_cli.py                  # Full benchmark (5 runs per size)
    ./scripts/benchmark_cli.py --quick          # Quick mode (1 run, no storage)
    ./scripts/benchmark_cli.py --quick --assert # Quick mode with assertions
    ./scripts/benchmark_cli.py --sizes 1kb 1mb  # Specific sizes only
    ./scripts/benchmark_cli.py --regenerate     # Force regenerate test files
    ./scripts/benchmark_cli.py --large          # Include 1GB and 4GB files (1 run each)
    ./scripts/benchmark_cli.py --extra 50GB      # Add custom size (reuses existing file)

Performance assertions (with --assert):
    1kb:   < 0.1s (100ms)
    10kb:  < 0.1s (100ms)
    500kb: < 0.5s (500ms)
    1mb:   < 1.0s (1 second)
    100mb: < 2.0s (2 seconds)
    1gb:   < 20s (with --large)
    4gb:   < 80s (with --large)
"""

import os
import subprocess
import sys
from pathlib import Path

# Determine paths early for potential re-exec
SCRIPT_DIR = Path(__file__).parent.resolve()
PROJECT_ROOT = SCRIPT_DIR.parent
VENV_PYTHON = PROJECT_ROOT / ".venv" / "bin" / "python"


def ensure_venv():
    """Re-exec with venv Python if not already running in it."""
    if not VENV_PYTHON.exists():
        print(f"Error: Project venv not found at {VENV_PYTHON}", file=sys.stderr)
        print("Please run: uv venv && uv pip install tqdm", file=sys.stderr)
        sys.exit(1)

    # Check if we're already running in the venv
    current_python = Path(sys.executable).resolve()
    if current_python != VENV_PYTHON.resolve():
        # Re-exec with venv Python
        os.execv(str(VENV_PYTHON), [str(VENV_PYTHON)] + sys.argv)


# Ensure we're running in the venv before importing dependencies
ensure_venv()

# Now safe to import dependencies that may only exist in venv
import argparse
import json
import time
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any

GENERATORS_DIR = SCRIPT_DIR / "generators"
sys.path.insert(0, str(GENERATORS_DIR))


@dataclass
class BenchmarkResult:
    """Result of a single benchmark run."""

    command: str
    file_size_bytes: int
    file_lines: int
    file_name: str
    duration_seconds: float
    lines_per_sec: float
    exit_code: int
    error: str | None = None


@dataclass
class BenchmarkSuite:
    """Collection of benchmark results."""

    timestamp: str
    wasot_version: str
    system_info: dict[str, Any]
    results: list[dict[str, Any]] = field(default_factory=list)


# File sizes to benchmark: (size_spec, max_runs override, max_time_seconds for assertions)
FILE_SIZES = {
    "1kb": ("1KB", None, 0.1),      # Must complete in 100ms
    "10kb": ("10KB", None, 0.1),    # Must complete in 100ms
    "500kb": ("500KB", None, 0.5),  # Must complete in 500ms
    "1mb": ("1MB", None, 1.0),      # Must complete in 1s
    "100mb": ("100MB", 3, 2.0),     # Must complete in 2s
}

# Large file sizes (only with --large flag): (size_spec, max_runs, max_time_seconds)
LARGE_FILE_SIZES = {
    "1gb": ("1GB", 1, 20.0),        # Must complete in 20s, 1 run only
    "4gb": ("4GB", 1, 80.0),        # Must complete in 80s, 1 run only
}

# Available output formats
AVAILABLE_FORMATS = {
    "xml": ["-t", "xml"],
    "jsonl": ["-t", "jsonl"],
    "wasot": ["-t", "wasot"],
    "html": ["-t", "html"],
    "binary": ["-t", "binary"],
}


def get_file_md5(path: Path) -> str:
    """Calculate MD5 hash of a file."""
    import hashlib

    md5 = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            md5.update(chunk)
    return md5.hexdigest()


def get_git_commit() -> str:
    """Get current git commit hash."""
    try:
        result = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=PROJECT_ROOT,
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode == 0:
            return result.stdout.strip()
    except Exception:
        pass
    return "unknown"


def get_git_dirty() -> bool:
    """Check if git working directory has uncommitted changes."""
    try:
        result = subprocess.run(
            ["git", "status", "--porcelain"],
            cwd=PROJECT_ROOT,
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode == 0:
            return len(result.stdout.strip()) > 0
    except Exception:
        pass
    return False


def build_wasot() -> Path:
    """Build wasot in release mode and return the binary path."""
    print("Building wasot (release mode)...")

    result = subprocess.run(
        ["cargo", "build", "--release", "--bin", "wasot"],
        cwd=PROJECT_ROOT,
        capture_output=True,
        text=True,
    )

    if result.returncode != 0:
        print("Build failed:", file=sys.stderr)
        print(result.stderr, file=sys.stderr)
        sys.exit(1)

    print("Build complete.")

    release_binary = PROJECT_ROOT / "target" / "release" / "wasot"
    if not release_binary.exists():
        raise FileNotFoundError(f"Binary not found after build: {release_binary}")

    return release_binary


def get_wasot_version(binary: Path) -> str:
    """Get wasot version string."""
    try:
        result = subprocess.run(
            [str(binary), "--version"],
            capture_output=True,
            text=True,
            timeout=10,
        )
        return result.stdout.strip()
    except Exception as e:
        return f"unknown ({e})"


def get_system_info() -> dict[str, Any]:
    """Gather system information for the benchmark report."""
    import platform

    info = {
        "platform": platform.system(),
        "platform_release": platform.release(),
        "platform_version": platform.version(),
        "architecture": platform.machine(),
        "processor": platform.processor(),
        "python_version": platform.python_version(),
    }

    # Try to get CPU info on Linux
    if platform.system() == "Linux":
        try:
            with open("/proc/cpuinfo") as f:
                for line in f:
                    if line.startswith("model name"):
                        info["cpu_model"] = line.split(":")[1].strip()
                        break
        except Exception:
            pass

        try:
            with open("/proc/meminfo") as f:
                for line in f:
                    if line.startswith("MemTotal"):
                        mem_kb = int(line.split()[1])
                        info["memory_gb"] = round(mem_kb / 1024 / 1024, 1)
                        break
        except Exception:
            pass

    return info


def generate_test_file(output_path: Path, size: str, seed: int = 42) -> int:
    """Generate a test file using the generator script.

    Returns the actual file size in bytes.
    """
    generator_script = GENERATORS_DIR / "generate_large_file.py"

    if not generator_script.exists():
        raise FileNotFoundError(f"Generator script not found: {generator_script}")

    # Ensure parent directory exists
    output_path.parent.mkdir(parents=True, exist_ok=True)

    cmd = [
        str(VENV_PYTHON),
        str(generator_script),
        "-o",
        str(output_path),
        "-s",
        size,
        "--seed",
        str(seed),
    ]

    print(f"  Generating {output_path.name} ({size})...", end=" ", flush=True)

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print("FAILED")
        print(f"Error: {result.stderr}", file=sys.stderr)
        raise RuntimeError(f"Failed to generate {output_path}")

    actual_size = output_path.stat().st_size
    print(f"done ({actual_size:,} bytes)")

    return actual_size


def run_benchmark(
    binary: Path,
    input_file: Path,
    command_args: list[str],
    num_lines: int,
    warmup: bool = False,
) -> BenchmarkResult:
    """Run a single benchmark and return the result."""
    file_size = input_file.stat().st_size

    cmd = [str(binary)] + command_args + [str(input_file)]

    start_time = time.perf_counter()

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=600,  # 10 minute timeout
        )
        end_time = time.perf_counter()

        duration = end_time - start_time
        lines_per_sec = num_lines / duration if duration > 0 else 0

        return BenchmarkResult(
            command=" ".join(command_args),
            file_size_bytes=file_size,
            file_lines=num_lines,
            file_name=input_file.name,
            duration_seconds=round(duration, 4),
            lines_per_sec=round(lines_per_sec, 1),
            exit_code=result.returncode,
            error=result.stderr if result.returncode != 0 else None,
        )

    except subprocess.TimeoutExpired:
        return BenchmarkResult(
            command=" ".join(command_args),
            file_size_bytes=file_size,
            file_lines=num_lines,
            file_name=input_file.name,
            duration_seconds=-1,
            lines_per_sec=0,
            exit_code=-1,
            error="Timeout after 600 seconds",
        )
    except Exception as e:
        return BenchmarkResult(
            command=" ".join(command_args),
            file_size_bytes=file_size,
            file_lines=num_lines,
            file_name=input_file.name,
            duration_seconds=-1,
            lines_per_sec=0,
            exit_code=-1,
            error=str(e),
        )


def format_size(size_bytes: int) -> str:
    """Format byte size as human-readable string."""
    for unit in ["B", "KB", "MB", "GB"]:
        if size_bytes < 1024:
            return f"{size_bytes:.1f} {unit}"
        size_bytes /= 1024
    return f"{size_bytes:.1f} TB"


def count_lines(file_path: Path) -> int:
    """Count lines in a file."""
    with open(file_path, "rb") as f:
        return sum(1 for _ in f)


def format_rate(lines_per_sec: float) -> str:
    """Format lines per second as human-readable string."""
    if lines_per_sec >= 1_000_000:
        return f"{lines_per_sec / 1_000_000:.1f}M lines/s"
    elif lines_per_sec >= 1_000:
        return f"{lines_per_sec / 1_000:.1f}K lines/s"
    else:
        return f"{lines_per_sec:.0f} lines/s"


def print_report(data: dict, assertion_failures: list | None = None) -> None:
    """Print a nice human-readable benchmark report."""
    print("\n" + "=" * 70)
    print("                        BENCHMARK REPORT")
    print("=" * 70)

    # Metadata
    print(f"\n  Date:        {data['timestamp']}")
    print(f"  Version:     {data['wasot_version']}")
    print(f"  Binary MD5:  {data['binary_md5']}")
    print(f"  Git Commit:  {data['git_commit'][:12]}{'  (dirty)' if data.get('git_dirty') else ''}")

    # System info
    sys_info = data.get("system_info", {})
    if sys_info:
        print(f"\n  System:      {sys_info.get('platform', 'Unknown')} {sys_info.get('platform_release', '')}")
        if "cpu_model" in sys_info:
            print(f"  CPU:         {sys_info['cpu_model']}")
        if "memory_gb" in sys_info:
            print(f"  Memory:      {sys_info['memory_gb']} GB")

    # Results table
    print("\n" + "-" * 70)
    print(f"  {'FILE':<10} {'LINES':<12} {'FORMAT':<10} {'TIME':<12} {'RATE':<18} {'STATUS'}")
    print("-" * 70)

    for result in data["results"]:
        file_name = result["file"]
        num_lines = f"{result.get('file_lines', 0):,}"
        command = result["command"].replace("to-", "")

        if "avg_duration_seconds" in result:
            time_str = f"{result['avg_duration_seconds']:.3f}s"
            rate = format_rate(result.get("avg_lines_per_sec", 0))

            # Show assertion status if available
            if result.get("assertion_failed"):
                max_time = result.get("max_time_assertion", 0)
                status = f"FAIL (>{max_time}s)"
            elif "max_time_assertion" in result:
                status = "OK"
            else:
                status = f"({result['successful_runs']}/{result['runs']})"
        else:
            time_str = "FAILED"
            rate = result.get("error", "Unknown error")[:20]
            status = "ERROR"

        print(f"  {file_name:<10} {num_lines:<12} {command:<10} {time_str:<12} {rate:<18} {status}")

    print("-" * 70)

    # Print assertion summary if there are failures
    if assertion_failures:
        print("\n  ASSERTION FAILURES:")
        for failure in assertion_failures:
            print(f"    - {failure['file']} ({failure['command']}): "
                  f"{failure['avg_duration']:.3f}s > {failure['max_allowed']}s limit")
        print()
    else:
        print()


def main():
    parser = argparse.ArgumentParser(
        description="Benchmark wasot CLI with various file sizes"
    )
    parser.add_argument(
        "-d",
        "--data-dir",
        type=Path,
        default=Path("data/bench"),
        help="Directory for generated test files and results (default: data/bench)",
    )
    parser.add_argument(
        "--regenerate",
        action="store_true",
        help="Force regeneration of test files even if they exist",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Run benchmark even if results for this binary already exist",
    )
    parser.add_argument(
        "--quick",
        action="store_true",
        help="Quick mode: one run per input, no file storage, just prints report",
    )
    parser.add_argument(
        "--warmup",
        action="store_true",
        help="Run a warmup pass before benchmarking",
    )
    parser.add_argument(
        "--runs",
        type=int,
        default=5,
        help="Number of benchmark runs per test (default: 5)",
    )
    parser.add_argument(
        "--formats",
        type=str,
        default="wasot",
        help="Comma-separated list of output formats to benchmark (default: wasot). "
             f"Available: {','.join(AVAILABLE_FORMATS.keys())}",
    )
    parser.add_argument(
        "--sizes",
        nargs="+",
        default=None,  # None means "all available sizes"
        help="File sizes to benchmark (default: all). "
             f"Standard: {', '.join(FILE_SIZES.keys())}. "
             f"With --large: also {', '.join(LARGE_FILE_SIZES.keys())}",
    )
    parser.add_argument(
        "--assert",
        action="store_true",
        dest="check_assertions",
        help="Check performance assertions (fail if benchmarks exceed time limits)",
    )
    parser.add_argument(
        "--large",
        action="store_true",
        help="Include large file sizes (1GB, 4GB) with 1 run each. "
             "Warning: generates multi-GB files and takes several minutes.",
    )
    parser.add_argument(
        "--extra",
        type=str,
        metavar="SIZE",
        help="Add a custom extra file size (e.g., 50GB). "
             "Treated as large file: 1 run, no assertion. "
             "Reuses existing file if found.",
    )
    args = parser.parse_args()

    # Merge large file sizes if --large flag is set
    all_file_sizes = dict(FILE_SIZES)
    if args.large:
        all_file_sizes.update(LARGE_FILE_SIZES)

    # Add custom extra size if provided
    extra_size_name = None
    if args.extra:
        # Normalize the size spec (e.g., "50gb" -> "50GB")
        extra_spec = args.extra.upper()
        extra_size_name = args.extra.lower().replace(" ", "")
        # No max_time assertion for custom sizes (use very large value)
        all_file_sizes[extra_size_name] = (extra_spec, 1, float("inf"))

    # Determine which sizes to benchmark
    if args.sizes is None:
        # Default: standard sizes (not including extra)
        args.sizes = list(FILE_SIZES.keys())
        if args.large:
            args.sizes += list(LARGE_FILE_SIZES.keys())

    # Add extra size to the benchmark list if specified
    if extra_size_name:
        if extra_size_name not in args.sizes:
            args.sizes.append(extra_size_name)
    else:
        # Validate user-specified sizes
        invalid_sizes = [s for s in args.sizes if s not in all_file_sizes]
        if invalid_sizes:
            print(f"Error: Unknown size(s): {', '.join(invalid_sizes)}", file=sys.stderr)
            available = list(FILE_SIZES.keys())
            if args.large:
                available += list(LARGE_FILE_SIZES.keys())
            print(f"Available sizes: {', '.join(available)}", file=sys.stderr)
            if not args.large and any(s in LARGE_FILE_SIZES for s in invalid_sizes):
                print("Hint: Use --large to enable 1gb and 4gb sizes", file=sys.stderr)
            sys.exit(1)

    # Parse and validate formats
    formats = [f.strip() for f in args.formats.split(",")]
    invalid_formats = [f for f in formats if f not in AVAILABLE_FORMATS]
    if invalid_formats:
        print(f"Error: Unknown format(s): {', '.join(invalid_formats)}", file=sys.stderr)
        print(f"Available formats: {', '.join(AVAILABLE_FORMATS.keys())}", file=sys.stderr)
        sys.exit(1)

    benchmark_commands = [
        {"name": f"to-{fmt}", "args": AVAILABLE_FORMATS[fmt]}
        for fmt in formats
    ]

    # Build wasot
    wasot_binary = build_wasot()
    print(f"Using wasot binary: {wasot_binary}")

    wasot_version = get_wasot_version(wasot_binary)
    print(f"Wasot version: {wasot_version}")

    # Calculate binary hash and check for existing results
    binary_md5 = get_file_md5(wasot_binary)
    print(f"Binary MD5: {binary_md5}")

    results_base = args.data_dir / "results"
    binary_md5_short = binary_md5[:8]
    if results_base.exists() and not args.force and not args.quick:
        existing = [f for f in results_base.iterdir() if f.is_file() and binary_md5_short in f.name]
        if existing:
            existing_file = existing[0]
            print(f"\nBenchmark already exists for this binary: {existing_file.name}")
            print("Use --force to re-run anyway.")

            # Load and print existing report
            with open(existing_file) as f:
                existing_data = json.load(f)
            print_report(existing_data, assertion_failures=[])
            sys.exit(0)

    # Ensure data directory exists
    args.data_dir.mkdir(parents=True, exist_ok=True)

    # Generate test files
    print("\n=== Generating test files ===")
    test_files: dict[str, Path] = {}
    file_sizes: dict[str, int] = {}
    file_lines: dict[str, int] = {}

    for size_name in args.sizes:
        size_spec, _, _ = all_file_sizes[size_name]
        file_path = args.data_dir / f"{size_name}.wasot"

        if file_path.exists() and not args.regenerate:
            print(f"  Using existing {file_path.name}")
            file_sizes[size_name] = file_path.stat().st_size
        else:
            file_sizes[size_name] = generate_test_file(file_path, size_spec)

        file_lines[size_name] = count_lines(file_path)
        test_files[size_name] = file_path

    # Run benchmarks
    print("\n=== Running benchmarks ===")

    suite = BenchmarkSuite(
        timestamp=datetime.now().isoformat(),
        wasot_version=wasot_version,
        system_info=get_system_info(),
    )

    total_benchmarks = len(args.sizes) * len(benchmark_commands) * args.runs
    current = 0

    assertion_failures = []

    for size_name in args.sizes:
        file_path = test_files[size_name]
        file_size = file_sizes[size_name]
        num_lines = file_lines[size_name]
        _, max_runs_override, max_time = all_file_sizes[size_name]
        if args.quick:
            num_runs = 1
        elif max_runs_override:
            num_runs = min(args.runs, max_runs_override)
        else:
            num_runs = args.runs

        print(f"\n  File: {file_path.name} ({format_size(file_size)}, {num_lines:,} lines)")

        for cmd_info in benchmark_commands:
            cmd_name = cmd_info["name"]
            cmd_args = cmd_info["args"]

            # Warmup run
            if args.warmup:
                print(f"    {cmd_name}: warming up...", end=" ", flush=True)
                run_benchmark(wasot_binary, file_path, cmd_args, num_lines, warmup=True)
                print("done")

            # Actual benchmark runs
            run_results = []
            for run_num in range(num_runs):
                current += 1
                print(
                    f"    {cmd_name}: run {run_num + 1}/{num_runs}...",
                    end=" ",
                    flush=True,
                )

                result = run_benchmark(wasot_binary, file_path, cmd_args, num_lines)
                run_results.append(result)

                if result.error:
                    print(f"ERROR: {result.error}")
                else:
                    print(
                        f"{result.duration_seconds:.3f}s "
                        f"({format_rate(result.lines_per_sec)})"
                    )

            # Calculate statistics
            successful_runs = [r for r in run_results if r.exit_code == 0]

            if successful_runs:
                durations = [r.duration_seconds for r in successful_runs]
                rates = [r.lines_per_sec for r in successful_runs]

                avg_duration = sum(durations) / len(durations)
                min_duration = min(durations)
                max_duration = max(durations)
                avg_lines_per_sec = sum(rates) / len(rates)

                result_entry = {
                    "file": size_name,
                    "file_size_bytes": file_size,
                    "file_lines": num_lines,
                    "command": cmd_name,
                    "runs": num_runs,
                    "successful_runs": len(successful_runs),
                    "avg_duration_seconds": round(avg_duration, 4),
                    "min_duration_seconds": round(min_duration, 4),
                    "max_duration_seconds": round(max_duration, 4),
                    "avg_lines_per_sec": round(avg_lines_per_sec, 1),
                    "max_time_assertion": max_time,
                }

                # Check assertion
                if avg_duration > max_time:
                    result_entry["assertion_failed"] = True
                    assertion_failures.append({
                        "file": size_name,
                        "command": cmd_name,
                        "avg_duration": round(avg_duration, 4),
                        "max_allowed": max_time,
                    })

                suite.results.append(result_entry)
            else:
                suite.results.append(
                    {
                        "file": size_name,
                        "file_size_bytes": file_size,
                        "file_lines": num_lines,
                        "command": cmd_name,
                        "runs": num_runs,
                        "successful_runs": 0,
                        "error": run_results[0].error if run_results else "No runs",
                    }
                )

    # Gather metadata
    git_commit = get_git_commit()
    git_dirty = get_git_dirty()

    output_data = {
        "timestamp": suite.timestamp,
        "wasot_version": suite.wasot_version,
        "binary_md5": binary_md5,
        "git_commit": git_commit,
        "git_dirty": git_dirty,
        "system_info": suite.system_info,
        "benchmark_config": {
            "runs_per_test": args.runs,
            "warmup": args.warmup,
            "file_sizes": {k: all_file_sizes[k][0] for k in args.sizes},
            "commands": [c["name"] for c in benchmark_commands],
        },
        "results": suite.results,
    }

    # Save results (skip in quick mode)
    if not args.quick:
        timestamp_str = datetime.now().strftime("%Y-%m-%d_%H%M%S")
        results_dir = args.data_dir / "results"
        results_dir.mkdir(parents=True, exist_ok=True)

        output_file = results_dir / f"{timestamp_str}_{binary_md5[:8]}.json"

        with open(output_file, "w") as f:
            json.dump(output_data, f, indent=2)

        print(f"\nResults saved to: {output_file}")

    # Print report
    print_report(output_data, assertion_failures)

    # Check assertions and exit with error if any failed
    if args.check_assertions and assertion_failures:
        print(f"ERROR: {len(assertion_failures)} assertion(s) failed!")
        sys.exit(1)


if __name__ == "__main__":
    main()
