#!/usr/bin/env python3
"""Freeze the unseen-trace 2x2 active intervention development surface."""

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


SCHEMA = "active-intervention-prospective-manifest-v0"
TP = 4
REPETITIONS = (1, 2, 3)
DURATION_S = 300.0
REPLAY_TIME_SCALE = 0.5
OFFERED_RATE_PER_GPU = 2.75
TARGET_COUNT = round(OFFERED_RATE_PER_GPU * DURATION_S * TP)
WINDOW_ID = "chat_w20260313_1000"
ENGINE_VERSION = "0.24.1.dev3+g668cfb7e2"


def atomic_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    os.replace(temporary, 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 configs() -> list[dict[str, Any]]:
    return [
        {
            "id": "source_mns32_mbbt4096",
            "mns": 32,
            "mbbt": 4096,
            "repetition_order": [1, 2, 3],
        },
        {
            "id": "mns64_mbbt4096",
            "mns": 64,
            "mbbt": 4096,
            "repetition_order": [2, 3, 1],
        },
        {
            "id": "mns32_mbbt8192",
            "mns": 32,
            "mbbt": 8192,
            "repetition_order": [3, 1, 2],
        },
        {
            "id": "joint_mns64_mbbt8192",
            "mns": 64,
            "mbbt": 8192,
            "repetition_order": [1, 3, 2],
        },
    ]


def partition_trace(source: Path, output_root: Path) -> dict[str, Any]:
    source_sha = sha256_file(source)
    output_root.mkdir(parents=True, exist_ok=True)
    paths = {rep: output_root / f"rep{rep}.jsonl" for rep in REPETITIONS}
    temporary = {rep: path.with_suffix(".jsonl.tmp") for rep, path in paths.items()}
    handles = {rep: temporary[rep].open("w", encoding="utf-8") for rep in REPETITIONS}
    counts = {rep: 0 for rep in REPETITIONS}
    id_digests = {rep: hashlib.sha256() for rep in REPETITIONS}
    total = 0
    try:
        with source.open(encoding="utf-8") as input_file:
            for line_number, line in enumerate(input_file, start=1):
                if not line.strip():
                    continue
                row = json.loads(line)
                original_id = str(row.get("request_id") or row.get("id") or line_number)
                digest = hashlib.sha256(
                    f"{source_sha}:{line_number}:{original_id}".encode()
                ).hexdigest()
                repetition = int(digest[:16], 16) % len(REPETITIONS) + 1
                row["request_id"] = f"active-r{repetition}-{digest}"
                handles[repetition].write(json.dumps(row, ensure_ascii=False) + "\n")
                counts[repetition] += 1
                total += 1
                id_digests[repetition].update(row["request_id"].encode() + b"\n")
    finally:
        for handle in handles.values():
            handle.close()
    for repetition in REPETITIONS:
        os.replace(temporary[repetition], paths[repetition])
    partitions = {
        str(rep): {
            "path": str(paths[rep]),
            "rows": counts[rep],
            "bytes": paths[rep].stat().st_size,
            "sha256": sha256_file(paths[rep]),
            "request_id_order_sha256": id_digests[rep].hexdigest(),
        }
        for rep in REPETITIONS
    }
    return {
        "source": str(source),
        "source_sha256": source_sha,
        "source_rows": total,
        "partition_rule": "sha256(source_sha:line_number:original_id) modulo 3",
        "partitions": partitions,
    }


def materialize_study(
    base_study: Path,
    target: Path,
    *,
    repetition: int,
    trace_path: Path,
    windows_path: Path,
) -> None:
    payload = json.loads(base_study.read_text(encoding="utf-8"))
    payload["study_id"] = f"active-intervention-trace13-rep{repetition}"
    payload["hardware"]["host_candidates"] = ["dash0"]
    payload["engine"]["engine_version"] = ENGINE_VERSION
    trace = payload["trace"]
    trace.update(
        {
            "windows_path": str(windows_path),
            "window_id": WINDOW_ID,
            "trace_file_override": str(trace_path),
            "completion_tokens_override": 128,
            "replay_time_scale": REPLAY_TIME_SCALE,
            "early_stop_max_lag_s": None,
            "early_stop_max_elapsed_s": 360.0,
            "restart_engine_after_early_stop": False,
            "adaptive_stop": {"enabled": False},
        }
    )
    atomic_json(target, payload)


def attainable_anchor(requests: list[Any], target_count: int) -> tuple[float, list[Any]]:
    ordered = sorted(float(request.sampling_u) for request in requests)
    if target_count <= 0 or target_count > len(ordered):
        raise ValueError(
            f"target count {target_count} is outside available range 1..{len(ordered)}"
        )
    candidates = []
    for index in sorted({target_count - 1, min(target_count, len(ordered) - 1)}):
        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], item[2])
    )
    return anchor, selected


