#!/usr/bin/env python3
"""Single-task authoring probe: can a foreign LLM (Qwen3-Coder-Next via
IONOS) write the Series SMA worked example in AILang Form-A, given only the
Form-A mini-spec and the Series library API?

Multi-turn: each turn, the model's program is run through `ail check` →
`ail run`; on failure the diagnostic is fed back and the model retries, up
to MAX_TURNS. Every turn is logged to a markdown progress doc.

Reads the token from $IONOS_API_TOKEN (caller sources it from ~/.ionos_token;
the value is never printed). Live IONOS calls — run only with consent.
"""
from __future__ import annotations
import json, os, re, subprocess, sys, urllib.request, urllib.error
from pathlib import Path

ENDPOINT = "https://openai.inference.de-txl.ionos.com/v1/chat/completions"
MODEL    = "Qwen/Qwen3-Coder-Next"
MAX_TURNS = 6
REPO = Path(__file__).resolve().parents[2]
AIL  = os.environ.get("AIL_BIN", str(REPO / "target/debug/ail"))
SPEC = (REPO / "experiments/2026-05-12-cross-model-authoring/rendered/ail.md").read_text()
EXPECTED = "3.0\n5.33333\n5.66667\n5.33333\n"
DOC = REPO / "experiments/2026-05-12-cross-model-authoring/sma-probe-run3.md"

GRAMMAR = """\
## AILang Form-A grammar (complete construct grammar)

Every program is fully-parenthesised S-expressions. ARITY IS FIXED PER HEAD —
count your parens. `seq` takes EXACTLY 2 subterms; `if` EXACTLY 3; `let`
EXACTLY 3.

Module:  (module NAME DECL*)
Decl:
  (data NAME (vars V*)? (doc S)? (ctor CNAME FIELDTYPE*)*)
  (fn NAME (doc S)? (type TYPE) (params P*) (body TERM))
Type:
  (con NAME TYPEARG*)             ; type constructor: (con Int), (con Series (con Float))
  V                               ; a bare identifier is a type variable, e.g. a
  (fn-type (params PT*) (ret PT) (effects E*)?)
  (forall (vars V*) TYPE)
  (own TYPE) | (borrow TYPE)      ; parameter / return MODE wrappers (mandatory)
Term (head forms):
  (app FN ARG*)                   ; call FN with N value args
  (tail-app FN ARG*)              ; tail call (use for recursion)
  (term-ctor TYPENAME CNAME ARG*) ; build an ADT value
  (new TYPENAME TYPEARG* VALARG*) ; BUILT-IN allocator term (NOT a function, never (app new ...)):
                                  ;   (new Series (con Float) 3)
  (if COND THEN ELSE)             ; exactly 3 subterms
  (let NAME VALUE BODY)           ; exactly 3: bind NAME=VALUE, then BODY
  (let-rec NAME (type T) (params P*) (body B) IN)   ; local recursive fn, then IN
  (match SCRUT (case PAT ARM)*)   ; one or more arms
  (seq LHS RHS)                   ; EXACTLY 2; for >2 effects, NEST: (seq A (seq B C))
  (do OP ARG*)                    ; effect op: (do io/print_str "\\n")
  (loop (NAME TYPE INIT)* BODY)   ; typed loop binders, then body
  (recur ARG*)                    ; re-enter the enclosing loop
  literals: 42 (Int)  1.0 (Float)  "txt" (Str)  true | false (Bool)
  NAME                            ; a bare identifier is a variable reference
Pattern:
  (pat-ctor CNAME FIELDPAT*)      ; FIELDPAT = bare name (binds) or pat-wild
  (pat-lit LIT) | pat-wild
Params:  (params a b c)           ; bare parameter names

Prelude functions you will need: `+ - * /` (`forall a. (a,a)->a`),
`eq ge gt lt le` (comparisons -> Bool), `int_to_float : (Int)->Float`,
`print : forall a. (a) -> Unit !IO` (NO trailing newline; emit one with
`(do io/print_str "\\n")`).
"""

SERIES_LIB = """\
## Series library reference (this module already exists — CALL these functions, do not redefine them)

`Series a` is a bounded ring buffer; element type `a` is `Int` or `Float`;
financial indexing (index 0 = newest). The library source (its public fns and
their exact types/behaviour):

""" + (REPO / "crates/ailang-kernel/src/series/source.ail").read_text() + """

Call its functions as `Series.new`, `Series.push`, `Series.at`, `Series.len`,
`Series.total_count`. Allocate a fresh one with the `new` term:
`(new Series (con Float) N)`.
"""

