"""Benchmark harness — warm-up, measurement, and statistics."""

from __future__ import annotations

import json
import os
import platform
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

import numpy as np

from benchmarks.config import MEASURED_ITERATIONS, WARMUP_ITERATIONS, ScalePreset


@dataclass
class BenchmarkResult:
    """Result of a single workload run against a single database."""

    db: str
    workload_id: str
    scale: str
    times_ms: list[float] = field(default_factory=list)
    result_value: Any = None
    error: str | None = None
    skipped: bool = False
    skip_reason: str | None = None

    @property
    def p50(self) -> float:
        if not self.times_ms:
            return 0.0
        return float(np.percentile(self.times_ms, 50))

    @property
    def p95(self) -> float:
        if not self.times_ms:
            return 0.0
        return float(np.percentile(self.times_ms, 95))

    @property
    def p99(self) -> float:
        if not self.times_ms:
            return 0.0
        return float(np.percentile(self.times_ms, 99))

    @property
    def mean(self) -> float:
        if not self.times_ms:
            return 0.0
        return float(np.mean(self.times_ms))

    @property
    def stddev(self) -> float:
        if not self.times_ms:
            return 0.0
        return float(np.std(self.times_ms))

    def to_dict(self) -> dict[str, Any]:
        return {
            "db": self.db,
            "workload_id": self.workload_id,
            "scale": self.scale,
            "times_ms": self.times_ms,
            "p50": self.p50,
            "p95": self.p95,
            "p99": self.p99,
            "mean": self.mean,
            "stddev": self.stddev,
            "result_value": _safe_serialize(self.result_value),
            "error": self.error,
            "skipped": self.skipped,
            "skip_reason": self.skip_reason,
        }


@dataclass
class BenchmarkRun:
    """Complete benchmark run with system info and all results."""

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

    def save(self, results_dir: str | Path) -> Path:
        """Save results to JSON."""
        results_dir = Path(results_dir)
        results_dir.mkdir(parents=True, exist_ok=True)
        filename = f"bench_{self.scale}_{self.timestamp}.json"
        path = results_dir / filename
        data = {
            "timestamp": self.timestamp,
            "system_info": self.system_info,
            "scale": self.scale,
            "results": [r.to_dict() for r in self.results],
        }
        path.write_text(json.dumps(data, indent=2, default=str))
        return path


def collect_system_info() -> dict[str, Any]:
    """Gather system information for result context."""
    info: dict[str, Any] = {
        "platform": platform.platform(),
        "processor": platform.processor(),
        "cpu_count": os.cpu_count(),
        "python_version": platform.python_version(),
    }

    # Try to get memory info on Linux
    try:
        with open("/proc/meminfo") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    kb = int(line.split()[1])
                    info["memory_gb"] = round(kb / 1_048_576, 1)
                    break
    except OSError:
        pass

    return info


def run_workload(
    fn: Callable[[], Any],
    *,
    warmup: int = WARMUP_ITERATIONS,
    iterations: int = MEASURED_ITERATIONS,
) -> tuple[list[float], Any]:
    """Run a workload function with warm-up and measurement.

    Returns (times_ms, last_result).
    """
    last_result = None

    # Warm-up
    for _ in range(warmup):
        fn()

    # Measured runs
    times: list[float] = []
    for _ in range(iterations):
        start = time.perf_counter()
        last_result = fn()
        elapsed = (time.perf_counter() - start) * 1000.0
        times.append(elapsed)

    return times, last_result


def execute_benchmark(
    adapter: Any,
    workload_fns: dict[str, Callable[[], Any]],
    scale: ScalePreset,
    *,
    warmup: int = WARMUP_ITERATIONS,
    iterations: int = MEASURED_ITERATIONS,
) -> BenchmarkRun:
    """Run all workloads against a single adapter.

    workload_fns maps workload_id → callable.
    """
    run = BenchmarkRun(
        timestamp=datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S"),
        system_info=collect_system_info(),
        scale=scale.name,
    )

    for wid, fn in workload_fns.items():
        print(f"  [{adapter.name}] {wid} ...", end=" ", flush=True)
        result = BenchmarkResult(
            db=adapter.name,
            workload_id=wid,
            scale=scale.name,
        )

        if fn is None:
            result.skipped = True
            result.skip_reason = f"{wid} not supported by {adapter.name}"
            print("SKIPPED")
        else:
            try:
                times, val = run_workload(fn, warmup=warmup, iterations=iterations)
                result.times_ms = times
                result.result_value = val
                print(f"{result.p50:.1f}ms (p50)")
            except Exception as exc:
                result.error = f"{type(exc).__name__}: {exc}"
                print(f"ERROR: {exc}")
                # Reset transaction state so subsequent workloads aren't blocked
                if hasattr(adapter, "db") and hasattr(adapter.db, "connection"):
                    try:
                        adapter.db.connection.rollback()
                    except Exception:
                        pass

        run.results.append(result)

    return run


def _safe_serialize(val: Any) -> Any:
    """Make a value JSON-serializable."""
    if val is None:
        return None
    if isinstance(val, (int, float, str, bool)):
        return val
    if isinstance(val, (list, tuple)):
        return [_safe_serialize(v) for v in val]
    if isinstance(val, dict):
        return {str(k): _safe_serialize(v) for k, v in val.items()}
    if isinstance(val, np.ndarray):
        return val.tolist()
    return str(val)
