#!/usr/bin/env python3
"""Ablation probe: WHY does Qwen3-Coder-Next fail at AILang Form-A?

Try-and-error from a minimum upward. A ladder of tasks of increasing
complexity (L0 trivial → L7 full SMA), each run INDEPENDENTLY: fresh,
minimal few-shot context, single shot (no retry, no dialog accumulation —
prior probes showed retries don't converge and long context degenerates).
The first level that breaks localises the wall.

Reads $IONOS_API_TOKEN (never printed). Live IONOS — run with consent.
"""
from __future__ import annotations
import json, os, re, subprocess, urllib.request
from pathlib import Path

ENDPOINT = "https://openai.inference.de-txl.ionos.com/v1/chat/completions"
MODEL    = os.environ.get("PROBE_MODEL","Qwen/Qwen3-Coder-Next")
REPO = Path(__file__).resolve().parents[2]
AIL  = os.environ.get("AIL_BIN", str(REPO / "target/debug/ail"))
DOC  = REPO / f"experiments/2026-05-12-cross-model-authoring/{os.environ.get('PROBE_TAG','qwen')}-ablation-ctor.md"

DEMO_A = """(module demo_a
  (fn double
    (type (fn-type (params (own (con Int))) (ret (own (con Int)))))
    (params x)
    (body (app + x x)))
  (fn main
    (type (fn-type (params) (ret (own (con Unit))) (effects IO)))
    (params)
    (body (seq (app print (app double 21)) (do io/print_str "\\n")))))"""

DEMO_B = """(module demo_b
  (data IntList (ctor Nil) (ctor Cons (con Int) (con IntList)))
  (fn sum
    (type (fn-type (params (own (con IntList))) (ret (own (con Int)))))
    (params xs)
    (body
      (match xs
        (case (pat-ctor Nil) 0)
        (case (pat-ctor Cons h t) (app + h (app sum t))))))
  (fn main
    (type (fn-type (params) (ret (own (con Unit))) (effects IO)))
    (params)
    (body
      (let xs (term-ctor IntList Cons 1 (term-ctor IntList Cons 2 (term-ctor IntList Nil)))
        (seq (app print (app sum xs)) (do io/print_str "\\n"))))))"""

DEMO_S = """(module demo_s
  (fn main
    (type (fn-type (params) (ret (own (con Unit))) (effects IO)))
    (params)
    (body
      (let s (new Series (con Float) 3)
        (let s2 (app Series.push s 1.0)
          (let s3 (app Series.push s2 2.0)
            (seq (app print (app Series.len s3)) (do io/print_str "\\n"))))))))"""

INTRO = ("AILang is a fully-parenthesised S-expression language (\"Form-A\"). "
         "Study these complete, valid example modules, then write the requested "
         "module. Return ONLY the AILang source for the module — no prose, no "
         "markdown fences.\n\n")

BASE   = INTRO + "Example 1:\n" + DEMO_A + "\n\nExample 2:\n" + DEMO_B + "\n"
SERIES = (BASE + "\nExample 3 (uses the Series ring-buffer library — "
          "`(new Series (con Float) N)` allocates one of capacity N; "
          "`Series.push : (own (Series a)) (own a) -> (own (Series a))` appends "
          "and returns the new series; `Series.at : (borrow (Series a)) (own Int) "
          "-> (own a)` reads index i back from newest (0 = newest); "
          "`Series.len`/`Series.total_count : (borrow (Series a)) -> (own Int)`):\n"
          + DEMO_S + "\n")

SMA_OUT = "3.0\n5.33333\n5.66667\n5.33333\n"
SMA_TASK = ("Write a module `m` computing a simple moving average, window 3, over "
            "the float stream 1.0, 5.0, 3.0, 8.0, 6.0, 2.0 using the Series library: "
            "create a Float Series lookback 3, push each value in turn, and after each "
            "push once at least 3 values have been pushed, print the average of the "
            "most recent 3 (their sum / 3.0) on its own line. Expected:\n" + SMA_OUT)

