#!/usr/bin/env python3
"""Shared helpers for Fold benchmark runners."""

from __future__ import annotations

import json
import os
import platform
import re
import shutil
import subprocess
import sys
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any


class MeasuredProcess:
    def __init__(
        self,
        returncode: int,
        output: str,
        wall_seconds: float,
        cpu_seconds: float | None = None,
        rss_mb: float | None = None,
        timed_out: bool = False,
    ) -> None:
        self.returncode = returncode
        self.output = output
        self.wall_seconds = wall_seconds
        self.cpu_seconds = cpu_seconds
        self.rss_mb = rss_mb
        self.timed_out = timed_out

    def metrics(self) -> dict[str, float | bool]:
        out: dict[str, float | bool] = {"wall_seconds": round(self.wall_seconds, 3)}
        if self.cpu_seconds is not None:
            out["cpu_seconds"] = round(self.cpu_seconds, 3)
        if self.rss_mb is not None:
            out["rss_mb"] = round(self.rss_mb, 3)
        if self.timed_out:
            out["timed_out"] = True
        return out


def load_json(path: Path) -> dict[str, Any]:
    try:
        with path.open("r", encoding="utf-8") as f:
            data = json.load(f)
    except (OSError, json.JSONDecodeError) as exc:
        raise ValueError(f"invalid JSON file {path}: {exc}") from exc
    if not isinstance(data, dict):
        raise ValueError(f"JSON file must contain an object: {path}")
    return data


def resolve_path(value: str, base: Path) -> Path:
    path = Path(value)
    if path.is_absolute():
        return path
    return (base / path).resolve()


def specs_from_args(paths: list[Path], default_dirs: list[Path]) -> list[Path]:
    if paths:
        return paths
    specs: list[Path] = []
    for spec_dir in default_dirs:
        if spec_dir.exists():
            specs.extend(sorted(spec_dir.glob("*.json")))
    return specs


def find_fold_binary(
    repo_root: Path, build: Path, override: Path | None
) -> Path | None:
    fold_bin = override or (build / "src" / "cli" / "fold")
    if not fold_bin.is_absolute():
        fold_bin = repo_root / fold_bin
    return fold_bin.resolve() if fold_bin.exists() else None


def expectation(spec: dict[str, Any]) -> dict[str, Any]:
    expect = spec.get("expect", {})
    return expect if isinstance(expect, dict) else {}


def budget_timeout_seconds(spec: dict[str, Any]) -> float | None:
    budgets = expectation(spec).get("budgets", {})
    if not isinstance(budgets, dict):
        return None
    limit = budgets.get("max_wall_seconds")
    if limit is None:
        return None
    try:
        return max(float(limit) * 2.0, float(limit) + 5.0)
    except (TypeError, ValueError):
        return None


def _parse_float(value: str) -> float | None:
    try:
        return float(value)
    except ValueError:
        return None


def _parse_macos_time(output: str) -> tuple[float | None, float | None]:
    user = None
    sys_time = None
    rss_mb = None
    for line in output.splitlines():
        match = re.search(
            r"([0-9.]+)\s+real\s+([0-9.]+)\s+user\s+([0-9.]+)\s+sys", line
        )
        if match:
            user = _parse_float(match.group(2))
            sys_time = _parse_float(match.group(3))
            continue
        match = re.search(r"([0-9]+)\s+maximum resident set size", line)
        if match:
            rss_kb = _parse_float(match.group(1))
            rss_mb = rss_kb / 1024 if rss_kb is not None else None
    cpu = None if user is None or sys_time is None else user + sys_time
    return cpu, rss_mb


def _parse_posix_time(output: str) -> tuple[float | None, float | None]:
    """Parse portable ``time -p`` output.

    macOS ``time -l`` calls ``sysctl kern.clockrate`` after the child exits.
    Sandboxed CI can deny that query and make ``time`` return 1 even when the
    measured command returned 0. Portable mode preserves the child status and
    CPU measurement; RSS remains unavailable rather than being fabricated.
    """

    user = None
    sys_time = None
    for line in output.splitlines():
        fields = line.split()
        if len(fields) != 2:
            continue
        if fields[0] == "user":
            user = _parse_float(fields[1])
        elif fields[0] == "sys":
            sys_time = _parse_float(fields[1])
    cpu = None if user is None or sys_time is None else user + sys_time
    return cpu, None


