#!/usr/bin/env python3
"""Control probe: the SAME SMA task, but in Python instead of AILang Form-A.

If the model nails this zero-shot, the AILang failures are about the unfamiliar
surface (structure-tracking in a never-seen language), NOT about coding ability.
If it fails here too, the model is just weak. Single-shot, temperature 0.2, same
as the AILang ablation — but no few-shot, because Python is the model's home turf.

This harness ONLY fetches the model's code and writes it to /tmp + a doc. It does
NOT execute it (LLM-generated code is run deliberately by the operator, after a
read). Reads $IONOS_API_TOKEN (never printed). Live IONOS — run with consent.
"""
from __future__ import annotations
import json, os, re, 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")
TAG      = os.environ.get("PROBE_TAG", "qwen")
REPO = Path(__file__).resolve().parents[2]
DOC  = REPO / f"experiments/2026-05-12-cross-model-authoring/{TAG}-python-sma.md"
OUT  = Path("/tmp/qwen_sma.py")

TASK = ("Write a Python 3 program that computes a simple moving average with "
        "window 3 over the float stream 1.0, 5.0, 3.0, 8.0, 6.0, 2.0: feed the "
        "values in one at a time, and after each value — once at least 3 values "
        "have been seen — print the average of the most recent 3 (their sum / 3) "
        "on its own line. The four printed lines should be the averages "
        "3.0, 5.33333..., 5.66667..., 5.33333... (exact float formatting is up to "
        "you). Return the program in a single ```python code block.")

def call(task):
    body = json.dumps({"model": MODEL, "temperature": 0.2, "max_tokens": 1500,
        "messages": [
            {"role": "system", "content": "You are an expert Python programmer."},
            {"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"```(?:python)?\n(.*?)```", resp, re.S)
    return (m.group(1) if m else resp).strip()

def main():
    resp, usage = call(TASK)
    code = extract(resp)
    OUT.write_text(code + "\n")
    log = [f"# Python control — same SMA task, in Python ({MODEL})\n",
           f"tokens: prompt={usage.get('prompt_tokens','?')} completion={usage.get('completion_tokens','?')}\n",
           "---\n", "extracted Python program:\n```python\n" + code + "\n```\n",
           f"\nwritten to: {OUT} (NOT executed by this harness)\n"]
    DOC.write_text("\n".join(log))
    print(f"completion_tokens={usage.get('completion_tokens','?')}")
    print(f"code → {OUT}")
    print(f"doc  → {DOC}")

if __name__ == "__main__":
    main()