# (id, context, task, expected_stdout_or_None[=check-only])
STAGES = [
 ("L0_const",  BASE,   "Write a module `m` with a function `answer` taking no parameters that returns the Int 42.", None),
 ("L1_arith",  BASE,   "Write a module `m` with a function `triple` taking one Int parameter, returning it multiplied by 3.", None),
 ("L2_print",  BASE,   "Write a module `m` whose `main` prints the Int 7 on its own line.", "7\n"),
 ("L3_adt",    BASE,   "Write a module `m` with an ADT `Light` (ctors Red, Green) and a function `code` returning Int 0 for Red and 1 for Green via match.", None),
 ("L4_rec",    BASE,   "Write a module `m` with an IntList ADT (Nil/Cons), a function `length` counting elements by recursion, and a `main` that prints the length of the list [1,2,3] on its own line.", "3\n"),
 ("L5_series_len", SERIES, "Write a module `m` using the Series library: create a Float Series lookback 3, push 1.0 then 2.0 then 3.0, then print its Series.len on its own line.", "3\n"),
 ("L6_series_at",  SERIES, "Write a module `m` using the Series library: create a Float Series lookback 3, push 1.0 then 2.0 then 3.0, then print (Series.at on index 0, the newest) on its own line.", "3.0\n"),
 ("L7_sma",        SERIES, SMA_TASK, SMA_OUT),
]

def call(system, task):
    body = json.dumps({"model": MODEL, "temperature": 0.2, "max_tokens": 1000, "stop": ["assistant"],
        "messages": [{"role":"system","content":system},{"role":"user","content":task}]}).encode()
    req = urllib.request.Request(ENDPOINT, data=body, headers={
        "Authorization": f"Bearer {os.environ['IONOS_API_TOKEN']}", "Content-Type":"application/json"})
    with urllib.request.urlopen(req, timeout=300) as r:
        d = json.loads(r.read())
    return d["choices"][0]["message"]["content"], d.get("usage", {})

def extract(resp):
    m=re.search(r"```(?:[a-zA-Z]*)?\n(.*?)```",resp,re.S)
    code=m.group(1) if m else resp
    i=code.find("(#0")
    if i<0: i=code.find("(module")
    if i<0: return code.strip()
    depth=0
    for j in range(i,len(code)):
        if code[j]=="(": depth+=1
        elif code[j]==")":
            depth-=1
            if depth==0: return code[i:j+1]
    return code[i:].strip()

def evaluate(code, expected):
    m = re.search(r"\(module\s+([A-Za-z_][A-Za-z0-9_]*)", code)
    name = m.group(1) if m else "m"
    f = Path(f"/tmp/abl/{name}.ail"); f.parent.mkdir(exist_ok=True); f.write_text(code+"\n")
    chk = subprocess.run([AIL,"check",str(f)], capture_output=True, text=True)
    if chk.returncode != 0: return "check", (chk.stderr or chk.stdout).strip()
    if expected is None: return "green", "(check-only)"
    run = subprocess.run([AIL,"run",str(f)], capture_output=True, text=True)
    if run.returncode != 0: return "run", (run.stderr or run.stdout).strip()
    if run.stdout != expected: return "output", f"got {run.stdout!r}, expected {expected!r}"
    return "green", run.stdout

def main():
    log = ["# Qwen3-Coder-Next ablation — where does Form-A authoring break?\n",
           f"**Model:** {MODEL}; single-shot per level; minimal few-shot context.\n",
           "Each level is independent (fresh context, one call, no retry).\n", "---\n"]
    summary = []
    for sid, ctx, task, exp in STAGES:
        try:
            resp, usage = call(ctx, task)
            code = extract(resp)
            stage, fb = evaluate(code, exp)
        except Exception as e:
            usage={}; code=""; stage="api-error"; fb=str(e)[:200]
        ok = stage == "green"
        mark = "✅ green" if ok else f"❌ {stage}"
        summary.append((sid, mark))
        print(f"{sid}: {mark}")
        log += [f"## {sid} — {mark}",
                f"\nctx={'BASE' if ctx is BASE else 'SERIES'}  tokens: "
                f"prompt={usage.get('prompt_tokens','?')} completion={usage.get('completion_tokens','?')}\n",
                "```\n"+code+"\n```\n",
                ("" if ok else f"**`ail {stage}`:**\n```\n{fb[:800]}\n```\n")]
    log += ["\n---\n## Ladder\n"] + [f"- {s}: {m}" for s,m in summary]
    DOC.write_text("\n".join(log))
    print("doc →", DOC)

if __name__ == "__main__":
    main()
