#!/usr/bin/env python3
"""
Plot benchmark progressions from data/bench/results.

Reads benchmark JSON files and creates visualizations of performance
over time, or exports data as CSV.

Usage:
    ./scripts/plot_benchmarks.py                    # Plot all benchmarks
    ./scripts/plot_benchmarks.py --csv              # Export as CSV
    ./scripts/plot_benchmarks.py --csv -o out.csv   # Export to specific file
    ./scripts/plot_benchmarks.py --sizes 1mb 100mb  # Plot specific sizes only
    ./scripts/plot_benchmarks.py --metric duration  # Plot duration instead of rate
"""

import argparse
import csv
import json
import sys
from dataclasses import dataclass
from datetime import datetime
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."""
    import os

    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 matplotlib", file=sys.stderr)
        sys.exit(1)

    current_python = Path(sys.executable).resolve()
    if current_python != VENV_PYTHON.resolve():
        os.execv(str(VENV_PYTHON), [str(VENV_PYTHON)] + sys.argv)


ensure_venv()


@dataclass
class BenchmarkPoint:
    """A single benchmark data point."""

    timestamp: datetime
    git_commit: str
    binary_md5: str
    file_size: str
    file_lines: int
    command: str
    avg_duration: float
    min_duration: float
    max_duration: float
    avg_lines_per_sec: float


def load_benchmark_files(results_dir: Path) -> list[BenchmarkPoint]:
    """Load all benchmark JSON files and return data points."""
    points = []

    if not results_dir.exists():
        print(f"Error: Results directory not found: {results_dir}", file=sys.stderr)
        return points

    json_files = sorted(results_dir.glob("*.json"))

    for json_file in json_files:
        try:
            with open(json_file) as f:
                data = json.load(f)

            timestamp = datetime.fromisoformat(data["timestamp"])
            git_commit = data.get("git_commit", "unknown")[:8]
            binary_md5 = data.get("binary_md5", "unknown")[:8]

            for result in data.get("results", []):
                # Skip failed results
                if "error" in result or result.get("successful_runs", 0) == 0:
                    continue

                # Handle both old (throughput) and new (lines_per_sec) formats
                avg_lines_per_sec = result.get("avg_lines_per_sec")
                if avg_lines_per_sec is None:
                    # Old format: calculate from file_lines and duration
                    file_lines = result.get("file_lines", 0)
                    avg_duration = result.get("avg_duration_seconds", 0)
                    if avg_duration > 0 and file_lines > 0:
                        avg_lines_per_sec = file_lines / avg_duration
                    else:
                        avg_lines_per_sec = 0

                points.append(
                    BenchmarkPoint(
                        timestamp=timestamp,
                        git_commit=git_commit,
                        binary_md5=binary_md5,
                        file_size=result["file"],
                        file_lines=result.get("file_lines", 0),
                        command=result["command"],
                        avg_duration=result.get("avg_duration_seconds", 0),
                        min_duration=result.get("min_duration_seconds", 0),
                        max_duration=result.get("max_duration_seconds", 0),
                        avg_lines_per_sec=avg_lines_per_sec,
                    )
                )

        except (json.JSONDecodeError, KeyError, ValueError) as e:
            print(f"Warning: Failed to parse {json_file}: {e}", file=sys.stderr)

    return points


def export_csv(points: list[BenchmarkPoint], output_path: Path | None) -> None:
    """Export benchmark data as CSV."""
    if output_path:
        f = open(output_path, "w", newline="")
    else:
        f = sys.stdout

    try:
        writer = csv.writer(f)
        writer.writerow(
            [
                "timestamp",
                "git_commit",
                "binary_md5",
                "file_size",
                "file_lines",
                "command",
                "avg_duration_seconds",
                "min_duration_seconds",
                "max_duration_seconds",
                "avg_lines_per_sec",
            ]
        )

        for p in sorted(points, key=lambda x: (x.timestamp, x.file_size)):
            writer.writerow(
                [
                    p.timestamp.isoformat(),
                    p.git_commit,
                    p.binary_md5,
                    p.file_size,
                    p.file_lines,
                    p.command,
                    p.avg_duration,
                    p.min_duration,
                    p.max_duration,
                    p.avg_lines_per_sec,
                ]
            )
    finally:
        if output_path:
            f.close()
            print(f"CSV exported to: {output_path}", file=sys.stderr)


def plot_benchmarks(
    points: list[BenchmarkPoint],
    sizes: list[str] | None,
    metric: str,
    output_path: Path | None,
) -> None:
    """Plot benchmark progressions using matplotlib."""
    try:
        import matplotlib.pyplot as plt
        import matplotlib.dates as mdates
    except ImportError:
        print(
            "Error: matplotlib not installed. Install with: uv pip install matplotlib",
            file=sys.stderr,
        )
        sys.exit(1)

    # Group points by file size
    by_size: dict[str, list[BenchmarkPoint]] = {}
    for p in points:
        if sizes and p.file_size not in sizes:
            continue
        if p.file_size not in by_size:
            by_size[p.file_size] = []
        by_size[p.file_size].append(p)

    if not by_size:
        print("No data to plot.", file=sys.stderr)
        return

    # Sort each group by timestamp
    for size in by_size:
        by_size[size].sort(key=lambda x: x.timestamp)

    # Define size order for consistent colors
    size_order = ["1kb", "10kb", "500kb", "1mb", "100mb"]
    ordered_sizes = [s for s in size_order if s in by_size]
    ordered_sizes += [s for s in by_size if s not in size_order]

    # Create figure
    fig, ax = plt.subplots(figsize=(12, 6))

    # Color map
    colors = plt.cm.viridis([i / len(ordered_sizes) for i in range(len(ordered_sizes))])

    for i, size in enumerate(ordered_sizes):
        pts = by_size[size]
        timestamps = [p.timestamp for p in pts]

        if metric == "duration":
            values = [p.avg_duration for p in pts]
            ylabel = "Duration (seconds)"
            title = "Benchmark Duration Over Time"
        else:  # rate
            values = [p.avg_lines_per_sec for p in pts]
            ylabel = "Lines per Second"
            title = "Benchmark Throughput Over Time"

        ax.plot(
            timestamps,
            values,
            marker="o",
            linestyle="-",
            label=size,
            color=colors[i],
            markersize=6,
        )

        # Add error bars for duration
        if metric == "duration":
            min_vals = [p.min_duration for p in pts]
            max_vals = [p.max_duration for p in pts]
            yerr_lower = [v - m for v, m in zip(values, min_vals)]
            yerr_upper = [m - v for v, m in zip(values, max_vals)]
            ax.errorbar(
                timestamps,
                values,
                yerr=[yerr_lower, yerr_upper],
                fmt="none",
                color=colors[i],
                alpha=0.5,
                capsize=3,
            )

    ax.set_xlabel("Date")
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.legend(title="File Size", loc="best")
    ax.grid(True, alpha=0.3)

    # Format x-axis dates
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m-%d\n%H:%M"))
    plt.xticks(rotation=0)

    # Use log scale for rate if values span multiple orders of magnitude
    if metric == "rate":
        all_values = [p.avg_lines_per_sec for pts in by_size.values() for p in pts]
        if max(all_values) / max(min(all_values), 1) > 100:
            ax.set_yscale("log")
            ax.set_ylabel("Lines per Second (log scale)")

    plt.tight_layout()

    if output_path:
        plt.savefig(output_path, dpi=150)
        print(f"Plot saved to: {output_path}")
    else:
        plt.show()


def main():
    parser = argparse.ArgumentParser(
        description="Plot benchmark progressions or export as CSV"
    )
    parser.add_argument(
        "-d",
        "--data-dir",
        type=Path,
        default=Path("data/bench/results"),
        help="Directory containing benchmark JSON files (default: data/bench/results)",
    )
    parser.add_argument(
        "--csv",
        action="store_true",
        help="Export data as CSV instead of plotting",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        help="Output file path (for CSV or plot image). If not specified, "
        "CSV goes to stdout and plot is shown interactively.",
    )
    parser.add_argument(
        "--sizes",
        nargs="+",
        help="Filter to specific file sizes (e.g., 1mb 100mb)",
    )
    parser.add_argument(
        "--metric",
        choices=["rate", "duration"],
        default="rate",
        help="Metric to plot: 'rate' (lines/sec) or 'duration' (default: rate)",
    )
    parser.add_argument(
        "--list",
        action="store_true",
        dest="list_files",
        help="List available benchmark files and exit",
    )

    args = parser.parse_args()

    # List mode
    if args.list_files:
        if not args.data_dir.exists():
            print(f"No results directory: {args.data_dir}")
            return

        json_files = sorted(args.data_dir.glob("*.json"))
        if not json_files:
            print("No benchmark files found.")
            return

        print(f"Found {len(json_files)} benchmark file(s) in {args.data_dir}:\n")
        for f in json_files:
            try:
                with open(f) as jf:
                    data = json.load(jf)
                ts = data.get("timestamp", "unknown")
                commit = data.get("git_commit", "unknown")[:8]
                md5 = data.get("binary_md5", "unknown")[:8]
                n_results = len(data.get("results", []))
                print(f"  {f.name}")
                print(f"    Timestamp: {ts}")
                print(f"    Commit: {commit}, MD5: {md5}")
                print(f"    Results: {n_results} benchmarks")
                print()
            except Exception as e:
                print(f"  {f.name} (error: {e})")
        return

    # Load data
    points = load_benchmark_files(args.data_dir)

    if not points:
        print("No benchmark data found.", file=sys.stderr)
        sys.exit(1)

    print(f"Loaded {len(points)} benchmark data points.", file=sys.stderr)

    # Filter by size if specified
    if args.sizes:
        original_count = len(points)
        points = [p for p in points if p.file_size in args.sizes]
        print(
            f"Filtered to {len(points)} points for sizes: {', '.join(args.sizes)}",
            file=sys.stderr,
        )

    if args.csv:
        export_csv(points, args.output)
    else:
        plot_benchmarks(points, args.sizes, args.metric, args.output)


if __name__ == "__main__":
    main()
