#!/usr/bin/env python3
"""Run source first, select the next intervention, then annotate the 2x2 surface."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any, Mapping


HERE = Path(__file__).resolve().parent
ACTION_DIR = HERE.parent / "action-aware-v0"
sys.path.insert(0, str(ACTION_DIR))
sys.path.insert(0, str(HERE))

import pilot_controller as action_controller  # noqa: E402
import prospective_decision  # noqa: E402


SCHEMA = "active-intervention-prospective-state-v0"


def validate_inputs(args: argparse.Namespace, manifest: Mapping[str, Any]) -> None:
    if manifest.get("schema") != "active-intervention-prospective-manifest-v0":
        raise RuntimeError("unexpected active intervention manifest schema")
    if manifest.get("status") != "PASS" or manifest["sanity"]["red_flags"]:
        raise RuntimeError("active intervention manifest did not pass preflight")
    required = {
        "manifest": args.manifest,
        "policy": args.policy,
        "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["trace"]["path"])
    missing = {name: str(path) for name, path in required.items() if not path.exists()}
    if missing:
        raise RuntimeError(f"active intervention input paths missing: {missing}")
    if prospective_decision.sha256_file(args.policy) != manifest["policy"]["sha256"]:
        raise RuntimeError("active intervention policy hash mismatch")


def dry_run(args: argparse.Namespace, manifest: Mapping[str, Any]) -> dict[str, Any]:
    plan = action_controller.dry_run_plan(args, manifest)
    return {
        "schema": "active-intervention-prospective-dry-run-v0",
        "status": "PASS",
        "manifest": str(args.manifest),
        "policy": str(args.policy),
        "source_first": manifest["source_config_id"],
        "post_source_order": "selected by telemetry policy; all remaining cells then annotated",
        "candidate_actions": manifest["actions"],
        "projected_h20_hours": plan["projected_h20_hours"],
        "hard_cap_h20_hours": plan["hard_cap_h20_hours"],
        "sessions": plan["sessions"],
    }


def load_or_build_decision(
    *, args: argparse.Namespace, run_root: Path
) -> dict[str, Any]:
    path = run_root / "active-decision.json"
    if path.exists():
        decision = json.loads(path.read_text(encoding="utf-8"))
        if decision.get("manifest_sha256") != prospective_decision.sha256_file(
            args.manifest
        ):
            raise RuntimeError("existing active decision has a different manifest")
        if decision.get("policy_sha256") != prospective_decision.sha256_file(args.policy):
            raise RuntimeError("existing active decision has a different policy")
        return decision
    decision = prospective_decision.build_decision(
        manifest_path=args.manifest,
        policy_path=args.policy,
        run_root=run_root,
    )
    prospective_decision.atomic_json(path, decision)
    return decision


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser()
    result.add_argument("--manifest", type=Path, required=True)
    result.add_argument("--policy", 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)
    action_controller.configure(args, manifest)
    action_controller.base.MARKER = "active-intervention-prospective-v0"
    if args.dry_run:
        print(json.dumps(dry_run(args, manifest), indent=2, sort_keys=True))
        return

    args.run_root.mkdir(parents=True, exist_ok=True)
    copied_manifest = args.run_root / "prospective-manifest.json"
    if not copied_manifest.exists():
        action_controller.atomic_json(copied_manifest, manifest)
    state_path = args.run_root / "controller-state.json"
    state = action_controller.load_state(
        state_path, float(manifest["budget"]["hard_cap_h20_hours"])
    )
    state["schema"] = SCHEMA
    state["status"] = "running"
    action_controller.atomic_json(state_path, state)

    configs = {str(item["id"]): dict(item) for item in manifest["configs"]}
    config_indexes = {
        str(item["id"]): index for index, item in enumerate(manifest["configs"])
    }
    source_id = str(manifest["source_config_id"])
    action_controller.execute_session(
        args=args,
        manifest=manifest,
        config=configs[source_id],
        index=config_indexes[source_id],
        state=state,
        state_path=state_path,
    )
    decision = load_or_build_decision(args=args, run_root=args.run_root)
    state["active_decision"] = {
        "path": str(args.run_root / "active-decision.json"),
        "status": decision["status"],
        "outcome_only": {
            key: decision["decisions"]["outcome_only"][key]
            for key in ("selected_cutoff_s", "decision_kind", "selected_action")
        },
        "telemetry": {
            key: decision["decisions"]["telemetry"][key]
            for key in ("selected_cutoff_s", "decision_kind", "selected_action")
        },
    }
    action_controller.atomic_json(state_path, state)
    if decision["status"] != "SELECTED":
        state["status"] = decision["status"].lower()
        state["completed_at"] = action_controller.time.time()
        action_controller.atomic_json(state_path, state)
        action_controller.wait_all_idle()
        print(json.dumps({"status": state["status"], "decision": decision["status"]}))
        return

    action_order = decision["decisions"]["telemetry"]["intervention_order"]
    execution_order = [source_id]
    for action_id in action_order:
        target_id = str(manifest["actions"][action_id])
        if target_id not in execution_order:
            execution_order.append(target_id)
    for config_id in configs:
        if config_id not in execution_order:
            execution_order.append(config_id)
    state["execution_order"] = execution_order
    action_controller.atomic_json(state_path, state)
    for config_id in execution_order[1:]:
        action_controller.execute_session(
            args=args,
            manifest=manifest,
            config=configs[config_id],
            index=config_indexes[config_id],
            state=state,
            state_path=state_path,
        )
    state["status"] = "complete"
    state["completed_at"] = action_controller.time.time()
    action_controller.atomic_json(state_path, state)
    action_controller.wait_all_idle()
    print(
        json.dumps(
            {
                "status": state["status"],
                "completed_sessions": state["completed_sessions"],
                "gpu_hours_total": state["gpu_hours_total"],
                "execution_order": execution_order,
            },
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