def _parse_gnu_time(output: str) -> tuple[float | None, float | None]:
    user = None
    sys_time = None
    rss_mb = None
    for line in output.splitlines():
        if line.startswith("User time (seconds):"):
            user = _parse_float(line.split(":", 1)[1].strip())
        elif line.startswith("System time (seconds):"):
            sys_time = _parse_float(line.split(":", 1)[1].strip())
        elif line.startswith("Maximum resident set size (kbytes):"):
            rss_kb = _parse_float(line.split(":", 1)[1].strip())
            rss_mb = rss_kb / 1024 if rss_kb is not None else None
    cpu = None if user is None or sys_time is None else user + sys_time
    return cpu, rss_mb


def run_measured(
    cmd: list[str],
    cwd: Path,
    env: dict[str, str] | None = None,
    timeout_seconds: float | None = None,
) -> MeasuredProcess:
    time_bin = shutil.which("time")
    system = platform.system()
    timed_cmd = cmd
    parser: Callable[[str], tuple[float | None, float | None]] | None = None
    if time_bin and system == "Darwin":
        timed_cmd = [time_bin, "-p", *cmd]
        parser = _parse_posix_time
    elif time_bin and system == "Linux":
        timed_cmd = [time_bin, "-v", *cmd]
        parser = _parse_gnu_time

    started = time.monotonic()
    executable = Path(cmd[0])
    if executable.is_absolute() and not executable.exists():
        return MeasuredProcess(
            127, f"failed to launch {cmd[0]}: executable does not exist", 0.0
        )
    try:
        proc = subprocess.run(
            timed_cmd,
            cwd=cwd,
            env=env,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            timeout=timeout_seconds,
        )
    except subprocess.TimeoutExpired as exc:
        wall = time.monotonic() - started
        output = exc.output if isinstance(exc.output, str) else ""
        if not output:
            output = f"timed out after {timeout_seconds} seconds"
        return MeasuredProcess(124, output, wall, timed_out=True)
    except OSError as exc:
        wall = time.monotonic() - started
        return MeasuredProcess(127, f"failed to launch {cmd[0]}: {exc}", wall)

    wall = time.monotonic() - started
    cpu = None
    rss = None
    if parser is not None:
        cpu, rss = parser(proc.stdout)
    return MeasuredProcess(proc.returncode, proc.stdout, wall, cpu, rss)


def budget_failures(
    expect: dict[str, Any], metrics: dict[str, float | bool]
) -> list[str]:
    budgets = expect.get("budgets", {}) if isinstance(expect, dict) else {}
    if not isinstance(budgets, dict):
        return []
    failures: list[str] = []
    checks = [
        ("max_wall_seconds", "wall_seconds"),
        ("max_cpu_seconds", "cpu_seconds"),
        ("max_rss_mb", "rss_mb"),
    ]
    for budget_key, metric_key in checks:
        limit = budgets.get(budget_key)
        value = metrics.get(metric_key)
        if limit is None or value is None or isinstance(value, bool):
            continue
        measured = _parse_float(str(value))
        maximum = _parse_float(str(limit))
        if measured is None or maximum is None:
            failures.append(f"invalid {metric_key}/{budget_key} budget")
        elif measured > maximum:
            failures.append(f"{metric_key} {value} exceeds {budget_key} {limit}")
    return failures


def merge_budget_result(
    comparison: dict[str, Any], spec: dict[str, Any], metrics: dict[str, float | bool]
) -> bool:
    comparison.update(metrics)
    failures = comparison.setdefault("failures", [])
    if not isinstance(failures, list):
        failures = []
        comparison["failures"] = failures
    failures.extend(budget_failures(expectation(spec), metrics))
    return not failures


