"""LM Studio and artifact preflight for frozen E11/E12 sensitivities."""

from __future__ import annotations

from dataclasses import asdict, replace
import json
from pathlib import Path
import time

from preflight_study3 import _tool_probe
from agent_harness.lm_studio import LMStudioClient
from agent_harness.lm_studio_embeddings import LMStudioEmbeddingClient
from agent_harness.lm_studio_management import LMStudioResidencyManager, LMStudioServer
from agent_harness.pilot import research_code_revision
from agent_harness.specs import load_embeddings, load_models, validate_configuration_tree


def run(root: Path) -> dict[str, object]:
    errors, warnings = validate_configuration_tree(root)
    if errors or warnings:
        raise RuntimeError(f"configuration failed: errors={errors}, warnings={warnings}")
    audit = json.loads((root / "docs/STUDY4_ANCILLARY_AUDIT.json").read_text())
    if not audit.get("outcome_independent"):
        raise RuntimeError("ancillary design audit did not pass")
    for experiment_id in ("E11", "E12"):
        raw = root / "results/raw" / experiment_id
        if raw.exists() and any(raw.rglob("*")):
            raise RuntimeError(f"{experiment_id} data exist before preflight")
    models = load_models(root)
    embedding = load_embeddings(root)["EMB002"]
    server = LMStudioServer(port=1234)
    start = server.ensure_running()
    residency = LMStudioResidencyManager(
        models["M002"].base_url, models["M002"].api_token_env, timeout_seconds=1800
    )
    report: dict[str, object] = {
        "schema_version": 1,
        "study": "Study 4 E11/E12 ancillaries",
        "research_code_revision": research_code_revision(root),
        "audit": audit,
        "server_start": start,
        "torch_used": False,
        "started_unix": time.time(),
        "models": {},
    }
    try:
        residency.unload_all()
        transition = residency.ensure_exclusive(
            embedding.model_key, embedding.loaded_context_length
        )
        client = LMStudioEmbeddingClient(embedding, timeout_seconds=1800)
        report["embedding"] = {
            "spec": asdict(embedding),
            "transition": transition.to_dict(),
            "resolved": client.resolve(),
            "probe": client.probe().to_dict(),
            "unload": residency.unload_all().to_dict(),
        }
        model_reports: dict[str, object] = {}
        profiles = [("M002", 16384), ("M002", 65536), ("M003", 65536), ("M004", 65536)]
        for model_id, context in profiles:
            model = replace(models[model_id], context_length=context)
            transition = residency.ensure_exclusive(model.expected_inference_key, context)
            model_client = LMStudioClient(model, timeout_seconds=1800)
            discovery, resolved = model_client.resolve()
            model_reports[f"{model_id}_{context}"] = {
                "spec": asdict(model),
                "transition": transition.to_dict(),
                "resolved": resolved.to_dict(),
                "discovery_errors": discovery.endpoint_errors,
                "tool_probe": _tool_probe(model_client, resolved.inference_key),
                "unload": residency.unload_all().to_dict(),
            }
        report["models"] = model_reports
        report["passed"] = True
        return report
    finally:
        cleanup: list[str] = []
        try:
            report["final_unload"] = residency.unload_all().to_dict()
        except Exception as exc:
            cleanup.append(f"unload_all: {exc}")
        try:
            report["server_stop"] = server.stop()
        except Exception as exc:
            cleanup.append(f"server_stop: {exc}")
        report["cleanup_errors"] = cleanup
        report["finished_unix"] = time.time()


if __name__ == "__main__":
    root = Path(__file__).resolve().parents[1]
    output = root / "results/reports/study4_ancillary_preflight.json"
    output.parent.mkdir(parents=True, exist_ok=True)
    report: dict[str, object] = {}
    try:
        report = run(root)
    except Exception as exc:
        report = {**report, "passed": False, "error": repr(exc)}
        output.write_text(json.dumps(report, indent=2, sort_keys=True, default=str) + "\n")
        raise
    output.write_text(json.dumps(report, indent=2, sort_keys=True, default=str) + "\n")
    print(json.dumps({"passed": True, "report": str(output)}, indent=2))