#!/usr/bin/env python3
"""A-P6-2 serialized solo controller for authoritative Phase-6 frontiers."""

from __future__ import annotations

import json
import math
import os
import shlex
import subprocess
import time
from pathlib import Path
from typing import Any

import opprof_phase6_controller as base


SOLO_ROOT = base.RUN_ROOT / "solo-authoritative"
STATE = SOLO_ROOT / "controller-state.json"
CAMPAIGN_STATE = base.RUN_ROOT / "controller-state.json"
GPU_LIMIT = 6.0
SAFETY_HOURS = 0.20
MARKER = "opprof-phase6-solo-A-P6-2"
TRACE = base.AITUNER / "trace_windows/traces/chat_w20260311_1000.jsonl"

ORDER = [
    "tp4_mns32", "tp4_mns64", "tp2_mns32", "tp2_mns64",
    "tp4_mns16", "tp2_mns8", "tp2_mns16", "tp4_mns8",
    "tp1_mns8", "tp1_mns16", "tp1_mns32", "tp1_mns64",
]

# Exact co-located primaries are remeasured before adaptive crawl. W4 had no
# prior primary, so it starts at P and selects L/U from the solo result.
CORE = {
    "tp1_mns8": [base.CELLS["tp1_mns8"]["peak"], base.CELLS["tp1_mns8"]["lower"]],
    "tp1_mns16": [base.CELLS["tp1_mns16"]["peak"], base.CELLS["tp1_mns16"]["upper"]],
    "tp1_mns32": [base.CELLS["tp1_mns32"]["peak"], base.CELLS["tp1_mns32"]["upper"]],
    "tp1_mns64": [base.CELLS["tp1_mns64"]["peak"], base.CELLS["tp1_mns64"]["upper"]],
    "tp2_mns8": [base.CELLS["tp2_mns8"]["peak"], base.CELLS["tp2_mns8"]["lower"]],
    "tp2_mns16": [base.CELLS["tp2_mns16"]["peak"], base.CELLS["tp2_mns16"]["lower"]],
    "tp2_mns32": [base.CELLS["tp2_mns32"]["peak"], base.CELLS["tp2_mns32"]["lower"]],
    "tp2_mns64": [base.CELLS["tp2_mns64"]["peak"], base.CELLS["tp2_mns64"]["lower"]],
    "tp4_mns8": [base.CELLS["tp4_mns8"]["peak"], base.CELLS["tp4_mns8"]["lower"]],
    "tp4_mns16": [
        base.CELLS["tp4_mns16"]["peak"], base.CELLS["tp4_mns16"]["lower"],
        base.CELLS["tp4_mns16"]["upper"],
    ],
    "tp4_mns32": [base.CELLS["tp4_mns32"]["peak"]],
    "tp4_mns64": [base.CELLS["tp4_mns64"]["peak"]],
}

CELL_ESTIMATE = {cell: {1: .11, 2: .22, 4: .48}[cfg["tp"]] for cell, cfg in base.CELLS.items()}


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


def load_state() -> dict[str, Any]:
    if STATE.exists():
        return json.loads(STATE.read_text())
    campaign = json.loads(CAMPAIGN_STATE.read_text())
    return {
        "schema": 1, "amendment": "A-P6-2", "status": "initialized",
        "hard_cap_h20_hours": GPU_LIMIT,
        "prior_h20_hours": float(campaign["gpu_hours_total"]),
        "gpu_hours_total": float(campaign["gpu_hours_total"]),
        "solo_gpu_hours": 0.0, "completed_cells": 0,
        "primary_anchors": 0, "confirmations": 0,
        "cells": {}, "failures": [], "started_at": time.time(),
    }


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


def historical() -> tuple[dict[tuple[str, float], dict[str, Any]], dict[str, list[float]]]:
    ground = json.loads(base.GROUND.read_text())
    expected = {}
    histories = {}
    for cell in ground["cells"]:
        anchors = []
        for probe in cell["probe_history"]:
            anchor = float(probe["sampling_u"])
            expected[(cell["cell_id"], anchor)] = probe
            anchors.append(anchor)
        histories[cell["cell_id"]] = sorted(anchors)
    return expected, histories


def same_anchor(left: float, right: float) -> bool:
    return math.isclose(left, right, rel_tol=0, abs_tol=1e-15)


def colocated_primary(cell: str, anchor: float) -> dict[str, Any] | None:
    cell_dir = base.RUN_ROOT / "cells" / cell
    for path in cell_dir.glob("anchor-*/result.json"):
        item = json.loads(path.read_text())
        if same_anchor(float(item["anchor"]), anchor):
            return item
    return None


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


def wait_all_idle(timeout: float = 30.0) -> None:
    deadline = time.monotonic() + timeout
    error = None
    while time.monotonic() < deadline:
        try:
            base.assert_all_idle()
            return
        except RuntimeError as current:
            error = current
            time.sleep(1)
    raise error or RuntimeError("GPU cleanup did not reach idle")


