#!/usr/bin/env python3
"""Constraint probe for the ADAPTIVE-MEMORY research (docs/research/14-adaptive-memory.md).

Question it answers empirically (mandate: «проверяй заблокированность методов пробой,
а не по памяти» — ADAPTIVE_MEMORY_RESEARCH_PROMPT §Жёсткий контур п.1): what LOGIT-LEVEL
control do our closed flagship providers actually expose in 2026? This bounds which
"learns-during-the-book" mechanisms are runnable on the flagship TRANSLATOR:

  logit-level lever            needs from the API                available on flagship?
  ------------------------------------------------------------------------------------
  kNN-MT (Khandelwal 2021)     DECODER HIDDEN STATES per step    NO API exposes this
  constrained/beam decoding    beam access / token masking       logit_bias at best
  output-distribution reading  per-token logprobs                <- THIS we probe live
  hard steering                logit_bias honored                <- THIS we probe live

So even the BEST case (logprobs + logit_bias both honored) does NOT unblock kNN-MT: the
datastore key in kNN-MT is the decoder's context representation, which none of these APIs
return. This probe measures the ceiling (logprobs/logit_bias), and the ceiling is still
below what kNN-MT needs — a doubly-closed door. Local decoders (llama.cpp/vLLM) DO expose
logprobs/logits by design (n_probs / prompt_logprobs), so local kNN-MT is physically
possible; that half is asserted from the runtimes' documented APIs, not reprobed here (the
stand GPU is held by the polygon's resident model — do not evict, exp03/§coordination).

Cheap by design: ~1 short call per provider. Reuses eval/.env + eval/providers.json quirks.
Usage: eval/.venv/bin/python eval/adaptive_probe.py
Output: stdout table + eval/data/adaptive_probe.json
"""
from __future__ import annotations

import json
import os
import urllib.error
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent
OUT = ROOT / "data" / "adaptive_probe.json"


def load_env() -> None:
    for cand in (ROOT / ".env", Path("~/.textmachine.env").expanduser()):
        if not cand.exists():
            continue
        for line in cand.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, _, v = line.partition("=")
            if v.strip():
                os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))
        return


load_env()

# The models we actually probe. Includes the REAL translator/editor stack models
# (deepseek-v4-flash draft, grok-4.3 editor, gemini-3.1-pro premium) plus a clean
# non-reasoning path per family to give logprobs its best chance to appear.
PROBES = [
    {"name": "deepseek-v4-flash", "base": "https://api.deepseek.com/v1",
     "model": "deepseek-v4-flash", "key": "DEEPSEEK_API_KEY", "max": 200},
    {"name": "grok-non-reasoning", "base": "https://api.x.ai/v1",
     "model": "grok-4.20-0309-non-reasoning", "key": "XAI_API_KEY", "max": 32},
    {"name": "grok-4.3(editor)", "base": "https://api.x.ai/v1",
     "model": "grok-4.3", "key": "XAI_API_KEY", "max": 200},
    {"name": "glm-4.5-air", "base": "https://api.z.ai/api/paas/v4",
     "model": "glm-4.5-air", "key": "ZAI_API_KEY", "max": 64,
     "extra": {"thinking": {"type": "disabled"}}},
    {"name": "gemini-2.5-flash", "base": "https://generativelanguage.googleapis.com/v1beta/openai",
     "model": "gemini-2.5-flash", "key": "GEMINI_API_KEY", "max": 64,
     "extra": {"extra_body": {"google": {"thinking_config": {"thinking_budget": 0}}}}},
    {"name": "gemini-3.1-pro(premium)", "base": "https://generativelanguage.googleapis.com/v1beta/openai",
     "model": "gemini-3.1-pro-preview", "key": "GEMINI_API_KEY", "max": 8000},
    {"name": "gpt-5-mini", "base": "https://api.openai.com/v1",
     "model": "gpt-5-mini", "key": "OPENAI_API_KEY", "max": 8000, "reasoning": True},
]


