#!/usr/bin/env python3
"""Run the minimal CollectiveSpec P0 collective-phase premise experiment.

P0 is deliberately *not* a performance experiment.  It holds the EAGLE
drafter at Kmax=3, injects a pre-fixed per-request verification horizon after
candidate generation, and records the scheduler and worker-side phase traces.
The two policy cells are a constant-K=3 control and a deterministic
request-static heterogeneous policy.
"""

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("--source-trace", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--run-id", required=True)
    parser.add_argument("--request-count", type=int, default=64)
    parser.add_argument("--completion-tokens", type=int, default=64)
    parser.add_argument("--port", type=int, required=True)
    parser.add_argument("--seed", type=int, default=20260713)
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def read_rows(source_trace: Path, count: int) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with source_trace.open("r", encoding="utf-8") as handle:
        for raw in handle:
            if not raw.strip():
                continue
            row = json.loads(raw)
            if not isinstance(row, dict):
                continue
            rows.append(row)
            if len(rows) == count:
                break
    if len(rows) != count:
        raise SystemExit(f"source trace contains {len(rows)} usable rows, need {count}")
    return rows


def write_workload(args: argparse.Namespace) -> tuple[Path, Path]:
    rows = read_rows(args.source_trace, args.request_count)
    trace_path = args.output_root / "p0_trace.jsonl"
    with trace_path.open("w", encoding="utf-8") as handle:
        for index, row in enumerate(rows):
            row = dict(row)
            # A single burst makes both DP schedulers exercise heterogeneous
            # verification plans concurrently.  P0 does not measure QPS.
            row["request_id"] = f"p0-{index:04d}"
            row["sampling_u"] = 0.0
            row["temperature"] = 0.0
            row["timestamp"] = 0.0
            handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
    windows_path = args.output_root / "p0_windows.json"
    windows = {
        "kind": "collectivespec_p0_fixed_burst",
        "source_trace": str(args.source_trace),
        "source_trace_sha256": sha256(args.source_trace),
        "windows": [
            {
                "window_id": "collectivespec_p0_burst",
                "trace_file": str(trace_path),
                "trace_type": "decode_only",
                "window_start": 0.0,
                "window_end": 1.0,
                "num_requests": args.request_count,
                "sampling_strategy": "all_requests_fixed_burst",
            }
        ],
    }
    windows_path.write_text(
        json.dumps(windows, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return trace_path, windows_path


def derived_spec(
    base: dict[str, Any],
    *,
    args: argparse.Namespace,
    windows_path: Path,
    policy_name: str,
) -> dict[str, Any]:
    spec = json.loads(json.dumps(base))
    spec["study_id"] = f"collectivespec-p0-{policy_name}-{args.run_id}"
    spec["engine"]["cwd"] = str(args.source_root)
    spec["engine"]["port"] = args.port
    flags = spec["engine"]["base_flags"]
    flags["port"] = args.port
    flags["seed"] = args.seed
    flags["scheduler-cls"] = "collectivespec_p0.scheduler.P0Scheduler"
    flags["worker-cls"] = "collectivespec_p0.worker.P0Worker"
    raw_speculative = flags.get("speculative-config")
    if not isinstance(raw_speculative, str):
        raise SystemExit("base spec must have speculative-config for P0")
    speculative = json.loads(raw_speculative)
    speculative["num_speculative_tokens"] = 3
    flags["speculative-config"] = json.dumps(speculative, separators=(",", ":"))

    trace = spec["trace"]
    trace.update(
        {
            "completion_tokens_override": args.completion_tokens,
            "early_stop_max_elapsed_s": 1800.0,
            "early_stop_max_lag_s": None,
            "max_concurrency": args.request_count,
            "max_requests_per_probe": None,
            "replay_time_scale": 1.0,
            "request_mode": "decode_only",
            "timestamp_field": "timestamp",
            "u_field": "sampling_u",
            "window_id": "collectivespec_p0_burst",
            "windows_path": str(windows_path),
        }
    )
    spec["slo"] = {"target_pass_rate": 1.0}
    spec["search"] = {
        "low": 0.0,
        "high": 1.0,
        "tolerance": 1e-6,
        "max_probes": 1,
        "sample_seed": args.seed,
    }
    return spec


def policy_payload(policy_name: str, seed: int) -> dict[str, Any]:
    if policy_name == "control_k3":
        return {
            "version": "collectivespec-p0-v1-control-k3",
            "seed": seed,
            "max_k": 3,
            "policy": "constant-k",
            "constant_k": 3,
        }
    if policy_name == "heterogeneous":
        return {
            "version": "collectivespec-p0-v1-heterogeneous",
            "seed": seed,
            "max_k": 3,
            "policy": "blake2b-request-static-k",
        }
    raise ValueError(f"unknown P0 policy {policy_name}")


def main() -> int:
    args = parse_args()
    if args.request_count < 8:
        raise SystemExit("request-count must be at least 8 to exercise all K values")
    if args.completion_tokens < 8:
        raise SystemExit("completion-tokens must be at least 8")
    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 not args.source_trace.is_file():
        raise SystemExit(f"source trace does not exist: {args.source_trace}")
    overlay = args.source_root / "scripts" / "collectivespec" / "p0_overlay"
    if not overlay.is_dir():
        raise SystemExit(f"P0 overlay does not exist: {overlay}")

    args.output_root.mkdir(parents=True, exist_ok=True)
    trace_path, windows_path = write_workload(args)
    base = json.loads(args.base_spec.read_text(encoding="utf-8"))
    manifest = {
        "kind": "collectivespec_p0_phase_trace",
        "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),
        "source_trace": str(args.source_trace),
        "source_trace_sha256": sha256(args.source_trace),
        "materialized_trace": str(trace_path),
        "materialized_trace_sha256": sha256(trace_path),
        "windows_path": str(windows_path),
        "overlay": str(overlay),
        "run_id": args.run_id,
        "topology_claim": "inherited_from_base_spec",
        "gpu_count": base.get("hardware", {}).get("gpu_count"),
        "request_count": args.request_count,
        "completion_tokens": args.completion_tokens,
        "policies": ["control_k3", "heterogeneous"],
        "seed": args.seed,
        "port": args.port,
        "python": sys.executable,
    }
    (args.output_root / "manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )

    results: list[dict[str, Any]] = []
    for ordinal, policy_name in enumerate(manifest["policies"], start=1):
        cell = args.output_root / policy_name
        cell.mkdir(exist_ok=True)
        policy_path = cell / "policy.json"
        policy_path.write_text(
            json.dumps(policy_payload(policy_name, args.seed), indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        spec_path = cell / "study_spec.json"
        spec_path.write_text(
            json.dumps(
                derived_spec(
                    base,
                    args=args,
                    windows_path=windows_path,
                    policy_name=policy_name,
                ),
                indent=2,
                sort_keys=True,
            )
            + "\n",
            encoding="utf-8",
        )
        store_root = cell / "store"
        logs = cell / "p0_logs"
        command = [
            sys.executable,
            "-m",
            "aituner.cli",
            "study",
            "tune",
            "--spec",
            str(spec_path),
            "--store-root",
            str(store_root),
            "--max-trials",
            "1",
        ]
        record = {
            "ordinal": ordinal,
            "policy": policy_name,
            "policy_path": str(policy_path),
            "spec_path": str(spec_path),
            "p0_logs": str(logs),
            "command": command,
        }
        print(json.dumps(record, sort_keys=True), flush=True)
        if args.dry_run:
            results.append({**record, "returncode": None})
            continue
        env = os.environ.copy()
        inherited = env.get("PYTHONPATH", "")
        env["PYTHONPATH"] = os.pathsep.join(
            part for part in (str(overlay), str(args.source_root / "src"), inherited) if part
        )
        env["PYTHONDONTWRITEBYTECODE"] = "1"
        env["COLLECTIVESPEC_P0_LOG_DIR"] = str(logs)
        env["COLLECTIVESPEC_P0_POLICY_PATH"] = str(policy_path)
        log_path = cell / "driver.log"
        with log_path.open("w", encoding="utf-8") as handle:
            completed = subprocess.run(
                command,
                cwd=args.source_root,
                env=env,
                text=True,
                stdout=handle,
                stderr=subprocess.STDOUT,
            )
        results.append({**record, "driver_log": str(log_path), "returncode": completed.returncode})
        if completed.returncode:
            break
    result = {
        "finished_at_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
        "results": results,
        "status": "ok" if results and all(item["returncode"] in (0, None) for item in results) else "failed",
    }
    (args.output_root / "driver_result.json").write_text(
        json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    return 0 if result["status"] == "ok" else 1


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