def remaining_projection(index: int) -> float:
    return sum(CELL_ESTIMATE[cell] for cell in ORDER[index:]) + SAFETY_HOURS


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


def run_one(
    entry: dict[str, Any], anchor: float, out: Path, state: dict[str, Any],
    cell_state: dict[str, Any], role: str,
) -> dict[str, Any]:
    result = base.run_clients(
        [entry], [(entry, anchor, out)], state, f"solo-{entry['cell']}"
    )[0]
    expected_count = cell_state["expected_counts"][str(anchor)]
    if int(result["selection"]["count"]) != int(expected_count):
        raise RuntimeError(
            f"selection mismatch {entry['cell']} {anchor}: "
            f"{result['selection']['count']} != {expected_count}"
        )
    cell_state.setdefault("runs", []).append({
        "anchor": anchor, "role": role, "dir": str(out),
        "pass_rate": result["pass_rate"], "feasible": result["feasible"],
    })
    save_state(state)
    return result


def anchor_trials(cell_state: dict[str, Any], anchor: float) -> list[dict[str, Any]]:
    return [
        item for item in cell_state.get("runs", [])
        if same_anchor(float(item["anchor"]), anchor)
    ]


def accepted_feasible(cell_state: dict[str, Any], anchor: float) -> bool | None:
    trials = anchor_trials(cell_state, anchor)
    votes = [bool(item["feasible"]) for item in trials]
    if not votes:
        return None
    if len(votes) == 1 or len(set(votes)) == 1:
        return votes[0]
    if len(votes) >= 3:
        return sum(votes) >= 2
    return None


def optional_fits(
    state: dict[str, Any], entry: dict[str, Any], future_after: float,
) -> bool:
    replay = len(entry["gpus"]) * 80 / 3600
    projected = (
        float(state["gpu_hours_total"]) + base.live_gpu_hours([entry])
        + future_after + replay + SAFETY_HOURS
    )
    return projected < GPU_LIMIT


def maybe_confirm(
    entry: dict[str, Any], anchor: float, primary: dict[str, Any],
    state: dict[str, Any], cell_state: dict[str, Any], expected: dict[tuple[str, float], dict[str, Any]],
    future_after: float,
) -> None:
    old = expected[(entry["cell"], anchor)]
    coloc = colocated_primary(entry["cell"], anchor)
    disagreement = (
        bool(primary["feasible"]) != bool(old["feasible"])
        or (coloc is not None and bool(primary["feasible"]) != bool(coloc["feasible"]))
    )
    boundary = .93 <= float(primary["pass_rate"]) <= .97
    if not (disagreement or boundary):
        return
    while len(anchor_trials(cell_state, anchor)) < 3:
        trials = anchor_trials(cell_state, anchor)
        if len(trials) >= 2 and len({bool(item["feasible"]) for item in trials}) == 1:
            return
        if not optional_fits(state, entry, future_after):
            cell_state.setdefault("deferred_confirmations", []).append(anchor)
            return
        trial = len(trials) + 1
        out = entry["dir"] / f"confirm-{trial - 1}-anchor-{anchor}"
        run_one(entry, anchor, out, state, cell_state, f"confirmation-{trial}")
        state["confirmations"] += 1


def run_primary(
    entry: dict[str, Any], anchor: float, state: dict[str, Any],
    cell_state: dict[str, Any], expected: dict[tuple[str, float], dict[str, Any]],
    future_after: float, role: str,
) -> dict[str, Any]:
    existing = [item for item in cell_state.get("runs", []) if item["role"].startswith("primary")]
    if any(same_anchor(float(item["anchor"]), anchor) for item in existing):
        path = next(
            Path(item["dir"]) for item in existing
            if same_anchor(float(item["anchor"]), anchor)
        )
        return json.loads((path / "result.json").read_text())
    out = entry["dir"] / f"anchor-{anchor}"
    result = run_one(entry, anchor, out, state, cell_state, role)
    state["primary_anchors"] += 1
    maybe_confirm(entry, anchor, result, state, cell_state, expected, future_after)
    return result


def next_below(history: list[float], tested: set[float]) -> float | None:
    if not tested:
        return None
    candidates = [x for x in history if x < min(tested) and x not in tested]
    return max(candidates) if candidates else None


def next_above(history: list[float], tested: set[float]) -> float | None:
    if not tested:
        return None
    candidates = [x for x in history if x > max(tested) and x not in tested]
    return min(candidates) if candidates else None