TASK = """\
Write a complete AILang module named `sma` that computes a simple moving
average over a stream of floats with a fixed window of 3, using the Series
library.

`main` must:
1. Create a Float Series with lookback 3.
2. Process this input stream, one value at a time, pushing each into the
   Series: 1.0, 5.0, 3.0, 8.0, 6.0, 2.0
3. After each push, once at least 3 values have been pushed in total, print
   the average of the most recent 3 values (their sum divided by 3.0), each
   on its own line.

Expected stdout (one average per line):
3.0
5.33333
5.66667
5.33333

Return ONLY the AILang Form-A source for the module, nothing else.
"""

def extract_code(resp: str) -> str:
    m = re.search(r"```(?:[a-zA-Z]*)?\n(.*?)```", resp, re.S)
    code = m.group(1) if m else resp
    # keep from the first (module ... to the matching end
    i = code.find("(module")
    return code[i:].strip() if i >= 0 else code.strip()

def call(messages) -> tuple[str, dict]:
    body = json.dumps({"model": MODEL, "messages": messages,
                       "temperature": 0.2, "max_tokens": 4000}).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 evaluate(code: str):
    """Return (green: bool, stage: str, feedback: str)."""
    m = re.search(r"\(module\s+([A-Za-z_][A-Za-z0-9_]*)", code)
    name = m.group(1) if m else "sma"
    f = Path(f"/tmp/{name}.ail"); f.write_text(code + "\n")
    chk = subprocess.run([AIL, "check", str(f)], capture_output=True, text=True)
    if chk.returncode != 0:
        return False, "check", (chk.stderr or chk.stdout).strip()
    run = subprocess.run([AIL, "run", str(f)], capture_output=True, text=True)
    if run.returncode != 0:
        return False, "run", (run.stderr or run.stdout).strip()
    if run.stdout != EXPECTED:
        return False, "output", f"program ran but stdout was:\n{run.stdout!r}\nexpected:\n{EXPECTED!r}"
    return True, "green", run.stdout

def main():
    system = GRAMMAR + "\n\n" + SPEC + "\n\n" + SERIES_LIB
    messages = [{"role": "system", "content": system},
                {"role": "user", "content": TASK}]
    log = ["# SMA authoring probe — Qwen3-Coder-Next via IONOS\n",
           f"**Model:** {MODEL}  ", f"**Max turns:** {MAX_TURNS}  ",
           "**Task:** write the Series SMA worked example in AILang Form-A, "
           "given the Form-A mini-spec + Series API only (not the reference).\n",
           f"**Expected stdout:** `{EXPECTED!r}`\n", "---\n"]
    total_p = total_c = 0
    green = False
    for turn in range(1, MAX_TURNS + 1):
        resp, usage = call(messages)
        total_p += usage.get("prompt_tokens", 0); total_c += usage.get("completion_tokens", 0)
        code = extract_code(resp)
        ok, stage, fb = evaluate(code)
        log += [f"## Turn {turn} — `{stage}`" + ("  ✅ GREEN" if ok else "  ❌"),
                f"\nTokens this turn: prompt={usage.get('prompt_tokens','?')}, "
                f"completion={usage.get('completion_tokens','?')}\n",
                "```\n" + code + "\n```\n"]
        if ok:
            log += [f"**Result:** reached green on turn {turn}. Output:\n```\n{fb}```\n"]
            green = True
            print(f"GREEN on turn {turn}")
            break
        log += [f"**`ail {stage}` feedback:**\n```\n{fb[:1500]}\n```\n"]
        print(f"turn {turn}: {stage} fail")
        messages += [{"role": "assistant", "content": resp},
                     {"role": "user",
                      "content": f"That program failed at `ail {stage}`:\n\n{fb[:1500]}\n\n"
                                 "Fix it and return ONLY the corrected AILang Form-A module."}]
    log += ["\n---\n",
            f"**Outcome:** {'GREEN' if green else 'did NOT reach green'} after "
            f"{turn} turn(s).  ",
            f"**Total tokens:** prompt={total_p}, completion={total_c}.\n"]
    DOC.write_text("\n".join(log))
    print(f"progress doc → {DOC}")
    sys.exit(0 if green else 2)

if __name__ == "__main__":
    main()