#!/usr/bin/env python3
"""Serialized controller for the crossed-constraint action-aware pilot."""

from __future__ import annotations

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


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

import opprof_phase6_controller as base  # noqa: E402


SCHEMA = "action-aware-constraint-pilot-state-v0"


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(args: argparse.Namespace, manifest: Mapping[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["budget"]["hard_cap_h20_hours"])
    base.MARKER = "action-aware-constraint-pilot-v0"


def validate_inputs(args: argparse.Namespace, manifest: Mapping[str, Any]) -> None:
    if manifest.get("schema") not in {
        "action-aware-constraint-pilot-manifest-v0",
        "action-aware-constraint-pilot-manifest-v1",
    }:
        raise RuntimeError("unexpected action-aware manifest schema")
    if manifest.get("status") != "PASS":
        raise RuntimeError("action-aware manifest did not pass preflight")
    red_flags = manifest.get("sanity", {}).get("red_flags", [])
    if red_flags:
        raise RuntimeError(f"manifest red flags: {red_flags}")

    required = {
        "manifest": args.manifest,
        "aituner_root": args.aituner_root,
        "vllm_source": args.vllm_source,
        "venv_python": args.venv / "bin/python",
        "venv_vllm": args.venv / "bin/vllm",
        "model": args.model,
        "client": args.client,
        "burnin_study": Path(manifest["burnin"]["study"]),
    }
    for repetition, item in manifest["repetitions"].items():
        required[f"rep{repetition}_study"] = Path(item["study"])
        required[f"rep{repetition}_trace"] = Path(item["merged_trace"]["path"])
    missing = {name: str(path) for name, path in required.items() if not path.exists()}
    if missing:
        raise RuntimeError(f"action-aware input paths missing: {missing}")


def config_map(manifest: Mapping[str, Any]) -> dict[str, dict[str, Any]]:
    return {str(item["id"]): dict(item) for item in manifest["configs"]}


def server_command(
    config: Mapping[str, Any], *, gpus: tuple[int, ...], port: int
) -> list[str]:
    return [
        "taskset",
        "-c",
        base.cpu_mask(gpus),
        str(base.VENV / "bin/vllm"),
        "serve",
        str(base.MODEL),
        "--host",
        "127.0.0.1",
        "--port",
        str(port),
        "--served-model-name",
        "qwen3-30b-a3b-community",
        "--max-num-batched-tokens",
        str(config["mbbt"]),
        "--max-num-seqs",
        str(config["mns"]),
        "--tensor-parallel-size",
        "4",
        "--shutdown-timeout",
        "120",
    ]


def client_command(
    entry: Mapping[str, Any],
    config: Mapping[str, Any],
    *,
    study: str,
    anchor: float,
    output: Path,
    warmup: bool,
) -> list[str]:
    command = [
        "taskset",
        "-c",
        base.cpu_mask(entry["gpus"]),
        str(base.VENV / "bin/python"),
        str(base.CLIENT),
        "warmup" if warmup else "run-anchor",
        "--study",
        study,
        "--cell",
        str(config["id"]),
        "--anchor",
        str(anchor),
        "--tp",
        "4",
        "--mns",
        str(config["mns"]),
        "--mbbt",
        str(config["mbbt"]),
        "--base-url",
        f"http://127.0.0.1:{entry['port']}",
        "--result-dir",
        str(output),
        "--disable-slo-early-stop",
    ]
    return command


def remaining_projection(
    manifest: Mapping[str, Any], *, completed_sessions: int
) -> float:
    remaining = len(manifest["configs"]) - completed_sessions
    return (
        remaining * float(manifest["budget"]["session_estimate_h20_hours"])
        + float(manifest["budget"]["safety_h20_hours"])
    )


def dry_run_plan(
    args: argparse.Namespace, manifest: Mapping[str, Any]
) -> dict[str, Any]:
    sessions = []
    for index, config in enumerate(manifest["configs"]):
        entry = {"gpus": (0, 1, 2, 3), "port": 9050 + index}
        session_root = args.run_root / "sessions" / str(config["id"])
        first_repetition = str(config["repetition_order"][0])
        first = manifest["repetitions"][first_repetition]
        commands = {
            "server": server_command(config, gpus=entry["gpus"], port=entry["port"]),
            "warmup": client_command(
                entry,
                config,
                study=first["study"],
                anchor=float(first["selection"]["anchor"]),
                output=session_root / "warmup",
                warmup=True,
            ),
            "burnin": client_command(
                entry,
                config,
                study=manifest["burnin"]["study"],
                anchor=float(manifest["burnin"]["anchor"]),
                output=session_root / "burnin",
                warmup=False,
            ),
        }
        for repetition in config["repetition_order"]:
            item = manifest["repetitions"][str(repetition)]
            commands[f"rep{repetition}"] = client_command(
                entry,
                config,
                study=item["study"],
                anchor=float(item["selection"]["anchor"]),
                output=session_root / f"rep{repetition}",
                warmup=False,
            )
        sessions.append(
            {
                "config": config["id"],
                "mns": config["mns"],
                "mbbt": config["mbbt"],
                "port": entry["port"],
                "repetition_order": config["repetition_order"],
                "commands": {
                    role: shlex.join(command) for role, command in commands.items()
                },
            }
        )
    return {
        "schema": "action-aware-constraint-pilot-dry-run-v0",
        "status": "PASS",
        "manifest": str(args.manifest),
        "run_root": str(args.run_root),
        "projected_h20_hours": remaining_projection(
            manifest, completed_sessions=0
        ),
        "hard_cap_h20_hours": manifest["budget"]["hard_cap_h20_hours"],
        "sessions": sessions,
    }


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": SCHEMA,
        "status": "initialized",
        "hard_cap_h20_hours": hard_cap,
        "gpu_hours_total": 0.0,
        "completed_sessions": 0,
        "sessions": {},
        "failures": [],
        "started_at": time.time(),
    }


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 start_server(
    *,
    args: argparse.Namespace,
    config: Mapping[str, Any],
    index: int,
) -> dict[str, Any]:
    gpus = (0, 1, 2, 3)
    session_root = args.run_root / "sessions" / str(config["id"])
    session_root.mkdir(parents=True, exist_ok=True)
    port = 9050 + index
    command = server_command(config, gpus=gpus, port=port)
    with (session_root / "commands.log").open("a", encoding="utf-8") as log:
        log.write(f"SERVER {shlex.join(command)}\n")
    server_log = (session_root / "server.log").open("ab", buffering=0)
    environment = os.environ.copy()
    environment.update(
        {
            "CUDA_VISIBLE_DEVICES": "0,1,2,3",
            "VLLM_OPPROF_DIR": str(session_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": str(config["id"]),
        "gpus": gpus,
        "port": port,
        "dir": session_root,
        "server": server,
        "server_handle": server_log,
        "spawned_at": time.time(),
        "results": [],
    }


def validate_result(
    result: Mapping[str, Any],
    *,
    config: Mapping[str, Any],
    selection: Mapping[str, Any],
    role: str,
    warmup: bool,
) -> None:
    if result.get("schema") != "action-aware-pilot-result-v0":
        raise RuntimeError(f"unexpected result schema: {role}")
    if result.get("config_id") != config["id"]:
        raise RuntimeError(f"config id mismatch: {role}")
    if int(result["tp"]) != 4:
        raise RuntimeError(f"TP mismatch: {role}")
    if int(result["mns"]) != int(config["mns"]):
        raise RuntimeError(f"MNS mismatch: {role}")
    if int(result["mbbt"]) != int(config["mbbt"]):
        raise RuntimeError(f"MBBT mismatch: {role}")
    if result.get("slo_early_stop_disabled") is not True:
        raise RuntimeError(f"SLO early stop was not disabled: {role}")
    if warmup:
        if result["kind"] != "warmup" or int(result["selection"]["count"]) != 16:
            raise RuntimeError(f"invalid warmup: {role}")
        return
    if bool(result["early_stopped"]):
        raise RuntimeError(f"uncensored run early-stopped: {role}")
    if int(result["selection"]["count"]) != int(selection["selected_count"]):
        raise RuntimeError(f"selection count mismatch: {role}")
    if int(result["observed_count"]) != int(selection["selected_count"]):
        raise RuntimeError(f"request accounting mismatch: {role}")
    for result_key, selection_key in (
        ("request_id_order_sha256", "request_id_order_sha256"),
        ("arrival_order_sha256", "arrival_order_sha256"),
        ("raw_length_order_sha256", "input_length_order_sha256"),
    ):
        if result["selection"][result_key] != selection[selection_key]:
            raise RuntimeError(f"selection hash mismatch {result_key}: {role}")


def burnin_gate(
    result: Mapping[str, Any],
    *,
    expected_count: int,
    maximum_elapsed_s: float,
) -> dict[str, Any]:
    if result.get("kind") != "anchor":
        raise RuntimeError("burnin gate received a non-anchor result")
    if int(result["selection"]["count"]) != expected_count:
        raise RuntimeError("burnin gate received the wrong request set")
    elapsed_s = float(result["interval"]["elapsed_s"])
    summary = {
        "elapsed_s": elapsed_s,
        "pass_rate": float(result["pass_rate"]),
        "feasible": bool(result["feasible"]),
    }
    if elapsed_s > maximum_elapsed_s:
        raise RuntimeError(
            f"burnin throughput gate failed: {elapsed_s:.3f}s > "
            f"{maximum_elapsed_s:.3f}s"
        )
    return summary


def run_client(
    *,
    entry: dict[str, Any],
    config: Mapping[str, Any],
    role: str,
    study: str,
    selection: Mapping[str, Any],
    output: Path,
    state: Mapping[str, Any],
    timeout_s: float,
    warmup: bool = False,
) -> dict[str, Any]:
    command = client_command(
        entry,
        config,
        study=study,
        anchor=float(selection["anchor"]),
        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() + timeout_s
    try:
        while process.poll() is None:
            if time.monotonic() > deadline:
                raise TimeoutError(f"client timeout: {config['id']} {role}")
            if entry["server"].poll() is not None:
                raise RuntimeError(f"server exited during {config['id']} {role}")
            base.assert_no_other_compute()
            if state["gpu_hours_total"] + base.live_gpu_hours([entry]) >= base.GPU_LIMIT:
                raise RuntimeError("action-aware pilot H20-hour hard cap reached")
            time.sleep(1.0)
    except Exception:
        try:
            os.killpg(process.pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
        try:
            process.wait(timeout=10.0)
        except subprocess.TimeoutExpired:
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            process.wait(timeout=10.0)
        raise
    finally:
        handle.close()
    if process.returncode:
        raise RuntimeError(
            f"client failed: config={config['id']} role={role} rc={process.returncode}"
        )
    result = json.loads((output / "result.json").read_text(encoding="utf-8"))
    validate_result(
        result,
        config=config,
        selection=selection,
        role=role,
        warmup=warmup,
    )
    entry["results"].append(
        {"anchor": float(selection["anchor"]), "dir": str(output), "kind": result["kind"]}
    )
    return result


def execute_session(
    *,
    args: argparse.Namespace,
    manifest: Mapping[str, Any],
    config: Mapping[str, Any],
    index: int,
    state: dict[str, Any],
    state_path: Path,
) -> None:
    name = str(config["id"])
    if state["sessions"].get(name, {}).get("status") == "complete":
        return
    projection = remaining_projection(
        manifest, completed_sessions=int(state["completed_sessions"])
    )
    if float(state["gpu_hours_total"]) + projection > base.GPU_LIMIT:
        raise RuntimeError(f"projected cost exceeds cap before {name}")
    load_values = {
        float(item["selection"]["offered_req_s_per_gpu"])
        for item in manifest["repetitions"].values()
    }
    load_text = (
        f"{next(iter(load_values)):.6g}"
        if len(load_values) == 1
        else ",".join(f"{value:.6g}" for value in sorted(load_values))
    )
    echo = (
        f"ACTION_AWARE_SESSION_ECHO host=dash0 config={name} tp=4 "
        f"mns={config['mns']} mbbt={config['mbbt']} gpus=0-3 "
        f"workload={manifest['source']['window_id']} load_per_gpu={load_text} "
        f"duration_s={manifest['engine']['duration_s']} "
        f"repetitions={','.join(map(str, config['repetition_order']))} "
        f"source={args.manifest} output={args.run_root / 'sessions' / name} "
        f"spent_h20h={state['gpu_hours_total']:.6f} "
        f"remaining_projection_h20h={projection:.3f} cap_h20h={base.GPU_LIMIT:.1f}"
    )
    append_echo(args.run_root, echo)
    wait_all_idle()
    session_state = {
        "status": "starting",
        "mns": int(config["mns"]),
        "mbbt": int(config["mbbt"]),
        "repetition_order": list(config["repetition_order"]),
        "started_at": time.time(),
        "runs": [],
    }
    state["status"] = "running"
    state["sessions"][name] = session_state
    atomic_json(state_path, state)
    entry = start_server(args=args, config=config, index=index)
    failure: Exception | None = None
    try:
        base.wait_ready(entry)
        first = manifest["repetitions"][str(config["repetition_order"][0])]
        session_state["status"] = "warmup"
        atomic_json(state_path, state)
        run_client(
            entry=entry,
            config=config,
            role="warmup",
            study=first["study"],
            selection=first["selection"],
            output=entry["dir"] / "warmup",
            state=state,
            timeout_s=180.0,
            warmup=True,
        )
        session_state["status"] = "burnin"
        atomic_json(state_path, state)
        burnin = manifest["burnin"]
        burnin_result = run_client(
            entry=entry,
            config=config,
            role="burnin",
            study=burnin["study"],
            selection=burnin,
            output=entry["dir"] / "burnin",
            state=state,
            timeout_s=float(manifest["engine"]["client_timeout_s"]),
        )
        session_state["burnin"] = burnin_gate(
            burnin_result,
            expected_count=int(burnin["selected_count"]),
            maximum_elapsed_s=float(manifest["engine"]["burnin_max_elapsed_s"]),
        )
        atomic_json(state_path, state)
        session_state["status"] = "measured"
        atomic_json(state_path, state)
        for repetition in config["repetition_order"]:
            item = manifest["repetitions"][str(repetition)]
            role = f"rep{repetition}"
            result = run_client(
                entry=entry,
                config=config,
                role=role,
                study=item["study"],
                selection=item["selection"],
                output=entry["dir"] / role,
                state=state,
                timeout_s=float(manifest["engine"]["client_timeout_s"]),
            )
            session_state["runs"].append(
                {
                    "repetition": int(repetition),
                    "pass_rate": result["pass_rate"],
                    "feasible": result["feasible"],
                    "slo_pass_count": result["slo_pass_count"],
                    "elapsed_s": result["interval"]["elapsed_s"],
                }
            )
            atomic_json(state_path, state)
        session_state["status"] = "stopping"
        atomic_json(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

    session_hours = base.live_gpu_hours([entry])
    state["gpu_hours_total"] += session_hours
    session_state["gpu_hours"] = session_hours
    if failure is not None:
        session_state["status"] = "failed"
        session_state["failure"] = repr(failure)
        state["status"] = "failed"
        state["failures"].append({"session": name, "failure": repr(failure)})
        atomic_json(state_path, state)
        raise failure
    validation = base.validate_cell(entry)
    session_state["validation"] = validation
    session_state["status"] = "complete"
    session_state["completed_at"] = time.time()
    state["completed_sessions"] += 1
    atomic_json(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)
    result.add_argument("--dry-run", action="store_true")
    return result


def main() -> None:
    args = parser().parse_args()
    manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
    validate_inputs(args, manifest)
    configure(args, manifest)
    if args.dry_run:
        print(json.dumps(dry_run_plan(args, manifest), indent=2, sort_keys=True))
        return
    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)
    state_path = args.run_root / "controller-state.json"
    state = load_state(state_path, base.GPU_LIMIT)
    state["status"] = "running"
    atomic_json(state_path, state)
    for index, config in enumerate(manifest["configs"]):
        execute_session(
            args=args,
            manifest=manifest,
            config=config,
            index=index,
            state=state,
            state_path=state_path,
        )
    state["status"] = "complete"
    state["completed_at"] = time.time()
    atomic_json(state_path, state)
    wait_all_idle()
    print(
        json.dumps(
            {
                "status": state["status"],
                "completed_sessions": state["completed_sessions"],
                "gpu_hours_total": state["gpu_hours_total"],
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()