"""The Onyx-shaped CALL for the COBOL SML component — "how we call the model is
the kicker" (Luke, 2026-09-03).

v1 alone (one greedy shot) = ~6% after repair. This layer is the system on top:
  * BRANCHING  — top-N first-token candidates, each continued greedily (the
    TinkyBrain tiny-model diversity trick: argmax per branch, not sampling)
  * bounded_retry — symbolic repair on each candidate
  * verification — cobc -free -c; keep the FIRST candidate that compiles

Same weights, better call. Returns the first compiler-verified program, or the
best-effort draft with compiles=False (honest).
"""
import os, subprocess, tempfile
import mlx.core as mx
from cobol_tokenizer import PAD, BOS, EOS, SEP
from cobol_repair import repair


def _next_logits(model, toks):
    return model(mx.array([toks]))[0, -1, :]


def _continue_greedy(model, toks, max_tokens):
    for _ in range(max_tokens):
        nt = mx.argmax(_next_logits(model, toks)).item()
        if nt in (PAD, EOS):
            break
        toks.append(nt)
    return toks


def branched_drafts(model, tok, prompt, branches=6, max_tokens=400):
    """Top-`branches` first-token candidates, each greedily continued."""
    base = [BOS] + tok.encode(prompt) + [SEP]
    logits = _next_logits(model, base)
    order = mx.argsort(logits)  # ascending; take the tail for top-k
    out = []
    for i in range(branches * 2):  # over-fetch; skip control tokens
        ft = int(order[-1 - i].item())
        if ft in (PAD, EOS, SEP):
            continue
        toks = _continue_greedy(model, base + [ft], max_tokens)
        out.append(tok.decode(toks[toks.index(SEP) + 1:]))
        if len(out) >= branches:
            break
    return out


def compiles(src):
    with tempfile.TemporaryDirectory() as d:
        f = os.path.join(d, "p.cob"); open(f, "w").write(src + "\n")
        r = subprocess.run(["cobc", "-free", "-c", "-o", os.path.join(d, "p.o"), f],
                           capture_output=True, text=True)
        return r.returncode == 0, (r.stderr or r.stdout).strip()


def system_call(model, tok, prompt, branches=6):
    """Onyx-shaped: branch -> repair -> verify; first verified wins."""
    tried = 0
    first_draft = None
    for draft in branched_drafts(model, tok, prompt, branches):
        if first_draft is None:
            first_draft = draft
        for cand, path in ((draft, f"v1-branch{tried}"),
                           (repair(draft), f"v1-branch{tried}+repair")):
            ok, _err = compiles(cand)
            if ok:
                return {"cobol": cand, "compiles": True, "path": path,
                        "branches_tried": tried + 1, "verifier": "cobc -free -c"}
        tried += 1
    # nothing verified — honest best-effort
    best = repair(first_draft) if first_draft else ""
    return {"cobol": best, "compiles": False, "path": "unresolved",
            "branches_tried": tried, "verifier": "cobc -free -c"}