"""Probe candidate checkpoints: load, record provenance, verify generation and context-length support, write evidence/models.json."""
from __future__ import annotations
import json, yaml, time
import mlx.core as mx
from .backends.mlx_backend import MLXModel, memory_snapshot, reset_peak
from .cache.sparse import cache_tensor_bytes
from .util import write_json, utc_now

def probe(config_path: str) -> int:
    cfg = yaml.safe_load(open(config_path))
    out = {"timestamp_utc": utc_now(), "models": {}}
    for alias, spec in cfg["models"].items():
        reset_peak()
        m = MLXModel(spec["repo"])
        info = m.info.to_dict()
        ids = m.chat_prompt_tokens([{"role": "user", "content": cfg["probe"]["prompt"] + "\n" + cfg["probe"]["question"]}])
        g = m.generate_text(ids, max_tokens=12)
        info["probe_answer"] = g["text"]; info["probe_ok"] = "ZR-77Q-0931" in g["text"]
        info["load_ms"] = m.load_ms
        lens = {}
        for L in cfg["probe"]["context_lengths"]:
            if m.info.max_position_embeddings and L > m.info.max_position_embeddings:
                lens[str(L)] = {"supported": False, "reason": "exceeds max_position_embeddings"}; continue
            toks = [m.tokenizer.encode("the", add_special_tokens=False)[0]] * L
            reset_peak()
            try:
                c, t = m.prefill(toks)
                lens[str(L)] = {"supported": True, "prefill_ms": t["prefill_ms"], "cache_bytes": cache_tensor_bytes(c), "peak_bytes": int(mx.get_peak_memory())}
            except Exception as e:
                lens[str(L)] = {"supported": False, "reason": f"{type(e).__name__}: {str(e)[:120]}"}
            del c
        info["context_lengths"] = lens
        info["memory_after_probe"] = memory_snapshot()
        out["models"][alias] = info
        print(alias, spec["repo"], "probe_ok", info["probe_ok"], {k: (round(v.get("prefill_ms", 0)), v.get("cache_bytes")) for k, v in lens.items()})
        del m
        mx.clear_cache()
    write_json("evidence/models.json", out)
    return 0 if all(v["probe_ok"] for v in out["models"].values()) else 1