#!/usr/bin/env python3
"""Run the survey-note prompt sequence against models via OpenRouter, with
our own fetch tool as the harness.

Why: in consumer apps the platform's fetch tool decides what a model can
reach (HTML yes, PDF no, sandbox offline). Here the only tool is
`fetch_url`, executed by this script with the repo's helpers, so all four
source types work and the comparison isolates the model's gathering and
writing discipline. The allowlist is enforced on our side: the model
cannot reach archives, mirrors or other sites even if it asks.

    OPENROUTER_API_KEY=... python3 modeltest.py --model x-ai/grok-4 \
        --prompts ../prompts/04-survey-synthesis.sequence.md \
        --out ../tests/runs --budget-usd 2.50

Transcripts (prompts, tool calls, tool results, answers, token usage) go
to <out>/<model-slug>/transcript.md and .json. Nothing is committed here.
"""
import argparse, datetime as dt, json, os, re, subprocess, sys, time, urllib.request

API = "https://openrouter.ai/api/v1/chat/completions"
ALLOW = ("https://caucasusbarometer.org/", "https://www.caucasusbarometer.org/",
         "https://pc-axis.geostat.ge/", "https://geostat.ge/", "https://www.geostat.ge/")
CRAWL_DELAY = 3
HERE = os.path.dirname(os.path.abspath(__file__))
VENV_PY = os.environ.get("GEORESEARCH_VENV_PY", "")  # python with pypdf + xlrd, for PDF/XLS

TOOL = {"type": "function", "function": {
    "name": "fetch_url",
    "description": "Fetch a URL on caucasusbarometer.org or geostat.ge and return its content as text: HTML pages as visible text, JSON as-is, PDF as extracted text, XLS as sheet dumps. Other hosts are refused.",
    "parameters": {"type": "object", "properties": {"url": {"type": "string"}}, "required": ["url"]}}}


def fetch(url):
    if not url.startswith(ALLOW):
        return f"REFUSED: {url} is outside the allowed sources (caucasusbarometer.org, geostat.ge)."
    time.sleep(CRAWL_DELAY)
    try:
        if "/PXWeb/api/" in url:
            return subprocess.run(["curl", "-s", "--max-time", "30", "-A", "Mozilla/5.0", url],
                                  capture_output=True, text=True, timeout=40).stdout[:30000]
        if "/downloads/attachments/" in url or url.lower().endswith((".pdf", ".xls", ".xlsx")):
            path = "/tmp/modeltest-" + re.sub(r"\W", "_", url)[-60:]
            subprocess.run(["curl", "-s", "-L", "--max-time", "60", "-A", "Mozilla/5.0", "-o", path, url], timeout=70)
            with open(path, "rb") as f:
                head = f.read(8)
            if head.startswith(b"%PDF"):
                code = "import sys,pypdf;r=pypdf.PdfReader(sys.argv[1]);print('\\n'.join(p.extract_text() or '' for p in r.pages))"
            elif head.startswith(b"\xd0\xcf\x11\xe0"):
                code = ("import sys,xlrd;wb=xlrd.open_workbook(sys.argv[1])\n"
                        "for sh in wb.sheets():\n print('SHEET',sh.name,sh.nrows,'x',sh.ncols)\n"
                        " for r in range(sh.nrows): print(' | '.join(str(sh.cell_value(r,c)).strip() for c in range(sh.ncols)))")
            else:
                return f"DOWNLOADED {os.path.getsize(path)} bytes but the file is neither PDF nor XLS (starts with {head!r})."
            py = VENV_PY or sys.executable
            out = subprocess.run([py, "-c", code, path], capture_output=True, text=True, timeout=120)
            return (out.stdout or out.stderr)[:30000]
        out = subprocess.run(["node", os.path.join(HERE, "fetch.js"), url, "--lenient", "--wait", "2500"],
                             capture_output=True, text=True, timeout=90)
        text = (out.stdout or out.stderr)[:30000]
        # visible text has no hrefs; append the page's links so the model can follow
        # attachments and variable pages the way a browser user could
        links = subprocess.run(["node", os.path.join(HERE, "links.js"), url], capture_output=True, text=True, timeout=90).stdout
        try:
            pairs = [(l["text"][:60], l["href"]) for l in json.loads(links) if l["href"].startswith(ALLOW)]
        except Exception:
            pairs = []
        if pairs:
            seen, lines = set(), []
            for t, h in pairs:
                if h in seen: continue
                seen.add(h); lines.append(f"{t} -> {h}")
            text += "\n\nLINKS ON THIS PAGE (text -> url):\n" + "\n".join(lines[:150])
        return text
    except Exception as e:
        return f"FETCH ERROR: {e}"


