#!/usr/bin/env python3
"""Run an isolated, reproducible static-speculation screening study.

This driver deliberately does not implement CollectiveSpec.  It establishes
whether the current static-K deployment has enough headroom to justify such a
system.  Each K gets an immutable derived StudySpec and a separate store.
"""

from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import os
import subprocess
import sys
from pathlib import Path
from typing import Any


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-spec", type=Path, required=True)
    parser.add_argument("--source-root", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--run-id", required=True)
    parser.add_argument("--k-values", nargs="+", type=int, default=[0, 1, 2, 3])
    parser.add_argument(
        "--order",
        nargs="+",
        type=int,
        help="Sequential launch order. Defaults to --k-values order.",
    )
    parser.add_argument("--sampling-u-low", type=float, default=0.005)
    parser.add_argument("--sampling-u-high", type=float, default=0.020)
    parser.add_argument("--tolerance", type=float, default=0.003)
    parser.add_argument("--max-probes", type=int, default=3)
    parser.add_argument(
        "--completion-tokens-override",
        type=int,
        help="Force every request to this exact completion length for a controlled screen.",
    )
    parser.add_argument(
        "--trace-windows-path",
        type=Path,
        help="Override trace.windows_path in every derived StudySpec.",
    )
    parser.add_argument(
        "--trace-window-id",
        help="Override trace.window_id in every derived StudySpec.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        help="Explicit engine seed recorded in every derived StudySpec.",
    )
    parser.add_argument(
        "--port",
        type=int,
        help="Serving port. Defaults to the base StudySpec's engine/base flag port.",
    )
    parser.add_argument(
        "--continue-on-error",
        action="store_true",
        help="Record a failed K and proceed to the next independent K.",
    )
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def derived_spec(
    base: dict[str, Any], args: argparse.Namespace, k: int
) -> dict[str, Any]:
    spec = json.loads(json.dumps(base))
    label = f"collectivespec-static-k{k}-screen-{args.run_id}"
    spec["study_id"] = label
    spec["engine"]["cwd"] = str(args.source_root)
    flags = spec["engine"]["base_flags"]
    port = args.port
    if port is None:
        port = int(flags.get("port", spec["engine"]["port"]))
    spec["engine"]["port"] = port
    flags["port"] = port

    if k == 0:
        # vLLM's EAGLE config rejects num_speculative_tokens=0.  Removing the
        # entire flag is the real no-speculation baseline.
        flags.pop("speculative-config", None)
    else:
        raw_config = flags.get("speculative-config")
        if raw_config is None:
            raise ValueError("base spec has no speculative-config for K > 0")
        config = json.loads(raw_config)
        config["num_speculative_tokens"] = k
        flags["speculative-config"] = json.dumps(config, separators=(",", ":"))

    search = spec["search"]
    search["low"] = args.sampling_u_low
    search["high"] = args.sampling_u_high
    search["tolerance"] = args.tolerance
    search["max_probes"] = args.max_probes
    if args.completion_tokens_override is not None:
        spec["trace"]["completion_tokens_override"] = args.completion_tokens_override
    if args.trace_windows_path is not None:
        spec["trace"]["windows_path"] = str(args.trace_windows_path)
    if args.trace_window_id is not None:
        spec["trace"]["window_id"] = args.trace_window_id
    if args.seed is not None:
        flags["seed"] = args.seed
    return spec


def main() -> int:
    args = parse_args()
    if not args.base_spec.is_file():
        raise SystemExit(f"base spec does not exist: {args.base_spec}")
    if not args.source_root.is_dir():
        raise SystemExit(f"source root does not exist: {args.source_root}")
    if any(k < 0 for k in args.k_values):
        raise SystemExit("K must be non-negative")
    if not 0 <= args.sampling_u_low < args.sampling_u_high <= 1:
        raise SystemExit("sampling-u bounds must satisfy 0 <= low < high <= 1")
    if args.max_probes < 1:
        raise SystemExit("max-probes must be positive")
    if args.completion_tokens_override is not None and args.completion_tokens_override < 0:
        raise SystemExit("completion token override must be non-negative")
    if (args.trace_windows_path is None) != (args.trace_window_id is None):
        raise SystemExit(
            "--trace-windows-path and --trace-window-id must be provided together"
        )
    if args.trace_windows_path is not None and not args.trace_windows_path.is_file():
        raise SystemExit(f"trace windows file does not exist: {args.trace_windows_path}")

    order = args.order or args.k_values
    if sorted(order) != sorted(args.k_values) or len(order) != len(args.k_values):
        raise SystemExit("--order must be a permutation of --k-values")

    base = json.loads(args.base_spec.read_text())
    args.output_root.mkdir(parents=True, exist_ok=True)
    specs_dir = args.output_root / "specs"
    specs_dir.mkdir(exist_ok=True)

    manifest = {
        "kind": "collectivespec_static_k_screen",
        "created_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
        "base_spec": str(args.base_spec),
        "base_spec_sha256": sha256(args.base_spec),
        "source_root": str(args.source_root),
        "run_id": args.run_id,
        "k_values": args.k_values,
        "launch_order": order,
        "sampling_u": [args.sampling_u_low, args.sampling_u_high],
        "tolerance": args.tolerance,
        "max_probes": args.max_probes,
        "completion_tokens_override": args.completion_tokens_override,
        "trace_windows_path": (
            str(args.trace_windows_path) if args.trace_windows_path is not None else None
        ),
        "trace_window_id": args.trace_window_id,
        "seed": args.seed,
        "port": args.port
        if args.port is not None
        else int(base["engine"]["base_flags"].get("port", base["engine"]["port"])),
        "python": sys.executable,
    }
    (args.output_root / "manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n"
    )

    failures: list[dict[str, Any]] = []
    for ordinal, k in enumerate(order, start=1):
        spec = derived_spec(base, args, k)
        spec_path = specs_dir / f"{ordinal:02d}_k{k}.json"
        spec_path.write_text(json.dumps(spec, indent=2, sort_keys=True) + "\n")
        store_root = args.output_root / "stores" / f"k{k}"
        command = [
            sys.executable,
            "-m",
            "aituner.cli",
            "study",
            "tune",
            "--spec",
            str(spec_path),
            "--store-root",
            str(store_root),
            "--max-trials",
            "1",
        ]
        print(f"[{ordinal}/{len(order)}] K={k} command={json.dumps(command)}", flush=True)
        if args.dry_run:
            continue

        log_path = args.output_root / "logs" / f"{ordinal:02d}_k{k}.log"
        log_path.parent.mkdir(exist_ok=True)
        env = os.environ.copy()
        env["PYTHONPATH"] = str(args.source_root / "src")
        env["PYTHONDONTWRITEBYTECODE"] = "1"
        with log_path.open("w") as log:
            completed = subprocess.run(
                command,
                cwd=args.source_root,
                env=env,
                text=True,
                stdout=log,
                stderr=subprocess.STDOUT,
            )
        print(
            f"[{ordinal}/{len(order)}] K={k} returncode={completed.returncode} log={log_path}",
            flush=True,
        )
        if completed.returncode:
            failures.append({"k": k, "returncode": completed.returncode, "log": str(log_path)})
            if not args.continue_on_error:
                break

    result = {
        "finished_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
        "failures": failures,
        "status": "ok" if not failures else "completed_with_failures",
    }
    (args.output_root / "driver_result.json").write_text(
        json.dumps(result, indent=2, sort_keys=True) + "\n"
    )
    return 0 if not failures else 1


if __name__ == "__main__":
    raise SystemExit(main())