def execute_cell(
    index: int, cell: str, state: dict[str, Any],
    expected: dict[tuple[str, float], dict[str, Any]], histories: dict[str, list[float]],
) -> None:
    if state["cells"].get(cell, {}).get("status") == "complete":
        return
    future = remaining_projection(index)
    if float(state["gpu_hours_total"]) + future >= GPU_LIMIT:
        state["status"] = "budget_projection_stop"
        state["budget_stop"] = {
            "before_cell": cell, "spent_h20_hours": state["gpu_hours_total"],
            "remaining_projection_h20_hours": future,
            "projected_total_h20_hours": state["gpu_hours_total"] + future,
            "hard_cap_h20_hours": GPU_LIMIT,
        }
        save_state(state)
        raise RuntimeError(f"projected budget exceeds cap before {cell}")
    cfg = base.CELLS[cell]
    echo = (
        f"SOLO_WAVE_ECHO cell={cell} tp={cfg['tp']} mns={cfg['mns']} "
        f"gpus=0-{int(cfg['tp'])-1} mandatory={','.join(map(str, CORE[cell]))} "
        f"spent_h20h={state['gpu_hours_total']:.6f} cell_est_h20h={CELL_ESTIMATE[cell]:.3f} "
        f"remaining_projection_h20h={future:.3f} cap_h20h={GPU_LIMIT:.1f} "
        f"ground_truth={base.GROUND} trace={TRACE}"
    )
    append_echo(echo)
    wait_all_idle()
    cell_state = {
        "status": "starting", "started_at": time.time(), "tp": cfg["tp"], "mns": cfg["mns"],
        "mandatory": CORE[cell],
        "expected_counts": {
            str(anchor): expected[(cell, anchor)]["request_count"] for anchor in histories[cell]
        },
        "runs": [],
    }
    state["status"] = "running"
    state["cells"][cell] = cell_state
    save_state(state)
    entry = start_entry(cell, index)
    failure = None
    future_after = sum(CELL_ESTIMATE[item] for item in ORDER[index + 1:])
    try:
        base.wait_ready(entry)
        cell_state["status"] = "warmup"
        save_state(state)
        warm = base.run_clients(
            [entry], [(entry, cfg["peak"], entry["dir"] / "warmup")],
            state, f"solo-{cell}", warmup=True,
        )[0]
        cell_state["warmup"] = {
            "exact_output_count": warm["exact_output_count"],
            "long_gt4096": warm["selection"]["long_gt4096"],
        }
        cell_state["status"] = "mandatory"
        save_state(state)
        for anchor in CORE[cell]:
            run_primary(entry, anchor, state, cell_state, expected, future_after, "primary-mandatory")

        peak = float(cfg["peak"])
        peak_vote = accepted_feasible(cell_state, peak)
        if cell in {"tp4_mns32", "tp4_mns64"} and peak_vote is not None:
            direction = float(cfg["upper"] if peak_vote else cfg["lower"])
            run_primary(entry, direction, state, cell_state, expected, future_after, "primary-direction")

        cell_state["status"] = "crawl"
        save_state(state)
        while True:
            primary_anchors = {
                float(item["anchor"]) for item in cell_state["runs"]
                if item["role"].startswith("primary")
            }
            votes = {anchor: accepted_feasible(cell_state, anchor) for anchor in primary_anchors}
            pass_anchors = [anchor for anchor, vote in votes.items() if vote is True]
            fail_anchors = [anchor for anchor, vote in votes.items() if vote is False]
            if pass_anchors and fail_anchors and max(pass_anchors) < min(fail_anchors):
                break
            if any(vote is None for vote in votes.values()):
                cell_state["censor"] = "UNRESOLVED_SOLO_ANCHOR"
                break
            if pass_anchors and not fail_anchors:
                candidate = next_above(histories[cell], primary_anchors)
            elif fail_anchors and not pass_anchors:
                candidate = next_below(histories[cell], primary_anchors)
            else:
                # Non-monotonic anchors already contain both states but no valid bracket.
                cell_state["censor"] = "NONMONOTONIC_SOLO_ANCHORS"
                break
            if candidate is None:
                cell_state["censor"] = "HISTORY_EDGE"
                break
            if not optional_fits(state, entry, future_after):
                cell_state["censor"] = "BUDGET_CENSORED"
                break
            run_primary(entry, candidate, state, cell_state, expected, future_after, "primary-crawl")
        cell_state["status"] = "stopping"
        save_state(state)
    except Exception as error:
        failure = error
    finally:
        try:
            base.stop_entry(entry)
        except Exception as error:
            failure = failure or error
        time.sleep(2)
        try:
            wait_all_idle()
        except Exception as error:
            failure = failure or error
    hours = base.live_gpu_hours([entry])
    state["gpu_hours_total"] += hours
    state["solo_gpu_hours"] += hours
    cell_state["gpu_hours"] = 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)
        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)


def main() -> None:
    SOLO_ROOT.mkdir(parents=True, exist_ok=True)
    base.GPU_LIMIT = GPU_LIMIT
    base.MARKER = MARKER
    expected, histories = historical()
    state = load_state()
    state["status"] = "running"
    save_state(state)
    for index, cell in enumerate(ORDER):
        execute_cell(index, cell, state, expected, histories)
    state["status"] = "complete"
    state["completed_at"] = time.time()
    save_state(state)
    print(json.dumps({
        "status": state["status"], "cells": state["completed_cells"],
        "primary_anchors": state["primary_anchors"],
        "confirmations": state["confirmations"],
        "solo_gpu_hours": state["solo_gpu_hours"],
        "campaign_gpu_hours": state["gpu_hours_total"],
    }, sort_keys=True))


if __name__ == "__main__":
    main()