#!/usr/bin/env python3
"""Format probe — tool-calling / structured AST. Instead of writing surface
syntax (where bracket-tracking is the wall), Qwen submits the program as the
canonical AILang AST JSON via a `submit_program` tool. Hypothesis: the model
produces structured JSON through the tool channel (which it is trained hard on,
and which the API validates) more reliably than free-text syntax. No converter
risk: the tool argument IS the .ail.json, written and `ail check`ed directly.

Same ablation ladder. Reads $IONOS_API_TOKEN (never printed). Live IONOS.
"""
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-tool-8k.md"

def ail_ast(src):
    p=Path("/tmp/abl/_fs.ail"); p.write_text(src+"\n")
    return subprocess.run([AIL,"parse",str(p)],capture_output=True,text=True).stdout.strip()

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"))))))))"""

TOOL={"type":"function","function":{
  "name":"submit_program",
  "description":"Submit a complete AILang program as its canonical AST JSON.",
  "parameters":{"type":"object","properties":{
    "module":{"type":"object","description":"The full module AST: an object with keys schema, name, imports, defs."}},
    "required":["module"]}}}

def sys_base(extra=""):
    return ("You author AILang programs as a canonical Abstract Syntax Tree in JSON, "
      "submitted with the submit_program tool (the `module` argument is the full AST: "
      "an object with keys schema, name, imports, defs). The AST node forms are shown "
      "by these complete, valid example programs.\n\n"
      "Example 1 AST:\n"+ail_ast(A)+"\n\nExample 2 AST:\n"+ail_ast(B)+"\n"+extra+
      "\nAlways call submit_program with the full module AST.")
SYS=sys_base()
SYS_SERIES=sys_base("\nExample 3 AST (uses the Series library — `new` term with type-arg "
      "Float and capacity; Series.push/Series.at/Series.len/Series.total_count called via "
      "app on a `var` named e.g. \"Series.push\"):\n"+ail_ast(S)+"\n")

SMA_OUT="3.0\n5.33333\n5.66667\n5.33333\n"
SMA_TASK=("Write a program named 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 stdout: "+repr(SMA_OUT))

STAGES=[
 ("L0_const",SYS,"Write a program named m with a function answer taking no parameters that returns the Int 42.",None),
 ("L1_arith",SYS,"Write a program named m with a function triple taking one Int parameter, returning it multiplied by 3.",None),
 ("L2_print",SYS,"Write a program named m whose main prints the Int 7 on its own line.","7\n"),
 ("L3_adt",SYS,"Write a program named 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",SYS,"Write a program named 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",SYS_SERIES,"Write a program named 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",SYS_SERIES,"Write a program named 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",SYS_SERIES,SMA_TASK,SMA_OUT),
]

def call(system,task):
    body=json.dumps({"model":MODEL,"temperature":0.2,"max_tokens":8000,
        "messages":[{"role":"system","content":system},{"role":"user","content":task}],
        "tools":[TOOL],"tool_choice":{"type":"function","function":{"name":"submit_program"}}}).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"], d.get("usage",{})

def evaluate(msg,expected):
    tc=msg.get("tool_calls")
    if not tc: return "no-tool-call",f"model returned content instead: {str(msg.get('content'))[:200]}",""
    try: args=json.loads(tc[0]["function"]["arguments"])
    except Exception as e: return "bad-json",f"tool arguments not valid JSON: {e}",tc[0]["function"]["arguments"][:400]
    module=args.get("module",args)
    if not isinstance(module,dict): return "bad-shape",f"module not an object: {type(module).__name__}",json.dumps(args)[:400]
    module.setdefault("schema","ailang/v0"); module.setdefault("imports",[])
    name=module.get("name","m")
    f=Path(f"/tmp/abl/{name}.ail.json"); f.parent.mkdir(exist_ok=True)
    f.write_text(json.dumps(module,indent=2)+"\n")
    pretty=json.dumps(module,indent=1)
    chk=subprocess.run([AIL,"check",str(f)],capture_output=True,text=True)
    if chk.returncode!=0: return "check",(chk.stderr or chk.stdout).strip(),pretty
    if expected is None: return "green","(check-only)",pretty
    run=subprocess.run([AIL,"run",str(f)],capture_output=True,text=True)
    if run.returncode!=0: return "run",(run.stderr or run.stdout).strip(),pretty
    if run.stdout!=expected: return "output",f"got {run.stdout!r}, expected {expected!r}",pretty
    return "green",run.stdout,pretty

def main():
    log=["# Qwen format probe — tool-calling / structured AST\n",
         f"**Model:** {MODEL}; submit_program tool, forced; AST written directly to .ail.json.\n","---\n"]
    summary=[]
    for sid,sysmsg,task,exp in STAGES:
        msg,usage=call(sysmsg,task)
        stage,fb,pretty=evaluate(msg,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",
              "submitted AST:\n```json\n"+pretty[:1400]+"\n```\n",
              ("" 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()