def compare_report(
    *,
    compare_script: Path,
    report: Path,
    spec_path: Path,
    repo_root: Path,
    spec: dict[str, Any],
    name: str,
    metrics: dict[str, float | bool],
) -> tuple[bool, dict[str, Any]]:
    cmp_proc = subprocess.run(
        [
            str(sys.executable),
            str(compare_script),
            "--report",
            str(report),
            "--spec",
            str(spec_path),
        ],
        cwd=repo_root,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    try:
        loaded = json.loads(cmp_proc.stdout)
        comparison: dict[str, Any] = (
            loaded if isinstance(loaded, dict) else {"raw": loaded}
        )
    except (TypeError, json.JSONDecodeError):
        comparison = {"raw_compare_output": cmp_proc.stdout}
    comparison["name"] = name
    comparison["seconds"] = metrics["wall_seconds"]
    passed = cmp_proc.returncode == 0 and merge_budget_result(comparison, spec, metrics)
    comparison["status"] = "ok" if passed else "failed"
    return passed, comparison


def scan_and_compare(
    *,
    cmd: list[str],
    report: Path,
    spec: dict[str, Any],
    spec_path: Path,
    name: str,
    repo_root: Path,
    compare_script: Path,
    env: dict[str, str] | None = None,
) -> tuple[bool, dict[str, Any]]:
    proc = run_measured(
        cmd, repo_root, env=env, timeout_seconds=budget_timeout_seconds(spec)
    )
    metrics = proc.metrics()
    if proc.returncode not in (0, 1):
        return False, {
            "name": name,
            "status": "failed",
            "reason": f"fold exited {proc.returncode}",
            "output": proc.output[-4000:],
            "seconds": metrics["wall_seconds"],
            **metrics,
        }
    if not report.exists():
        return False, {"name": name, "status": "failed", "reason": "report not written"}
    return compare_report(
        compare_script=compare_script,
        report=report,
        spec_path=spec_path,
        repo_root=repo_root,
        spec=spec,
        name=name,
        metrics=metrics,
    )


def write_compile_db(
    spec: dict[str, Any],
    spec_path: Path,
    output_dir: Path,
    repo_root: Path,
    *,
    source_key: str = "sources",
) -> Path:
    base = spec_path.parent
    sources = spec.get(source_key, [])
    if not isinstance(sources, list) or not sources:
        raise ValueError(f"spec must contain non-empty {source_key} list: {spec_path}")
    standard = str(spec.get("standard", "c++20"))
    extra_args = spec.get("extra_compile_args", [])
    if not isinstance(extra_args, list):
        extra_args = []

    db = []
    for entry in sources:
        source = resolve_path(str(entry), base)
        if not source.exists():
            raise FileNotFoundError(source)
        args = ["clang++", f"-std={standard}", *map(str, extra_args), "-c", str(source)]
        db.append({"directory": str(repo_root), "arguments": args, "file": str(source)})

    compile_db = output_dir / f"{spec_path.stem}_compile_commands.json"
    compile_db.write_text(json.dumps(db, indent=2), encoding="utf-8")
    return compile_db


def build_c_fixture(
    spec: dict[str, Any],
    spec_path: Path,
    output_dir: Path,
    repo_root: Path,
    *,
    source_key: str,
    compiler_key: str = "compiler",
    cflags_key: str = "cflags",
) -> Path:
    source = resolve_path(str(spec[source_key]), spec_path.parent)
    if not source.exists():
        raise FileNotFoundError(source)
    compiler = str(spec.get(compiler_key, "cc"))
    compiler_path = shutil.which(compiler)
    if compiler_path is None:
        raise FileNotFoundError(f"compiler not found: {compiler}")
    cflags = spec.get(cflags_key, [])
    if not isinstance(cflags, list):
        cflags = []
    built_dir = output_dir / "built"
    built_dir.mkdir(parents=True, exist_ok=True)
    binary = built_dir / spec_path.stem
    try:
        proc = subprocess.run(
            [compiler_path, *map(str, cflags), str(source), "-o", str(binary)],
            cwd=repo_root,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
        )
    except OSError as exc:
        raise RuntimeError(
            f"fixture build could not start for {source}: {exc}"
        ) from exc
    if proc.returncode != 0:
        raise RuntimeError(f"fixture build failed for {source}: {proc.stdout[-2000:]}")
    return binary


def print_missing_fold(fold_bin: Path) -> None:
    print(
        json.dumps(
            {"status": "failed", "reason": f"fold binary not found: {fold_bin}"},
            indent=2,
        )
    )


def run_suite(
    *,
    suite_name: str,
    specs: list[Path],
    run_one: Callable[[Path], tuple[bool, dict[str, Any]]],
) -> int:
    results = []
    ok = True
    for spec_path in specs:
        passed, result = run_one(spec_path.resolve())
        ok = ok and passed
        results.append(result)
    summary = {
        "suite": suite_name,
        "status": "ok" if ok else "failed",
        "count": len(results),
        "results": results,
    }
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0 if ok else 1


def prepend_env_path(env: dict[str, str], key: str, value: Path) -> None:
    current = env.get(key)
    env[key] = str(value) if not current else str(value) + os.pathsep + current