#!/usr/bin/env python3
"""Format probe — M-expressions. McCarthy's original Lisp notation:
`head[arg1; arg2; ...]` — explicit structure (brackets + semicolons) but in
the function-call shape `f[x; y]` that LLMs know from C/Python/JS. Tests the
"more explicit structure helps" finding in a familiar guise.

Own literal parser (atoms verbatim, whitespace-agnostic), round-trip-verified
before judging output. Same ablation ladder.

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="Qwen/Qwen3-Coder-Next"
REPO=Path(__file__).resolve().parents[2]
AIL=os.environ.get("AIL_BIN", str(REPO/"target/debug/ail"))
DOC=REPO/"experiments/2026-05-12-cross-model-authoring/qwen-mexpr.md"

def tokenize(s):
    t=[];i=0
    while i<len(s):
        c=s[i]
        if c.isspace(): i+=1; continue
        if c in '()': t.append(c); i+=1; continue
        if c=='"':
            j=i+1
            while j<len(s):
                if s[j]=='\\': j+=2; continue
                if s[j]=='"': break
                j+=1
            t.append(s[i:j+1]); i=j+1; continue
        j=i
        while j<len(s) and not s[j].isspace() and s[j] not in '()"': j+=1
        t.append(s[i:j]); i=j
    return t
def parse_s(toks):
    t=toks.pop(0)
    if t=='(':
        n=[]
        while toks and toks[0]!=')': n.append(parse_s(toks))
        if not toks: raise ValueError("unbalanced")
        toks.pop(0); return n
    return t
def to_tree(s): return parse_s(tokenize(s))
def emit_sexpr(n): return '('+' '.join(emit_sexpr(x) for x in n)+')' if isinstance(n,list) else n

def emit_mexpr(node, ind=0):
    if not isinstance(node,list): return node
    head=node[0]; args=node[1:]
    if not args: return f"{head}[]"
    if all(not isinstance(a,list) for a in args):
        return f"{head}[" + "; ".join(args) + "]"
    sp="  "*(ind+1)
    inner=(";\n").join(sp+emit_mexpr(a,ind+1) for a in args)
    return f"{head}[\n{inner}\n{'  '*ind}]"

def mtokenize(s):
    t=[];i=0
    while i<len(s):
        c=s[i]
        if c.isspace(): i+=1; continue
        if c in '[];': t.append(c); i+=1; continue
        if c=='"':
            j=i+1
            while j<len(s):
                if s[j]=='\\': j+=2; continue
                if s[j]=='"': break
                j+=1
            t.append(s[i:j+1]); i=j+1; continue
        j=i
        while j<len(s) and not s[j].isspace() and s[j] not in '[];"': j+=1
        t.append(s[i:j]); i=j
    return t
def parse_m(toks):
    if not toks: raise ValueError("unexpected end")
    atom=toks.pop(0)
    if atom in '[];': raise ValueError(f"expected atom, got {atom!r}")
    if toks and toks[0]=='[':
        toks.pop(0); node=[atom]
        if toks and toks[0]==']': toks.pop(0); return node
        while True:
            node.append(parse_m(toks))
            if not toks: raise ValueError("unterminated [")
            sep=toks.pop(0)
            if sep==']': break
            if sep!=';': raise ValueError(f"expected ; or ] got {sep!r}")
            if toks and toks[0]==']':   # tolerate a trailing ';' before ']'
                toks.pop(0); break
        return node
    return atom
def from_mexpr(text): return emit_sexpr(parse_m(mtokenize(text)))

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")))))"""
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"))))))"""
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=("This program is written in M-expression notation. A node is "
       "`head[arg1; arg2; ...]` — the operation name, then SQUARE brackets "
       "containing its arguments separated by semicolons. An argument is either "
       "an atom written literally (a number, a name, an operator like + or * or "
       "/, or a quoted string like \"\\n\") or another `head[...]` node. A node "
       "with no arguments is `head[]`. Whitespace and line breaks inside the "
       "brackets are ignored. Study these complete, valid examples and write the "
       "requested program in the SAME M-expression notation. Return ONLY the "
       "M-expression — no prose, no code fences.\n\n")
BASE=INTRO+"Example 1:\n"+emit_mexpr(to_tree(A))+"\n\nExample 2:\n"+emit_mexpr(to_tree(B))+"\n"
SERIES=(BASE+"\nExample 3 (Series ring-buffer library: `new[Series; con[Float]; N]` "
        "builds a series of capacity N; `Series.push` takes an owned series and an "
        "element, returns the new series; `Series.at` takes a series and an Int "
        "index (0=newest); `Series.len`/`Series.total_count` take a series, return "
        "Int):\n"+emit_mexpr(to_tree(S))+"\n")

SMA_OUT="3.0\n5.33333\n5.66667\n5.33333\n"
SMA_TASK=("Write a program `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 divided by 3.0) on its own line. Expected:\n"+SMA_OUT)

STAGES=[
 ("L0_const",BASE,"Write a program `m` with a function `answer` taking no parameters that returns the Int 42.",None),
 ("L1_arith",BASE,"Write a program `m` with a function `triple` taking one Int parameter, returning it multiplied by 3.",None),
 ("L2_print",BASE,"Write a program `m` whose `main` prints the Int 7 on its own line.","7\n"),
 ("L3_adt",BASE,"Write a program `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 program `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 program `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 program `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":2000,
        "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=120) 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)
    body=(m.group(1) if m else resp)
    i=body.find("module[")
    return body[i:].strip() if i>=0 else body.strip()

def evaluate(mx,expected):
    try: plain=from_mexpr(mx)
    except Exception as e: return "convert",f"own parser rejected output: {e}",""
    m=re.search(r"\(module\s+([A-Za-z_][A-Za-z0-9_]*)",plain)
    name=m.group(1) if m else "m"
    f=Path(f"/tmp/abl/{name}.ail"); f.parent.mkdir(exist_ok=True); f.write_text(plain+"\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(),plain
    if expected is None: return "green","(check-only)",plain
    run=subprocess.run([AIL,"run",str(f)],capture_output=True,text=True)
    if run.returncode!=0: return "run",(run.stderr or run.stdout).strip(),plain
    if run.stdout!=expected: return "output",f"got {run.stdout!r}, expected {expected!r}",plain
    return "green",run.stdout,plain

def main():
    log=["# Qwen format probe — M-expressions `head[a; b]`\n",
         f"**Model:** {MODEL}; single-shot; own literal parser (atoms verbatim).\n","---\n"]
    summary=[]
    for sid,ctx,task,exp in STAGES:
        resp,usage=call(ctx,task); mx=extract(resp)
        stage,fb,plain=evaluate(mx,exp)
        ok=stage=="green"; mark="✅ green" if ok else f"❌ {stage}"
        summary.append((sid,mark)); print(f"{sid}: {mark}")
        log+=[f"## {sid} — {mark}",
              f"\ntokens: prompt={usage.get('prompt_tokens','?')} completion={usage.get('completion_tokens','?')}\n",
              "model output:\n```\n"+mx[:1400]+"\n```\n",
              ("converted Form-A:\n```\n"+plain[:1100]+"\n```\n" if plain else ""),
              ("" if ok else f"**{stage}:**\n```\n{fb[:600]}\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()