def one_call(p: dict, with_logprobs: bool, with_logit_bias: bool):
    """Return dict describing what the provider did with logprobs / logit_bias."""
    key = os.environ.get(p["key"], "")
    if not key:
        return {"skip": f"no key {p['key']}"}
    body = {
        "model": p["model"],
        "messages": [{"role": "user", "content": "Reply with exactly the single word: ok"}],
        "max_tokens": p["max"],
    }
    if p.get("reasoning"):
        body["max_completion_tokens"] = body.pop("max_tokens")
        body["reasoning_effort"] = "minimal"
    else:
        body["temperature"] = 0.0
    body.update(p.get("extra", {}))
    if with_logprobs:
        body["logprobs"] = True
        body["top_logprobs"] = 5
    if with_logit_bias:
        # bias an arbitrary token id hard; we only check acceptance, not effect.
        body["logit_bias"] = {"1734": -100}
    req = urllib.request.Request(
        p["base"].rstrip("/") + "/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"},
    )
    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            data = json.loads(resp.read())
    except urllib.error.HTTPError as e:
        return {"http_error": e.code, "detail": e.read().decode(errors="replace")[:300].replace("\n", " ")}
    except Exception as e:
        return {"transport_error": f"{type(e).__name__}: {e}"}
    ch = (data.get("choices") or [{}])[0]
    lp = ch.get("logprobs")
    # logprobs "present" = non-null object carrying per-token content entries
    has_lp = bool(lp) and bool((lp or {}).get("content"))
    sample = None
    if has_lp:
        first = lp["content"][0]
        sample = {"token": first.get("token"),
                  "logprob": round(first.get("logprob", 0.0), 3),
                  "n_top": len(first.get("top_logprobs", []))}
    return {
        "http": 200,
        "content": (ch.get("message", {}).get("content") or "")[:40],
        "finish": ch.get("finish_reason"),
        "logprobs_present": has_lp,
        "logprobs_sample": sample,
        "usage": {k: data.get("usage", {}).get(k) for k in ("prompt_tokens", "completion_tokens")},
    }


def main():
    results = []
    print(f"{'provider':<26} {'logprobs?':<11} {'logit_bias':<12} sample / note")
    print("-" * 88)
    for p in PROBES:
        r_lp = one_call(p, with_logprobs=True, with_logit_bias=False)
        r_lb = one_call(p, with_logprobs=False, with_logit_bias=True)
        rec = {"provider": p["name"], "model": p["model"], "logprobs": r_lp, "logit_bias": r_lb}
        results.append(rec)

        if "skip" in r_lp:
            print(f"{p['name']:<26} {'SKIP':<11} {'':<12} {r_lp['skip']}")
            continue
        if r_lp.get("http") == 200:
            lp_verdict = "YES" if r_lp["logprobs_present"] else "no(200)"
        elif "http_error" in r_lp:
            lp_verdict = f"HTTP{r_lp['http_error']}"
        else:
            lp_verdict = "ERR"
        # logit_bias: 200 = accepted (not rejected); HTTP4xx = rejected
        if r_lb.get("http") == 200:
            lb_verdict = "accepted"
        elif "http_error" in r_lb:
            lb_verdict = f"HTTP{r_lb['http_error']}"
        else:
            lb_verdict = "ERR"
        note = ""
        if r_lp.get("logprobs_sample"):
            s = r_lp["logprobs_sample"]
            note = f"tok={s['token']!r} lp={s['logprob']} top{s['n_top']}"
        elif "detail" in r_lp:
            note = r_lp["detail"][:60]
        elif r_lp.get("http") == 200 and not r_lp["logprobs_present"]:
            note = f"finish={r_lp.get('finish')} content={r_lp.get('content')!r}"
        print(f"{p['name']:<26} {lp_verdict:<11} {lb_verdict:<12} {note}")

    OUT.write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\nsaved -> {OUT}")
    print("\nReading: 'logprobs YES' = output distribution readable (still NOT decoder hidden\n"
          "states → kNN-MT still blocked). 'logit_bias accepted' = crude token steering only.")


if __name__ == "__main__":
    main()
