#!/usr/bin/env python3
"""GIF Forge generation cycle: prompt file -> model -> render.py -> run.
Usage: cycle.py <prompt-file> <round-tag>
Auth: FORGE_API_KEY env var. Defaults to the DeepSeek chat endpoint;
override with FORGE_API_BASE and FORGE_MODEL for any OpenAI-compatible API.
The model writes the whole render script; this harness runs it.
Config block at the top is the only thing you should need to edit.
"""
import json
import os
import re
import subprocess
import sys
import urllib.request

# ---- config ---------------------------------------------------------------
ALLOWED_IMPORTS = {"PIL", "math", "random"}  # the import gate
RENDER_TIMEOUT = 240                        # seconds for render.py to run
API_BASE = os.environ.get("FORGE_API_BASE", "https://api.deepseek.com")
API_MODEL = os.environ.get("FORGE_MODEL", "deepseek-chat")
API_KEY = os.environ.get("FORGE_API_KEY", "")
# ---------------------------------------------------------------------------

def generate(prompt: str) -> str:
    if not API_KEY:
        raise SystemExit("FORGE_API_KEY is not set")
    body = {"model": API_MODEL, "messages": [{"role": "user", "content": prompt}]}
    req = urllib.request.Request(
        API_BASE.rstrip("/") + "/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json",
                 "Authorization": f"Bearer {API_KEY}"},
        method="POST")
    resp = urllib.request.urlopen(req, timeout=600)
    payload = json.loads(resp.read().decode())
    return payload["choices"][0]["message"]["content"]

def main():
    prompt_file, tag = sys.argv[1], sys.argv[2]
    prompt = open(prompt_file).read()
    ans = generate(prompt)
    open(f"answer-{tag}.md", "w").write(ans)
    m = re.search(r"```python\s*\n(.*?)```", ans, re.S) or \
        re.search(r"```\s*\n(.*?)```", ans, re.S)
    if not m:
        print("NO CODE BLOCK FOUND; answer saved to answer-%s.md" % tag)
        sys.exit(2)
    code = m.group(1)
    open("render.py", "w").write(code)
    bad = set(re.findall(r"^\s*(?:import|from)\s+([A-Za-z_][A-Za-z0-9_]*)", code, re.M))
    ok = bad <= ALLOWED_IMPORTS
    print(json.dumps({"imports": sorted(bad), "imports_ok": ok,
                      "code_chars": len(code)}))
    if not ok:
        print("IMPORT GATE FAILED — repair must remove: " + ", ".join(sorted(bad)))
        sys.exit(3)
    run = subprocess.run([sys.executable, "render.py"], capture_output=True,
                         text=True, timeout=RENDER_TIMEOUT)
    print("render rc:", run.returncode)
    print("stdout tail:", run.stdout[-500:])
    print("stderr tail:", run.stderr[-1500:])
    sys.exit(0 if run.returncode == 0 else 4)

if __name__ == "__main__":
    main()