#!/usr/bin/env python3
"""
Unified performance tooling for wasot.

This script provides a single entry point for all performance-related tasks:
- Benchmarking with hyperfine (statistical rigor)
- CPU profiling with samply/flamegraph
- Memory profiling with DHAT
- Comparison between runs/commits

Usage:
    ./scripts/perf.py bench                    # Quick benchmark with hyperfine
    ./scripts/perf.py bench --compare HEAD~1   # Compare with previous commit
    ./scripts/perf.py cpu                      # CPU profile with samply
    ./scripts/perf.py cpu --flamegraph         # Generate flamegraph
    ./scripts/perf.py mem                      # Memory profile with DHAT
    ./scripts/perf.py report                   # Show last benchmark results
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Optional

# Paths
SCRIPT_DIR = Path(__file__).parent.resolve()
PROJECT_ROOT = SCRIPT_DIR.parent
VENV_PYTHON = PROJECT_ROOT / ".venv" / "bin" / "python"
DATA_DIR = PROJECT_ROOT / "data" / "bench"
RESULTS_DIR = DATA_DIR / "results"
PROFILES_DIR = DATA_DIR / "profiles"
MEMORY_DIR = DATA_DIR / "memory"

# Test files
TEST_FILES = {
    "small": DATA_DIR / "10kb.wasot",
    "medium": DATA_DIR / "1mb.wasot",
    "large": DATA_DIR / "100mb.wasot",
}

DEFAULT_TEST = "medium"


def ensure_venv():
    """Re-exec with venv Python if available."""
    if VENV_PYTHON.exists():
        current = Path(sys.executable).resolve()
        if current != VENV_PYTHON.resolve():
            os.execv(str(VENV_PYTHON), [str(VENV_PYTHON)] + sys.argv)


ensure_venv()


def run(cmd: list, **kwargs) -> subprocess.CompletedProcess:
    """Run command and return result."""
    return subprocess.run(cmd, **kwargs)


def run_output(cmd: list, **kwargs) -> str:
    """Run command and return stdout."""
    result = run(cmd, capture_output=True, text=True, **kwargs)
    return result.stdout.strip()


def get_git_commit() -> str:
    """Get current git commit hash."""
    try:
        return run_output(["git", "rev-parse", "--short", "HEAD"], cwd=PROJECT_ROOT)
    except Exception:
        return "unknown"


def get_git_dirty() -> bool:
    """Check if working directory is dirty."""
    try:
        result = run_output(["git", "status", "--porcelain"], cwd=PROJECT_ROOT)
        return len(result) > 0
    except Exception:
        return False


def check_tool(name: str) -> bool:
    """Check if a tool is available."""
    return shutil.which(name) is not None


def check_perf_permissions() -> tuple[bool, str]:
    """Check if perf profiling is allowed."""
    try:
        with open("/proc/sys/kernel/perf_event_paranoid") as f:
            level = int(f.read().strip())
        if level <= 1:
            return True, f"perf_event_paranoid={level} (OK)"
        else:
            return False, f"perf_event_paranoid={level} (too restrictive, need <=1)"
    except Exception as e:
        return False, f"Cannot read perf_event_paranoid: {e}"


def ensure_test_files():
    """Ensure test files exist."""
    missing = [name for name, path in TEST_FILES.items() if not path.exists()]
    if missing:
        print(f"Generating test files: {', '.join(missing)}")
        run([sys.executable, SCRIPT_DIR / "benchmark_cli.py", "--sizes"] + missing + ["--runs", "1", "--quick"],
            cwd=PROJECT_ROOT)


def build_release(features: list[str] = None) -> Path:
    """Build wasot in release mode with optional features."""
    cmd = ["cargo", "build", "--release", "--bin", "wasot"]
    if features:
        cmd.extend(["--features", ",".join(features)])

    env = os.environ.copy()
    # Always include debug info for profiling
    env["CARGO_PROFILE_RELEASE_DEBUG"] = "true"
    env["CARGO_PROFILE_RELEASE_STRIP"] = "false"

    print(f"Building wasot{'with ' + ','.join(features) if features else ''}...")
    result = run(cmd, cwd=PROJECT_ROOT, env=env, capture_output=True, text=True)
    if result.returncode != 0:
        print("Build failed:", file=sys.stderr)
        print(result.stderr, file=sys.stderr)
        sys.exit(1)

    return PROJECT_ROOT / "target" / "release" / "wasot"


def get_binary_info(binary: Path) -> dict:
    """Get information about binary."""
    import hashlib
    with open(binary, "rb") as f:
        md5 = hashlib.md5(f.read()).hexdigest()[:8]

    return {
        "path": str(binary),
        "md5": md5,
        "git_commit": get_git_commit(),
        "git_dirty": get_git_dirty(),
    }


# =============================================================================
# BENCHMARKING
# =============================================================================

def cmd_bench(args):
    """Run benchmarks with hyperfine."""
    if not check_tool("hyperfine"):
        print("Error: hyperfine not installed. Install with: cargo install hyperfine")
        sys.exit(1)

    ensure_test_files()

    # Build
    binary = build_release()
    info = get_binary_info(binary)
    print(f"Binary: {info['md5']} (commit: {info['git_commit']}{'*' if info['git_dirty'] else ''})")

    # Select test file
    test_file = TEST_FILES.get(args.size, TEST_FILES[DEFAULT_TEST])
    if not test_file.exists():
        print(f"Error: Test file not found: {test_file}")
        sys.exit(1)

    # Count lines
    with open(test_file, "rb") as f:
        num_lines = sum(1 for _ in f)

    print(f"\nBenchmarking with: {test_file.name} ({num_lines:,} lines)")

    # Build hyperfine command
    cmd = [
        "hyperfine",
        "--warmup", str(args.warmup),
        "--runs", str(args.runs),
        "--shell", "none",
    ]

    if args.export_json:
        cmd.extend(["--export-json", str(args.export_json)])

    if args.export_markdown:
        cmd.extend(["--export-markdown", str(args.export_markdown)])

    # The actual command to benchmark
    bench_cmd = f"{binary} {test_file} -t wasot"

    if args.compare:
        # Build comparison version
        print(f"\nBuilding comparison version ({args.compare})...")

        # Stash current changes if dirty
        was_dirty = info['git_dirty']
        if was_dirty:
            run(["git", "stash"], cwd=PROJECT_ROOT)

        try:
            # Checkout comparison ref
            run(["git", "checkout", args.compare], cwd=PROJECT_ROOT, capture_output=True)

            # Build comparison binary
            compare_binary = PROJECT_ROOT / "target" / "release" / "wasot-compare"
            run(["cargo", "build", "--release", "--bin", "wasot"], cwd=PROJECT_ROOT, capture_output=True)
            shutil.copy(PROJECT_ROOT / "target" / "release" / "wasot", compare_binary)

            # Restore original
            run(["git", "checkout", "-"], cwd=PROJECT_ROOT, capture_output=True)
            if was_dirty:
                run(["git", "stash", "pop"], cwd=PROJECT_ROOT)

            # Rebuild current
            build_release()

            # Add comparison command
            compare_cmd = f"{compare_binary} {test_file} -t wasot"
            cmd.extend([
                "--command-name", f"current ({info['git_commit']})",
                "--command-name", f"compare ({args.compare})",
                bench_cmd,
                compare_cmd,
            ])
        except Exception as e:
            # Restore on error
            run(["git", "checkout", "-"], cwd=PROJECT_ROOT, capture_output=True)
            if was_dirty:
                run(["git", "stash", "pop"], cwd=PROJECT_ROOT)
            raise e
    else:
        cmd.append(bench_cmd)

    # Run hyperfine
    print()
    result = run(cmd, cwd=PROJECT_ROOT)

    # Calculate lines/sec from hyperfine output
    if args.export_json and Path(args.export_json).exists():
        with open(args.export_json) as f:
            data = json.load(f)

        print("\n" + "=" * 60)
        print("THROUGHPUT ANALYSIS")
        print("=" * 60)

        for r in data.get("results", []):
            mean_time = r.get("mean", 0)
            if mean_time > 0:
                lines_per_sec = num_lines / mean_time
                name = r.get("command", "unknown")
                if lines_per_sec >= 1_000_000:
                    rate_str = f"{lines_per_sec/1_000_000:.2f}M lines/s"
                elif lines_per_sec >= 1_000:
                    rate_str = f"{lines_per_sec/1_000:.1f}K lines/s"
                else:
                    rate_str = f"{lines_per_sec:.0f} lines/s"
                print(f"  {name}: {rate_str} (mean: {mean_time*1000:.1f}ms)")
        print()

    return result.returncode


# =============================================================================
# CPU PROFILING
# =============================================================================

def cmd_cpu(args):
    """CPU profiling with samply or flamegraph."""
    ensure_test_files()

    # Check permissions
    perf_ok, perf_msg = check_perf_permissions()
    if not perf_ok:
        print(f"Error: {perf_msg}")
        print()
        print("To fix, run:")
        print("  echo -1 | sudo tee /proc/sys/kernel/perf_event_paranoid")
        print()
        print("To make permanent, add to /etc/sysctl.conf:")
        print("  kernel.perf_event_paranoid = -1")
        sys.exit(1)

    # Build with debug info
    binary = build_release()
    info = get_binary_info(binary)

    test_file = TEST_FILES.get(args.size, TEST_FILES[DEFAULT_TEST])

    # Create profiles directory
    PROFILES_DIR.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    if args.flamegraph:
        if not check_tool("flamegraph"):
            print("Error: flamegraph not installed. Install with: cargo install flamegraph")
            sys.exit(1)

        output_svg = PROFILES_DIR / f"flamegraph_{timestamp}_{info['md5']}.svg"

        print(f"Generating flamegraph: {output_svg.name}")
        cmd = [
            "flamegraph",
            "-o", str(output_svg),
            "--", str(binary), str(test_file), "-t", "wasot"
        ]

        result = run(cmd, cwd=PROJECT_ROOT, capture_output=True, text=True)
        if result.returncode == 0:
            print(f"Flamegraph saved: {output_svg}")
            print(f"Open in browser: file://{output_svg}")
        else:
            print("Flamegraph generation failed:", file=sys.stderr)
            print(result.stderr, file=sys.stderr)
            return 1
    else:
        if not check_tool("samply"):
            print("Error: samply not installed. Install with: cargo install samply")
            sys.exit(1)

        print(f"Starting samply profiler (will open browser)...")
        print(f"Profiling: {binary} {test_file}")

        cmd = ["samply", "record", str(binary), str(test_file), "-t", "wasot"]
        run(cmd, cwd=PROJECT_ROOT)

    return 0


# =============================================================================
# MEMORY PROFILING
# =============================================================================

def cmd_mem(args):
    """Memory profiling with DHAT."""
    ensure_test_files()

    # Build with DHAT
    binary = build_release(features=["dhat-heap"])
    info = get_binary_info(binary)

    test_file = TEST_FILES.get(args.size, TEST_FILES[DEFAULT_TEST])

    # Create memory directory
    MEMORY_DIR.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    print(f"Running DHAT profile on: {test_file.name}")

    # Run with DHAT
    result = run(
        [str(binary), str(test_file), "-t", "wasot"],
        cwd=PROJECT_ROOT,
        capture_output=True,
        text=True,
    )

    # Show DHAT output
    for line in result.stderr.split('\n'):
        if line.startswith('dhat:'):
            print(f"  {line}")

    # Move dhat-heap.json to profiles
    dhat_src = PROJECT_ROOT / "dhat-heap.json"
    if dhat_src.exists():
        dhat_dst = MEMORY_DIR / f"dhat_{timestamp}_{info['md5']}.json"
        shutil.copy(dhat_src, dhat_dst)
        print(f"\nProfile saved: {dhat_dst.name}")

        # Analyze
        if not args.no_analyze:
            print()
            run([sys.executable, SCRIPT_DIR / "profile.py", "--analyze"], cwd=PROJECT_ROOT)

    return 0


# =============================================================================
# STATUS/REPORT
# =============================================================================

def cmd_status(args):
    """Show profiling environment status."""
    print("=" * 60)
    print("PERFORMANCE TOOLING STATUS")
    print("=" * 60)
    print()

    # Tools
    print("Tools:")
    tools = [
        ("hyperfine", "Statistical benchmarking"),
        ("samply", "Interactive CPU profiling"),
        ("flamegraph", "CPU flamegraph generation"),
        ("perf", "Linux performance counters"),
    ]
    for tool, desc in tools:
        status = "✓" if check_tool(tool) else "✗"
        print(f"  {status} {tool:<12} - {desc}")

    print()

    # Perf permissions
    perf_ok, perf_msg = check_perf_permissions()
    status = "✓" if perf_ok else "✗"
    print(f"Permissions:")
    print(f"  {status} {perf_msg}")

    if not perf_ok:
        print()
        print("  To enable CPU profiling:")
        print("    echo -1 | sudo tee /proc/sys/kernel/perf_event_paranoid")

    print()

    # Test files
    print("Test Files:")
    for name, path in TEST_FILES.items():
        if path.exists():
            size = path.stat().st_size
            with open(path, "rb") as f:
                lines = sum(1 for _ in f)
            print(f"  ✓ {name:<8} - {path.name} ({lines:,} lines, {size/1024:.1f} KB)")
        else:
            print(f"  ✗ {name:<8} - not generated")

    print()

    # Recent profiles
    print("Recent Profiles:")

    if RESULTS_DIR.exists():
        results = sorted(RESULTS_DIR.glob("*.json"), reverse=True)[:3]
        for r in results:
            print(f"  📊 {r.name}")

    if PROFILES_DIR.exists():
        profiles = sorted(PROFILES_DIR.glob("*.svg"), reverse=True)[:3]
        for p in profiles:
            print(f"  🔥 {p.name}")

    if MEMORY_DIR.exists():
        memory = sorted(MEMORY_DIR.glob("*.json"), reverse=True)[:3]
        for m in memory:
            print(f"  💾 {m.name}")

    print()


# =============================================================================
# MAIN
# =============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Unified performance tooling for wasot",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Commands:
  bench     Run benchmarks with hyperfine
  cpu       CPU profiling with samply/flamegraph
  mem       Memory profiling with DHAT
  status    Show tooling status

Examples:
  %(prog)s bench                      # Quick benchmark
  %(prog)s bench --compare HEAD~1     # Compare with previous commit
  %(prog)s bench --size large         # Benchmark with 100MB file
  %(prog)s cpu                        # Interactive CPU profile
  %(prog)s cpu --flamegraph           # Generate flamegraph SVG
  %(prog)s mem                        # Memory profile with DHAT
  %(prog)s status                     # Show tool availability
"""
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # bench subcommand
    bench_parser = subparsers.add_parser("bench", help="Run benchmarks with hyperfine")
    bench_parser.add_argument("--runs", "-r", type=int, default=10, help="Number of runs (default: 10)")
    bench_parser.add_argument("--warmup", "-w", type=int, default=3, help="Warmup runs (default: 3)")
    bench_parser.add_argument("--size", "-s", choices=["small", "medium", "large"], default="medium",
                              help="Test file size (default: medium/1MB)")
    bench_parser.add_argument("--compare", "-c", metavar="REF",
                              help="Compare with git ref (e.g., HEAD~1, main)")
    bench_parser.add_argument("--export-json", "-j", type=Path, help="Export results to JSON")
    bench_parser.add_argument("--export-markdown", "-m", type=Path, help="Export results to Markdown")

    # cpu subcommand
    cpu_parser = subparsers.add_parser("cpu", help="CPU profiling")
    cpu_parser.add_argument("--flamegraph", "-f", action="store_true",
                           help="Generate flamegraph SVG instead of interactive samply")
    cpu_parser.add_argument("--size", "-s", choices=["small", "medium", "large"], default="medium",
                           help="Test file size (default: medium/1MB)")

    # mem subcommand
    mem_parser = subparsers.add_parser("mem", help="Memory profiling with DHAT")
    mem_parser.add_argument("--size", "-s", choices=["small", "medium", "large"], default="medium",
                           help="Test file size (default: medium/1MB)")
    mem_parser.add_argument("--no-analyze", action="store_true",
                           help="Don't analyze after profiling")

    # status subcommand
    subparsers.add_parser("status", help="Show tooling status")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        return 0

    if args.command == "bench":
        return cmd_bench(args)
    elif args.command == "cpu":
        return cmd_cpu(args)
    elif args.command == "mem":
        return cmd_mem(args)
    elif args.command == "status":
        return cmd_status(args)

    return 0


if __name__ == "__main__":
    sys.exit(main() or 0)
