#!/usr/bin/env python3
from __future__ import annotations

import json
import math
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path


HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import pilot_controller as controller  # noqa: E402
import prepare_pilot as prepare  # noqa: E402
from analyze_pilot import campaign_gpu_accounting  # noqa: E402


@dataclass
class Request:
    row_id: str
    sampling_u: float
    arrival_s: float = 0.0
    prompt_tokens_hint: int = 1


def main() -> None:
    requests = [
        Request("a", 0.1),
        Request("b", 0.2),
        Request("c", 0.2),
        Request("d", 0.9),
    ]
    anchor, selected = prepare.attainable_anchor(requests, target_count=2)
    assert anchor == 0.2
    assert [request.row_id for request in selected] == ["a", "b", "c"]

    with tempfile.TemporaryDirectory() as temporary:
        root = Path(temporary)
        source = root / "source.jsonl"
        rows = []
        for index, role in enumerate(prepare.ROLES):
            rows.append(
                {
                    "request_id": role,
                    "timestamp": float(index),
                    "sampling_u": (index + 0.5) / len(prepare.ROLES),
                    "input_length": 16 + index,
                    "messages": [{"role": "user", "content": role}],
                }
            )
        source.write_text(
            "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8"
        )
        windows, stats = prepare.materialize_bands(
            source,
            {
                "window_id": "source",
                "trace_type": "chat",
                "window_start": 0.0,
                "window_end": 600.0,
            },
            root / "private",
        )
        assert windows.is_file()
        assert all(stats[role]["rows"] == 1 for role in prepare.ROLES)
        for role in prepare.ROLES:
            row = json.loads((root / "private" / "traces" / f"{role}.jsonl").read_text())
            assert row["fidelity_pilot_band"] == role
            assert abs(float(row["sampling_u"]) - 0.5) < 1e-12

        prior = root / "prior-state.json"
        primary = root / "primary-state.json"
        prior.write_text(
            json.dumps(
                {
                    "status": "failed",
                    "gpu_hours_total": 0.02,
                    "hard_cap_h20_hours": 3.5,
                }
            ),
            encoding="utf-8",
        )
        primary.write_text(
            json.dumps(
                {
                    "status": "complete",
                    "gpu_hours_total": 1.5,
                    "hard_cap_h20_hours": 3.5,
                }
            ),
            encoding="utf-8",
        )
        accounting = campaign_gpu_accounting(primary, (prior,))
        assert math.isclose(accounting["aggregate_h20_hours"], 1.52)
        assert all(accounting["invariants"].values())

    assert len(controller.ORDER) == 6
    assert set(controller.ORDER) == set(prepare.CELLS)
    assert math.isclose(
        sum(
            controller.CELL_ESTIMATE_H20_HOURS[int(config["tp"])]
            for config in prepare.CELLS.values()
        ) + controller.SAFETY_H20_HOURS,
        3.0,
    )
    selection = {
        "selected_count": 122,
        "request_id_order_sha256": "request-hash",
        "arrival_order_sha256": "arrival-hash",
        "input_length_order_sha256": "length-hash",
    }
    warmup = {
        "kind": "warmup",
        "selection": {"count": 16},
        "invariants": {
            "warmup_16": True,
            "warmup_exact_16": True,
            "warmup_long": True,
        },
    }
    controller.validate_result_selection(
        result=warmup,
        selection=selection,
        cell="tp1_mns8",
        role="burnin",
        warmup=True,
    )
    measured = {
        "kind": "anchor",
        "selection": {
            "count": 122,
            "request_id_order_sha256": "request-hash",
            "arrival_order_sha256": "arrival-hash",
            "raw_length_order_sha256": "length-hash",
        },
        "invariants": {},
    }
    controller.validate_result_selection(
        result=measured,
        selection=selection,
        cell="tp1_mns8",
        role="low1",
        warmup=False,
    )
    try:
        controller.validate_result_selection(
            result=warmup,
            selection=selection,
            cell="tp1_mns8",
            role="low1",
            warmup=False,
        )
    except RuntimeError as error:
        assert "selection count mismatch" in str(error)
    else:
        raise AssertionError("measured selection accepted a warmup subset")
    print("fidelity pilot tools: PASS")


if __name__ == "__main__":
    main()