#!/usr/bin/env python3
"""Serialized dash0 controller for the exact-timestamp prefix pilot."""

from __future__ import annotations

import argparse
import json
import os
import shlex
import subprocess
import sys
import time
from pathlib import Path
from typing import Any


HERE = Path(__file__).resolve().parent
PHASE6 = HERE.parent / "opprof-phase6"
sys.path.insert(0, str(PHASE6))

import opprof_phase6_controller as base  # noqa: E402


ORDER = (
    "tp1_mns8",
    "tp1_mns64",
    "tp2_mns8",
    "tp2_mns64",
    "tp4_mns16",
    "tp4_mns64",
)
CELL_ESTIMATE_H20_HOURS = {1: 0.20, 2: 0.40, 4: 0.80}
SAFETY_H20_HOURS = 0.20


def atomic_json(path: Path, payload: Any) -> None:
    base.atomic_json(path, payload)


def wait_all_idle(timeout_s: float = 30.0) -> None:
    deadline = time.monotonic() + timeout_s
    last_error: Exception | None = None
    while time.monotonic() < deadline:
        try:
            base.assert_all_idle()
            return
        except RuntimeError as error:
            last_error = error
            time.sleep(1.0)
    raise last_error or RuntimeError("GPU idle timeout")


def configure_base(args: argparse.Namespace, manifest: dict[str, Any]) -> None:
    base.WORKDIR = args.run_root.parent
    base.RUN_ROOT = args.run_root
    base.STATE = args.run_root / "controller-state.json"
    base.SOURCE = args.vllm_source
    base.VENV = args.venv
    base.AITUNER = args.aituner_root
    base.MODEL = args.model
    base.CLIENT = args.client
    base.GPU_LIMIT = float(manifest["execution"]["hard_cap_h20_hours"])
    base.MARKER = "fidelity-prefix-pilot-20260714"
    base.CELLS = {
        cell: {"tp": int(config["tp"]), "mns": int(config["mns"])}
        for cell, config in manifest["cells"].items()
    }


def load_state(path: Path, hard_cap: float) -> dict[str, Any]:
    if path.exists():
        return json.loads(path.read_text(encoding="utf-8"))
    return {
        "schema": "fidelity-prefix-pilot-state-v1",
        "status": "initialized",
        "hard_cap_h20_hours": hard_cap,
        "gpu_hours_total": 0.0,
        "completed_cells": 0,
        "cells": {},
        "failures": [],
        "started_at": time.time(),
    }


def save_state(path: Path, state: dict[str, Any]) -> None:
    atomic_json(path, state)


def append_echo(run_root: Path, line: str) -> None:
    run_root.mkdir(parents=True, exist_ok=True)
    with (run_root / "launch-echo.log").open("a", encoding="utf-8") as target:
        target.write(line + "\n")
    print(line, flush=True)


def remaining_projection(manifest: dict[str, Any], index: int) -> float:
    return sum(
        CELL_ESTIMATE_H20_HOURS[int(manifest["cells"][cell]["tp"])]
        for cell in ORDER[index:]
    ) + SAFETY_H20_HOURS