def chat(model, messages, key):
    body = json.dumps({"model": model, "messages": messages, "tools": [TOOL], "temperature": 0, "usage": {"include": True}}).encode()
    req = urllib.request.Request(API, data=body, headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json",
                                                          "HTTP-Referer": "https://openclaw.ge", "X-Title": "GeoResearch model test"})
    with urllib.request.urlopen(req, timeout=600) as r:
        return json.load(r)


def split_prompts(path):
    text = open(path, encoding="utf-8").read()
    blocks = re.findall(r"## PROMPT (\d+)[^\n]*\n(?:.*?\n)*?```\n(.*?)```", text, re.S)
    return [(int(n), b.strip()) for n, b in blocks]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True); ap.add_argument("--prompts", required=True)
    ap.add_argument("--out", required=True); ap.add_argument("--budget-usd", type=float, default=3.0)
    ap.add_argument("--max-tool-calls", type=int, default=80)
    a = ap.parse_args()
    key = os.environ.get("OPENROUTER_API_KEY") or sys.exit("OPENROUTER_API_KEY not set")
    slug = re.sub(r"[^\w.-]+", "_", a.model)
    outdir = os.path.join(a.out, slug); os.makedirs(outdir, exist_ok=True)
    messages, log, spent, calls = [], [], 0.0, 0
    def save():
        json.dump(log, open(os.path.join(outdir, "transcript.json"), "w"), indent=1, ensure_ascii=False)
    for n, prompt in split_prompts(a.prompts):
        messages.append({"role": "user", "content": prompt}); log.append({"prompt": n, "text": prompt})
        while True:
            resp = None
            for attempt in range(6):
                try:
                    resp = chat(a.model, messages, key); break
                except Exception as e:
                    detail = getattr(e, "read", lambda: b"")()
                    if getattr(e, "code", None) in (402, 429) and attempt < 5:
                        log.append({"retry": f"prompt {n}: HTTP {e.code}, waiting 130 s (attempt {attempt+1})"}); save()
                        time.sleep(130); continue
                    log.append({"stopped": f"API error on prompt {n}: {e} {detail[:500]!r}"}); save(); spent = a.budget_usd + 1; break
            if resp is None:
                break
            if "choices" not in resp:
                log.append({"stopped": f"unexpected response on prompt {n}: {json.dumps(resp)[:500]}"}); save(); spent = a.budget_usd + 1; break
            usage = resp.get("usage", {}); spent += float(usage.get("cost", 0) or 0)
            msg = resp["choices"][0]["message"]; messages.append(msg)
            log.append({"prompt": n, "assistant": msg.get("content"), "tool_calls": msg.get("tool_calls"), "usage": usage})
            if not msg.get("tool_calls"):
                break
            for tc in msg["tool_calls"]:
                calls += 1
                url = json.loads(tc["function"]["arguments"]).get("url", "")
                result = fetch(url) if calls <= a.max_tool_calls else "REFUSED: tool-call limit reached."
                messages.append({"role": "tool", "tool_call_id": tc["id"], "content": result})
                log.append({"prompt": n, "tool": url, "result_chars": len(result), "result": result})
            if spent > a.budget_usd:
                log.append({"stopped": f"budget {a.budget_usd} USD exceeded at {spent:.2f}"}); break
        if spent > a.budget_usd and not (log and "stopped" in log[-1]):
            log.append({"stopped": f"budget {a.budget_usd} USD exceeded at {spent:.2f} after prompt {n}"})
        save()
        if spent > a.budget_usd:
            break
    with open(os.path.join(outdir, "transcript.md"), "w") as f:
        f.write(f"# {a.model} — {dt.datetime.now(dt.timezone.utc):%Y-%m-%d %H:%M} UTC — tool calls {calls}, cost ${spent:.3f}\n\n")
        for e in log:
            if "text" in e: f.write(f"\n## PROMPT {e['prompt']}\n\n```\n{e['text']}\n```\n")
            elif "tool" in e: f.write(f"- fetch `{e['tool']}` → {e['result_chars']} chars\n")
            elif "assistant" in e and e["assistant"]: f.write(f"\n### answer\n\n{e['assistant']}\n")
            elif "stopped" in e: f.write(f"\n**{e['stopped']}**\n")
    print(f"{a.model}: {calls} tool calls, ${spent:.3f}, transcript in {outdir}")


if __name__ == "__main__":
    main()