def selection_record(selected: list[Any]) -> dict[str, Any]:
    return {
        "anchor": max(float(request.sampling_u) for request in selected),
        "selected_count": len(selected),
        "target_count": TARGET_COUNT,
        "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(
    *,
    base_study: Path,
    base_action_manifest: Path,
    source_trace: Path,
    windows_path: Path,
    private_root: Path,
    policy_path: Path,
) -> dict[str, Any]:
    base_manifest = json.loads(base_action_manifest.read_text(encoding="utf-8"))
    if base_manifest.get("status") != "PASS":
        raise ValueError("base action-aware manifest did not pass")
    policy = json.loads(policy_path.read_text(encoding="utf-8"))
    if policy.get("schema") != "active-intervention-policy-v0":
        raise ValueError("unexpected frozen policy schema")
    if policy.get("sanity", {}).get("red_flags"):
        raise ValueError("frozen policy contains red flags")

    partition = partition_trace(source_trace, private_root / "traces")
    repetitions = {}
    selected_sets: list[set[str]] = []
    for repetition in REPETITIONS:
        trace_path = Path(partition["partitions"][str(repetition)]["path"])
        study_path = private_root / "studies" / f"rep{repetition}-tp4.json"
        materialize_study(
            base_study,
            study_path,
            repetition=repetition,
            trace_path=trace_path,
            windows_path=windows_path,
        )
        study = load_study_spec(study_path)
        window, requests = load_trace_requests(study, study_spec_path=study_path)
        duration_s = float(window.window_end - window.window_start)
        if not math.isclose(duration_s, DURATION_S, abs_tol=1e-9):
            raise ValueError(f"rep{repetition}: duration {duration_s} != {DURATION_S}")
        _anchor, selected = attainable_anchor(requests, TARGET_COUNT)
        record = selection_record(selected)
        selected_sets.append({request.row_id for request in selected})
        repetitions[str(repetition)] = {
            "study": str(study_path),
            "study_sha256": sha256_file(study_path),
            "trace": partition["partitions"][str(repetition)],
            "available_filtered_requests": len(requests),
            "selection": record,
        }

    frozen_configs = configs()
    config_ids = {str(config["id"]) for config in frozen_configs}
    invariants = {
        "three_nonempty_trace_partitions": all(
            int(item["rows"]) > 0 for item in partition["partitions"].values()
        ),
        "partition_rows_conserved": sum(
            int(item["rows"]) for item in partition["partitions"].values()
        )
        == int(partition["source_rows"]),
        "selected_sets_disjoint": all(
            not selected_sets[left] & selected_sets[right]
            for left in range(len(selected_sets))
            for right in range(left + 1, len(selected_sets))
        ),
        "target_count_attained": all(
            abs(int(item["selection"]["selected_count"]) - TARGET_COUNT) <= 1
            for item in repetitions.values()
        ),
        "four_unique_configs": len(config_ids) == 4,
        "two_by_two_surface": {
            (int(config["mns"]), int(config["mbbt"]))
            for config in frozen_configs
        }
        == {(32, 4096), (64, 4096), (32, 8192), (64, 8192)},
        "repetition_orders_are_permutations": all(
            sorted(config["repetition_order"]) == list(REPETITIONS)
            for config in frozen_configs
        ),
    }
    red_flags = [name for name, passed in invariants.items() if not passed]
    return {
        "schema": SCHEMA,
        "status": "PASS" if not red_flags else "STOP",
        "source": {
            "window_id": WINDOW_ID,
            "source_trace": str(source_trace),
            "source_trace_sha256": partition["source_sha256"],
            "windows_path": str(windows_path),
            "base_study": str(base_study),
            "base_study_sha256": sha256_file(base_study),
            "base_action_manifest": str(base_action_manifest),
            "base_action_manifest_sha256": sha256_file(base_action_manifest),
        },
        "policy": {
            "path": str(policy_path),
            "sha256": sha256_file(policy_path),
            "status": policy["status"],
            "training": policy["training"],
            "measurement_policy": policy["measurement_policy"],
            "launch_reason": (
                "bounded unseen-trace joint-action test after a negative narrow "
                "retrospective replay"
            ),
        },
        "engine": {
            "tp": TP,
            "duration_s": DURATION_S,
            "client_timeout_s": 450.0,
            "burnin_max_elapsed_s": 90.0,
            "disable_slo_early_stop": True,
        },
        "burnin": base_manifest["burnin"],
        "private": {"trace_partition": partition},
        "repetitions": repetitions,
        "configs": frozen_configs,
        "source_config_id": "source_mns32_mbbt4096",
        "actions": {
            "noop": "source_mns32_mbbt4096",
            "mns": "mns64_mbbt4096",
            "mbbt": "mns32_mbbt8192",
            "joint": "joint_mns64_mbbt8192",
        },
        "checkpoints": {
            "fractions": [0.25, 0.50, 0.75, 1.0],
            "seconds": [75.0, 150.0, 225.0, 300.0],
        },
        "gates": {
            "acceptable_regret": 0.02,
            "source_ceiling_normalized_goodput": 0.98,
            "confirmation_trigger_gpu_cost_reduction": 0.10,
            "contribution_gpu_cost_reduction": 0.20,
            "maximum_task_regret": 0.05,
        },
        "budget": {
            "hard_cap_h20_hours": 6.0,
            "session_estimate_h20_hours": 1.3,
            "safety_h20_hours": 0.3,
            "expected_h20_hours": [4.6, 5.5],
            "expected_wall_minutes": [75, 100],
        },
        "sanity": {"invariants": invariants, "red_flags": red_flags},
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-study", type=Path, required=True)
    parser.add_argument("--base-action-manifest", type=Path, required=True)
    parser.add_argument("--source-trace", type=Path, required=True)
    parser.add_argument("--windows-path", type=Path, required=True)
    parser.add_argument("--private-root", type=Path, required=True)
    parser.add_argument("--policy", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    payload = build(
        base_study=args.base_study,
        base_action_manifest=args.base_action_manifest,
        source_trace=args.source_trace,
        windows_path=args.windows_path,
        private_root=args.private_root,
        policy_path=args.policy,
    )
    atomic_json(args.output, payload)
    print(json.dumps({"status": payload["status"], "sanity": payload["sanity"]}))
    if payload["status"] != "PASS":
        raise SystemExit("prospective manifest preflight failed")


if __name__ == "__main__":
    main()