def start_server(
    *,
    cell: str,
    index: int,
    run_root: Path,
) -> dict[str, Any]:
    config = base.CELLS[cell]
    gpus = tuple(range(int(config["tp"])))
    cell_root = run_root / "cells" / cell
    cell_root.mkdir(parents=True, exist_ok=True)
    port = 8900 + index
    command = base.server_command(cell, gpus, port)
    with (cell_root / "commands.log").open("a", encoding="utf-8") as log:
        log.write(f"SERVER {shlex.join(command)}\n")
    server_log = (cell_root / "server.log").open("ab", buffering=0)
    environment = os.environ.copy()
    environment.update(
        {
            "CUDA_VISIBLE_DEVICES": ",".join(map(str, gpus)),
            "VLLM_OPPROF_DIR": str(cell_root / "opprof"),
            "OPPROF_PHASE6_MARKER": base.MARKER,
            "AITUNER_ROOT": str(base.AITUNER),
            "HF_HUB_OFFLINE": "1",
            "TRANSFORMERS_OFFLINE": "1",
            "PYTHONUNBUFFERED": "1",
        }
    )
    server = subprocess.Popen(
        command,
        cwd=base.SOURCE,
        env=environment,
        stdout=server_log,
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    base.OWNED_PGIDS.add(server.pid)
    return {
        "cell": cell,
        "gpus": gpus,
        "port": port,
        "dir": cell_root,
        "server": server,
        "server_handle": server_log,
        "spawned_at": time.time(),
        "results": [],
    }


def selection_for(
    manifest: dict[str, Any], cell: str, role: str
) -> tuple[str, dict[str, Any]]:
    level = "low" if role == "burnin" or role.startswith("low") else "high"
    return level, manifest["cells"][cell]["targets"][level]["selections"][role]


def client_command(
    entry: dict[str, Any],
    *,
    role: str,
    selection: dict[str, Any],
    output: Path,
    warmup: bool,
) -> list[str]:
    config = base.CELLS[entry["cell"]]
    return [
        "taskset",
        "-c",
        base.cpu_mask(entry["gpus"]),
        str(base.VENV / "bin/python"),
        str(base.CLIENT),
        "warmup" if warmup else "run-anchor",
        "--study",
        str(selection["study"]),
        "--cell",
        entry["cell"],
        "--anchor",
        str(selection["anchor"]),
        "--tp",
        str(config["tp"]),
        "--mns",
        str(config["mns"]),
        "--base-url",
        f"http://127.0.0.1:{entry['port']}",
        "--result-dir",
        str(output),
    ]


def run_client(
    *,
    entry: dict[str, Any],
    role: str,
    selection: dict[str, Any],
    output: Path,
    state: dict[str, Any],
    warmup: bool = False,
) -> dict[str, Any]:
    command = client_command(
        entry, role=role, selection=selection, output=output, warmup=warmup
    )
    with (entry["dir"] / "commands.log").open("a", encoding="utf-8") as log:
        log.write(f"CLIENT role={role} {shlex.join(command)}\n")
    handle = (output.parent / f"{output.name}.log").open("ab", buffering=0)
    environment = os.environ.copy()
    environment.update({"AITUNER_ROOT": str(base.AITUNER), "PYTHONUNBUFFERED": "1"})
    process = subprocess.Popen(
        command,
        cwd=base.WORKDIR,
        env=environment,
        stdout=handle,
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    deadline = time.monotonic() + 180.0
    try:
        while process.poll() is None:
            if time.monotonic() > deadline:
                process.terminate()
                raise TimeoutError(f"client timeout: {entry['cell']} {role}")
            if entry["server"].poll() is not None:
                raise RuntimeError(f"server exited during {entry['cell']} {role}")
            base.assert_no_other_compute()
            if state["gpu_hours_total"] + base.live_gpu_hours([entry]) >= base.GPU_LIMIT:
                process.terminate()
                raise RuntimeError("pilot H20-hour hard cap reached")
            time.sleep(1.0)
    finally:
        handle.close()
    if process.returncode:
        raise RuntimeError(
            f"client failed: cell={entry['cell']} role={role} rc={process.returncode}"
        )
    result = json.loads((output / "result.json").read_text(encoding="utf-8"))
    validate_result_selection(
        result=result,
        selection=selection,
        cell=entry["cell"],
        role=role,
        warmup=warmup,
    )
    entry["results"].append(
        {"anchor": float(selection["anchor"]), "dir": str(output), "kind": result["kind"]}
    )
    return result


def validate_result_selection(
    *,
    result: dict[str, Any],
    selection: dict[str, Any],
    cell: str,
    role: str,
    warmup: bool,
) -> None:
    if warmup:
        if result["kind"] != "warmup" or int(result["selection"]["count"]) != 16:
            raise RuntimeError(f"invalid warmup selection: {cell} {role}")
        for key in ("warmup_16", "warmup_exact_16", "warmup_long"):
            if not result["invariants"].get(key, False):
                raise RuntimeError(f"warmup invariant {key} failed: {cell} {role}")
        return

    if int(result["selection"]["count"]) != int(selection["selected_count"]):
        raise RuntimeError(f"selection count mismatch: {cell} {role}")
    for key in (
        "request_id_order_sha256",
        "arrival_order_sha256",
        "raw_length_order_sha256",
    ):
        manifest_key = (
            "input_length_order_sha256" if key == "raw_length_order_sha256" else key
        )
        if result["selection"][key] != selection[manifest_key]:
            raise RuntimeError(f"selection hash mismatch {key}: {cell} {role}")


def execute_cell(
    *,
    index: int,
    cell: str,
    manifest: dict[str, Any],
    run_root: Path,
    state_path: Path,
    state: dict[str, Any],
) -> None:
    if state["cells"].get(cell, {}).get("status") == "complete":
        return
    projection = remaining_projection(manifest, index)
    if state["gpu_hours_total"] + projection > base.GPU_LIMIT:
        state["status"] = "budget_projection_stop"
        state["budget_stop"] = {
            "before_cell": cell,
            "spent_h20_hours": state["gpu_hours_total"],
            "remaining_projection_h20_hours": projection,
            "hard_cap_h20_hours": base.GPU_LIMIT,
        }
        save_state(state_path, state)
        raise RuntimeError(f"projected pilot cost exceeds hard cap before {cell}")

    config = manifest["cells"][cell]
    echo = (
        f"PILOT_CELL_ECHO cell={cell} tp={config['tp']} mns={config['mns']} "
        f"gpus=0-{int(config['tp']) - 1} workload={manifest['source']['window_id']} "
        f"roles=burnin+low1/high1/low2/high2/low3/high3 "
        f"spent_h20h={state['gpu_hours_total']:.6f} "
        f"remaining_projection_h20h={projection:.3f} cap_h20h={base.GPU_LIMIT:.1f} "
        f"manifest={run_root / 'pilot-manifest.json'}"
    )
    append_echo(run_root, echo)
    wait_all_idle()
    cell_state = {
        "status": "starting",
        "tp": int(config["tp"]),
        "mns": int(config["mns"]),
        "started_at": time.time(),
        "runs": [],
    }
    state["status"] = "running"
    state["cells"][cell] = cell_state
    save_state(state_path, state)
    entry = start_server(cell=cell, index=index, run_root=run_root)
    failure: Exception | None = None
    try:
        base.wait_ready(entry)
        _level, burnin = selection_for(manifest, cell, "burnin")
        cell_state["status"] = "warmup"
        save_state(state_path, state)
        warmup = run_client(
            entry=entry,
            role="burnin",
            selection=burnin,
            output=entry["dir"] / "warmup",
            state=state,
            warmup=True,
        )
        cell_state["warmup"] = {
            "exact_output_count": warmup["exact_output_count"],
            "long_gt4096": warmup["selection"]["long_gt4096"],
        }
        cell_state["status"] = "burnin"
        save_state(state_path, state)
        burnin_result = run_client(
            entry=entry,
            role="burnin",
            selection=burnin,
            output=entry["dir"] / "burnin",
            state=state,
        )
        cell_state["burnin"] = {
            "pass_rate": burnin_result["pass_rate"],
            "feasible": burnin_result["feasible"],
        }
        role_order = manifest["execution"][
            "even_cell_order" if index % 2 == 0 else "odd_cell_order"
        ]
        cell_state["status"] = "measured"
        cell_state["role_order"] = role_order
        save_state(state_path, state)
        for role in role_order:
            level, selection = selection_for(manifest, cell, role)
            result = run_client(
                entry=entry,
                role=role,
                selection=selection,
                output=entry["dir"] / f"{level}-rep{role[-1]}",
                state=state,
            )
            cell_state["runs"].append(
                {
                    "role": role,
                    "level": level,
                    "anchor": selection["anchor"],
                    "selected_count": selection["selected_count"],
                    "pass_rate": result["pass_rate"],
                    "feasible": result["feasible"],
                    "elapsed_s": result["interval"]["elapsed_s"],
                }
            )
            save_state(state_path, state)
        cell_state["status"] = "stopping"
        save_state(state_path, state)
    except Exception as error:  # noqa: BLE001
        failure = error
    finally:
        try:
            base.stop_entry(entry)
        except Exception as error:  # noqa: BLE001
            failure = failure or error
        time.sleep(2.0)
        try:
            wait_all_idle()
        except Exception as error:  # noqa: BLE001
            failure = failure or error

    cell_hours = base.live_gpu_hours([entry])
    state["gpu_hours_total"] += cell_hours
    cell_state["gpu_hours"] = cell_hours
    if failure is not None:
        cell_state["status"] = "failed"
        cell_state["failure"] = repr(failure)
        state["status"] = "failed"
        state["failures"].append({"cell": cell, "failure": repr(failure)})
        save_state(state_path, state)
        raise failure
    validation = base.validate_cell(entry)
    cell_state["validation"] = validation
    cell_state["status"] = "complete"
    cell_state["completed_at"] = time.time()
    state["completed_cells"] += 1
    save_state(state_path, state)


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser()
    result.add_argument("--manifest", type=Path, required=True)
    result.add_argument("--run-root", type=Path, required=True)
    result.add_argument("--aituner-root", type=Path, required=True)
    result.add_argument("--vllm-source", type=Path, required=True)
    result.add_argument("--venv", type=Path, required=True)
    result.add_argument("--model", type=Path, required=True)
    result.add_argument("--client", type=Path, required=True)
    return result


def main() -> None:
    args = parser().parse_args()
    manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
    if manifest["status"] != "PASS":
        raise RuntimeError("pilot manifest did not pass preflight")
    args.run_root.mkdir(parents=True, exist_ok=True)
    copied_manifest = args.run_root / "pilot-manifest.json"
    if not copied_manifest.exists():
        atomic_json(copied_manifest, manifest)
    configure_base(args, manifest)
    state_path = args.run_root / "controller-state.json"
    state = load_state(state_path, base.GPU_LIMIT)
    state["status"] = "running"
    save_state(state_path, state)
    for index, cell in enumerate(ORDER):
        execute_cell(
            index=index,
            cell=cell,
            manifest=manifest,
            run_root=args.run_root,
            state_path=state_path,
            state=state,
        )
    state["status"] = "complete"
    state["completed_at"] = time.time()
    save_state(state_path, state)
    print(json.dumps({
        "status": state["status"],
        "completed_cells": state["completed_cells"],
        "gpu_hours_total": state["gpu_hours_total"],
    }, sort_keys=True))


if __name__ == "__main__":
    main()