#!/usr/bin/env python3
"""Materialize session-disjoint pilot repeats and freeze attainable anchors.

The private outputs retain prompt text and stay on the experiment host.  The
public manifest contains only aggregate counts, hashes, paths, and parameters.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import sys
from pathlib import Path
from typing import Any


AITUNER_ROOT = Path(os.environ.get("AITUNER_ROOT", Path(__file__).resolve().parents[2]))
sys.path.insert(0, str(AITUNER_ROOT / "src"))

from aituner.spec import load_study_spec  # noqa: E402
from aituner.trace import load_trace_requests, select_requests_for_threshold  # noqa: E402


ROLES = ("burnin", "low1", "high1", "low2", "high2", "low3", "high3")
CELLS = {
    "tp1_mns8": {"tp": 1, "mns": 8, "frontier_req_s_gpu": 2.3833333333333333},
    "tp1_mns64": {"tp": 1, "mns": 64, "frontier_req_s_gpu": 2.3833333333333333},
    "tp2_mns8": {"tp": 2, "mns": 8, "frontier_req_s_gpu": 2.2416666666666667},
    "tp2_mns64": {"tp": 2, "mns": 64, "frontier_req_s_gpu": 2.3},
    "tp4_mns16": {"tp": 4, "mns": 16, "frontier_req_s_gpu": 2.5},
    "tp4_mns64": {"tp": 4, "mns": 64, "frontier_req_s_gpu": 2.5},
}
TARGET_MULTIPLIERS = {"low": 0.85, "high": 1.25}


def atomic_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    os.replace(tmp, path)


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def order_hash(values: list[str]) -> str:
    return hashlib.sha256("\n".join(values).encode()).hexdigest()


def resolve_source_trace(windows_path: Path, window_id: str) -> tuple[dict[str, Any], Path]:
    payload = json.loads(windows_path.read_text(encoding="utf-8"))
    for window in payload["windows"]:
        if window["window_id"] != window_id:
            continue
        trace = Path(window["trace_file"])
        if not trace.is_absolute():
            trace = (windows_path.parent / trace).resolve()
        return window, trace
    raise ValueError(f"window not found: {window_id}")


def materialize_bands(
    source_trace: Path,
    source_window: dict[str, Any],
    private_root: Path,
) -> tuple[Path, dict[str, Any]]:
    traces_root = private_root / "traces"
    traces_root.mkdir(parents=True, exist_ok=True)
    temporary = {role: traces_root / f".{role}.jsonl.tmp" for role in ROLES}
    final = {role: traces_root / f"{role}.jsonl" for role in ROLES}
    handles = {role: temporary[role].open("w", encoding="utf-8") for role in ROLES}
    stats = {
        role: {
            "rows": 0,
            "sum_input_tokens": 0,
            "min_timestamp": None,
            "max_timestamp": None,
        }
        for role in ROLES
    }
    try:
        with source_trace.open(encoding="utf-8") as source:
            for line_number, line in enumerate(source):
                row = json.loads(line)
                value = float(row["sampling_u"])
                if not 0.0 <= value <= 1.0:
                    raise ValueError(f"sampling_u outside [0,1] at line {line_number}")
                band = min(len(ROLES) - 1, int(value * len(ROLES)))
                role = ROLES[band]
                remapped = value * len(ROLES) - band
                row["sampling_u"] = min(remapped, math.nextafter(1.0, 0.0))
                row["fidelity_pilot_band"] = role
                handles[role].write(json.dumps(row, ensure_ascii=False) + "\n")
                timestamp = float(row["timestamp"])
                item = stats[role]
                item["rows"] += 1
                item["sum_input_tokens"] += int(row.get("input_length") or 0)
                item["min_timestamp"] = (
                    timestamp if item["min_timestamp"] is None
                    else min(float(item["min_timestamp"]), timestamp)
                )
                item["max_timestamp"] = (
                    timestamp if item["max_timestamp"] is None
                    else max(float(item["max_timestamp"]), timestamp)
                )
    finally:
        for handle in handles.values():
            handle.close()
    for role in ROLES:
        os.replace(temporary[role], final[role])
        stats[role]["sha256"] = sha256_file(final[role])
        stats[role]["bytes"] = final[role].stat().st_size

    windows = []
    for role in ROLES:
        window = dict(source_window)
        window["window_id"] = f"fidelity_pilot_{role}"
        window["trace_file"] = f"traces/{role}.jsonl"
        window["num_requests"] = stats[role]["rows"]
        window["sum_input_length"] = stats[role]["sum_input_tokens"]
        window["sampling_strategy"] = "session_uniform_seven_disjoint_bands_remapped"
        window["fidelity_pilot_role"] = role
        windows.append(window)
    private_windows = private_root / "windows.json"
    atomic_json(
        private_windows,
        {
            "schema": "fidelity-pilot-private-windows-v1",
            "roles": list(ROLES),
            "windows": windows,
        },
    )
    return private_windows, stats


def write_studies(
    *,
    base_primary: Path,
    base_tp4: Path,
    private_windows: Path,
    private_root: Path,
) -> dict[str, dict[str, Path]]:
    bases = {
        "primary": json.loads(base_primary.read_text(encoding="utf-8")),
        "tp4": json.loads(base_tp4.read_text(encoding="utf-8")),
    }
    result: dict[str, dict[str, Path]] = {}
    for role in ROLES:
        result[role] = {}
        for tier, base in bases.items():
            payload = json.loads(json.dumps(base))
            payload["study_id"] = f"fidelity-prefix-pilot-{role}-{tier}"
            payload["hardware"]["host_candidates"] = ["dash0"]
            payload["engine"]["engine_version"] = "0.24.1.dev3+opprof"
            payload["trace"]["windows_path"] = str(private_windows)
            payload["trace"]["window_id"] = f"fidelity_pilot_{role}"
            path = private_root / "studies" / f"{role}-{tier}.json"
            atomic_json(path, payload)
            result[role][tier] = path
    return result


def attainable_anchor(requests: list[Any], target_count: int) -> tuple[float, list[Any]]:
    ordered = sorted(float(request.sampling_u) for request in requests)
    if not ordered:
        raise ValueError("no requests after study filtering")
    candidate_indices = sorted({
        max(0, min(len(ordered) - 1, target_count - 1)),
        max(0, min(len(ordered) - 1, target_count)),
    })
    candidates = []
    for index in candidate_indices:
        anchor = ordered[index]
        selected = select_requests_for_threshold(requests, threshold=anchor)
        candidates.append((abs(len(selected) - target_count), len(selected), anchor, selected))
    _error, _count, anchor, selected = min(candidates, key=lambda item: (item[0], item[1]))
    return anchor, selected


def selected_record(selected: list[Any], *, tp: int, duration_s: float) -> dict[str, Any]:
    return {
        "anchor": max(float(request.sampling_u) for request in selected),
        "selected_count": len(selected),
        "offered_req_s": len(selected) / duration_s,
        "offered_req_s_per_gpu": len(selected) / duration_s / tp,
        "request_id_order_sha256": order_hash([request.row_id for request in selected]),
        "arrival_order_sha256": order_hash([f"{request.arrival_s:.12f}" for request in selected]),
        "input_length_order_sha256": order_hash(
            [str(request.prompt_tokens_hint) for request in selected]
        ),
    }


def build_manifest(
    *,
    studies: dict[str, dict[str, Path]],
    private_windows: Path,
    band_stats: dict[str, Any],
    source_trace: Path,
    source_windows: Path,
    source_window_id: str,
) -> dict[str, Any]:
    loaded = {}
    durations = {}
    for role, tiers in studies.items():
        loaded[role] = {}
        for tier, path in tiers.items():
            study = load_study_spec(path)
            window, requests = load_trace_requests(study, study_spec_path=path)
            loaded[role][tier] = requests
            durations[role] = float(window.window_end - window.window_start)

    cells = {}
    all_hashes = []
    for cell, config in CELLS.items():
        tp = int(config["tp"])
        tier = "tp4" if tp == 4 else "primary"
        targets = {}
        for level, multiplier in TARGET_MULTIPLIERS.items():
            target_rate = float(config["frontier_req_s_gpu"]) * multiplier
            target_count = round(target_rate * durations["low1"] * tp)
            roles = [
                role
                for role in ROLES
                if role.startswith(level) or (level == "low" and role == "burnin")
            ]
            selections = {}
            for role in roles:
                anchor, selected = attainable_anchor(loaded[role][tier], target_count)
                record = selected_record(selected, tp=tp, duration_s=durations[role])
                record["anchor"] = anchor
                record["study"] = str(studies[role][tier])
                selections[role] = record
                all_hashes.append(record["request_id_order_sha256"])
            targets[level] = {
                "multiplier": multiplier,
                "target_req_s_per_gpu": target_rate,
                "target_count": target_count,
                "selections": selections,
            }
        cells[cell] = {**config, "targets": targets}

    red_flags = []
    for cell, config in cells.items():
        for level, target in config["targets"].items():
            if not target["selections"]:
                red_flags.append(f"missing_{cell}_{level}")
            for selection in target["selections"].values():
                if selection["selected_count"] <= 0:
                    red_flags.append(f"empty_{cell}_{level}")
    per_cell_distinct = {}
    for cell, config in cells.items():
        hashes = [
            selection["request_id_order_sha256"]
            for target in config["targets"].values()
            for selection in target["selections"].values()
        ]
        per_cell_distinct[cell] = len(hashes) == len(set(hashes))
        if not per_cell_distinct[cell]:
            red_flags.append(f"session_bands_overlap_{cell}")
    return {
        "schema": "fidelity-prefix-pilot-manifest-v1",
        "status": "PASS" if not red_flags else "STOP",
        "source": {
            "windows": str(source_windows),
            "window_id": source_window_id,
            "trace": str(source_trace),
            "trace_sha256": sha256_file(source_trace),
        },
        "private": {
            "windows": str(private_windows),
            "windows_sha256": sha256_file(private_windows),
            "band_stats": band_stats,
            "studies": {
                role: {tier: str(path) for tier, path in tiers.items()}
                for role, tiers in studies.items()
            },
        },
        "roles": list(ROLES),
        "cells": cells,
        "execution": {
            "cutoff_s": 5.0,
            "replicates_per_level": 3,
            "label": "2-of-3 session-disjoint repetitions",
            "even_cell_order": ["low1", "high1", "high2", "low2", "low3", "high3"],
            "odd_cell_order": ["high1", "low1", "low2", "high2", "high3", "low3"],
            "hard_cap_h20_hours": 3.5,
        },
        "sanity": {
            "red_flags": red_flags,
            "n_cells": len(cells),
            "n_roles": len(ROLES),
            "selected_sets": len(all_hashes),
            "distinct_selected_sets": len(set(all_hashes)),
            "per_cell_selected_sets_distinct": per_cell_distinct,
            "invariants": {
                "cells_6": len(cells) == 6,
                "roles_7": len(ROLES) == 7,
                "band_rows_nonzero": all(stats["rows"] > 0 for stats in band_stats.values()),
                "session_bands_disjoint_per_cell": all(per_cell_distinct.values()),
            },
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-windows", type=Path, required=True)
    parser.add_argument("--source-window-id", default="chat_w20260312_1000")
    parser.add_argument("--base-primary-study", type=Path, required=True)
    parser.add_argument("--base-tp4-study", type=Path, required=True)
    parser.add_argument("--private-root", type=Path, required=True)
    parser.add_argument("--public-manifest", type=Path, required=True)
    args = parser.parse_args()

    source_window, source_trace = resolve_source_trace(
        args.source_windows, args.source_window_id
    )
    private_windows, band_stats = materialize_bands(
        source_trace, source_window, args.private_root
    )
    studies = write_studies(
        base_primary=args.base_primary_study,
        base_tp4=args.base_tp4_study,
        private_windows=private_windows,
        private_root=args.private_root,
    )
    manifest = build_manifest(
        studies=studies,
        private_windows=private_windows,
        band_stats=band_stats,
        source_trace=source_trace,
        source_windows=args.source_windows,
        source_window_id=args.source_window_id,
    )
    atomic_json(args.public_manifest, manifest)
    print(json.dumps({
        "status": manifest["status"],
        "manifest": str(args.public_manifest),
        "sanity": manifest["sanity"],
    }, sort_keys=True))
    if manifest["status"] != "PASS":
        raise RuntimeError(f"pilot preflight failed: {manifest['sanity']['red_flags']}")


if __name__ == "__main__":
    main()