#!/usr/bin/env python3
"""braindump synthesize: turn an anonymous survey export into a review document.

Pipeline (deterministic Python around two model stages):

  load    CryptPad CSV export -> responses r01, r02, ... with per-question answers
  tag     model extracts atomic points per response, with verbatim quotes that
          Python verifies against the source text
  cluster model groups points into themes, flags novel points and dissent
  render  Python writes review.md from the structured data, so every count and
          quote is exact; the model never writes the final document

Runs on the Inference Cooperative's own gateway. Default model is a Tinfoil
TEE model so responses are processed inside an enclave. Stdlib only.

Usage:
  op run --env-file .env.op -- python3 synthesize.py responses.csv --out out/
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import re
import sys
import time
import urllib.error
import urllib.request
from collections import Counter, defaultdict
from datetime import date
from pathlib import Path

HERE = Path(__file__).resolve().parent
DEFAULT_BASE = "https://gateway.inference.coop/v1"
DEFAULT_MODEL = "tinfoil/glm-5.3"
KEY_ENV = "INFERENCE_COOP_API_KEY"
HTTP_TIMEOUT = 190  # gateway request_timeout is 180s; TimeoutErrorRetries is 0

# USD per 1M tokens (input, output), mirrored from code/litellm/config.yaml.
PRICES = {
    "tinfoil/glm-5.3": (1.80, 5.75),
    "tinfoil/deepseek-v4-1-flash": (0.65, 1.45),
    "tinfoil/gpt-oss-120b": (0.15, 0.60),
    "greenpt/glm-5.3": (1.80, 5.75),
    "greenpt/glm-5.3-flash": (0.11, 0.44),
    "greenpt/deepseek-v4.1-flash": (0.22, 1.10),
}

KINDS = {"idea", "want", "need", "goal", "concern"}
STANCES = {"for", "against", "neutral"}
META_COL_RE = re.compile(r"^(id|#|time|timestamp|date|submitted|created|author|user)\b", re.I)


# --------------------------------------------------------------------------- #
# Loading
# --------------------------------------------------------------------------- #

def _norm_ws(s: str) -> str:
    return re.sub(r"\s+", " ", s).strip()


def looks_like_choice_column(values: list[str]) -> bool:
    """Heuristic: a checkbox/choice export column has short, repeating tokens."""
    cells = [v for v in values if v.strip()]
    if len(cells) < 3:
        return False
    tokens: list[str] = []
    for c in cells:
        parts = [p.strip() for p in re.split(r"[;|,\n]", c) if p.strip()]
        if not parts or any(len(p) > 40 for p in parts):
            return False
        tokens.extend(parts)
    distinct = {t.lower() for t in tokens}
    return len(distinct) <= 20 and len(tokens) >= len(cells)


def split_choices(cell: str) -> list[str]:
    return [p.strip() for p in re.split(r"[;|,\n]", cell) if p.strip()]


def load_responses(path: Path, drop_cols: list[str] | None = None,
                   choice_cols: list[str] | None = None, limit: int | None = None) -> dict:
    """Return {"questions": [...], "choice_questions": [...], "responses": [...]}.

    Each response: {"id": "r01", "answers": {question: text}, "choices": {question: [..]}}.
    """
    with path.open(newline="", encoding="utf-8-sig") as fh:
        rows = list(csv.reader(fh))
    if not rows:
        raise SystemExit(f"{path}: empty file")
    header = [h.strip() for h in rows[0]]
    body = [r for r in rows[1:] if any(c.strip() for c in r)]
    if limit:
        body = body[:limit]

    drop = {c.lower() for c in (drop_cols or [])}
    keep_idx = []
    for i, h in enumerate(header):
        if not h or h.lower() in drop or META_COL_RE.match(h):
            continue
        if all((r[i].strip() == "" if i < len(r) else True) for r in body):
            continue  # entirely empty column
        keep_idx.append(i)

    forced_choice = {c.lower() for c in (choice_cols or [])}
    choice_idx = set()
    for i in keep_idx:
        col = [r[i] if i < len(r) else "" for r in body]
        if header[i].lower() in forced_choice or looks_like_choice_column(col):
            choice_idx.add(i)

    questions = [header[i] for i in keep_idx if i not in choice_idx]
    choice_questions = [header[i] for i in keep_idx if i in choice_idx]

    responses = []
    width = max(2, len(str(len(body))))
    for n, r in enumerate(body, start=1):
        rid = f"r{n:0{width}d}"
        answers, choices = {}, {}
        for i in keep_idx:
            cell = r[i].strip() if i < len(r) else ""
            if not cell:
                continue
            if i in choice_idx:
                choices[header[i]] = split_choices(cell)
            else:
                answers[header[i]] = cell
        if answers or choices:
            responses.append({"id": rid, "answers": answers, "choices": choices})
    return {"questions": questions, "choice_questions": choice_questions, "responses": responses}


# --------------------------------------------------------------------------- #
# Gateway client with on-disk cache
# --------------------------------------------------------------------------- #

class Gateway:
    def __init__(self, base: str, model: str, cache_dir: Path | None, use_cache: bool = True,
                 no_thinking: bool = False, reasoning_effort: str | None = None,
                 raw_dir: Path | None = None, verbose: bool = False):
        self.base = base.rstrip("/")
        self.model = model
        self.cache_dir = cache_dir
        self.use_cache = use_cache and cache_dir is not None
        self.no_thinking = no_thinking
        self.reasoning_effort = reasoning_effort
        self.raw_dir = raw_dir
        self.verbose = verbose
        self.calls = 0
        self.cache_hits = 0
        self.usage = Counter()

    # Separated so tests can monkeypatch the network.
    def _http_post(self, body: dict) -> dict:
        key = os.environ.get(KEY_ENV, "")
        if not key:
            raise SystemExit(f"{KEY_ENV} is not set. Run via: op run --env-file .env.op -- python3 synthesize.py ...")
        req = urllib.request.Request(
            f"{self.base}/chat/completions",
            data=json.dumps(body).encode("utf-8"),
            headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json",
                     "User-Agent": "braindump-synthesize/0.1"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as e:
            detail = e.read().decode("utf-8", "replace")[:800]
            raise SystemExit(f"gateway HTTP {e.code}: {detail}")

    def chat(self, system: str, user: str, max_tokens: int = 8000, tools: list | None = None,
             tag: str = "") -> dict:
        body: dict = {
            "model": self.model,
            "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}],
            "max_tokens": max_tokens,
            "temperature": 0.2,
        }
        if tools:
            body["tools"] = tools
            body["tool_choice"] = {"type": "function", "function": {"name": tools[0]["function"]["name"]}}
        if self.no_thinking:
            body["chat_template_kwargs"] = {"enable_thinking": False}
        if self.reasoning_effort:
            body["reasoning_effort"] = self.reasoning_effort

        digest = hashlib.sha256(json.dumps(body, sort_keys=True).encode()).hexdigest()
        cache_file = self.cache_dir / f"{digest}.json" if self.cache_dir else None
        if self.use_cache and cache_file and cache_file.exists():
            self.cache_hits += 1
            data = json.loads(cache_file.read_text())
        else:
            self.calls += 1
            t0 = time.time()
            data = self._http_post(body)
            if self.verbose:
                print(f"  [{tag or 'call'}] {time.time() - t0:.1f}s finish={_finish(data)} usage={data.get('usage')}", file=sys.stderr)
            if cache_file:
                cache_file.parent.mkdir(parents=True, exist_ok=True)
                cache_file.write_text(json.dumps(data, indent=1))
            if self.raw_dir:
                self.raw_dir.mkdir(parents=True, exist_ok=True)
                stamp = time.strftime("%Y%m%dT%H%M%S")
                (self.raw_dir / f"{stamp}-{tag or 'call'}-{digest[:8]}.json").write_text(
                    json.dumps({"request": _redact(body), "response": data}, indent=1))
        u = data.get("usage") or {}
        self.usage["prompt_tokens"] += u.get("prompt_tokens", 0)
        self.usage["completion_tokens"] += u.get("completion_tokens", 0)
        return data


def _redact(body: dict) -> dict:
    # Request bodies carry survey text; keep them for fixtures but never a key.
    return body


def _finish(data: dict) -> str:
    try:
        return data["choices"][0].get("finish_reason", "?")
    except (KeyError, IndexError, TypeError):
        return "?"


def message_text(data: dict) -> str:
    try:
        msg = data["choices"][0]["message"]
    except (KeyError, IndexError, TypeError):
        raise ValueError(f"unexpected response shape: {json.dumps(data)[:300]}")
    if msg.get("tool_calls"):
        return msg["tool_calls"][0]["function"]["arguments"]
    return msg.get("content") or ""


# --------------------------------------------------------------------------- #
# JSON extraction
# --------------------------------------------------------------------------- #

FENCE_RE = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.S)


def extract_json(text: str) -> dict:
    """Pull one JSON object out of model output that may be wrapped in prose or a fence."""
    text = text.strip()
    m = FENCE_RE.search(text)
    candidates = [m.group(1)] if m else []
    if text.startswith("{"):
        candidates.append(text)
    start, end = text.find("{"), text.rfind("}")
    if start != -1 and end > start:
        candidates.append(text[start:end + 1])
    last_err: Exception | None = None
    for c in candidates:
        try:
            obj = json.loads(c)
            if isinstance(obj, dict):
                return obj
        except json.JSONDecodeError as e:
            last_err = e
            try:  # tolerate trailing commas
                obj = json.loads(re.sub(r",\s*([}\]])", r"\1", c))
                if isinstance(obj, dict):
                    return obj
            except json.JSONDecodeError as e2:
                last_err = e2
    raise ValueError(f"no JSON object found ({last_err})")


def _check_reasoning_exhaustion(data: dict, max_tokens: int, tag: str) -> None:
    """Reasoning models can spend the whole completion budget thinking and return no content.

    Retrying the same request would just bill the same tokens again, so stop here with advice.
    """
    if _finish(data) != "length":
        return
    try:
        msg = data["choices"][0]["message"]
        usage = data.get("usage") or {}
        reasoning = (usage.get("completion_tokens_details") or {}).get("reasoning_tokens") or 0
    except (KeyError, IndexError, TypeError):
        return
    if not (msg.get("content") or msg.get("tool_calls")) and reasoning:
        raise SystemExit(
            f"{tag}: the model used all {max_tokens} completion tokens on hidden reasoning "
            f"({reasoning} reasoning tokens) and returned no content. Not retrying. Try "
            f"--no-thinking, --reasoning-effort low, a larger --max-tokens, or a different --model.")


def call_json(gw: Gateway, system: str, user: str, required: list[str], tag: str,
              tools: list | None = None, max_tokens: int = 8000) -> dict:
    """Ask for JSON; on a parse/shape failure retry once with the error appended."""
    data = gw.chat(system, user, max_tokens=max_tokens, tools=tools, tag=tag)
    _check_reasoning_exhaustion(data, max_tokens, tag)
    try:
        obj = extract_json(message_text(data))
        missing = [k for k in required if k not in obj]
        if missing:
            raise ValueError(f"missing keys {missing}")
        return obj
    except ValueError as e:
        if _finish(data) == "length":
            e = ValueError(f"{e}; output hit max_tokens={max_tokens} (finish_reason=length)")
        retry_user = (f"{user}\n\n---\nYour previous reply could not be parsed: {e}. "
                      f"Reply again with exactly one fenced ```json block containing the keys {required} and nothing else.")
        data = gw.chat(system, retry_user, max_tokens=max_tokens, tools=tools, tag=f"{tag}-retry")
        obj = extract_json(message_text(data))
        missing = [k for k in required if k not in obj]
        if missing:
            raise SystemExit(f"{tag}: model output still missing keys {missing} after retry")
        return obj


# --------------------------------------------------------------------------- #
# Stage A: tag
# --------------------------------------------------------------------------- #

def q_index(questions: list[str]) -> dict[str, str]:
    return {q: f"Q{i}" for i, q in enumerate(questions, start=1)}


def format_batch(batch: list[dict], qidx: dict[str, str]) -> str:
    parts = []
    for r in batch:
        parts.append(f"## Response {r['id']}")
        for q, text in r["answers"].items():
            parts.append(f"### {r['id']} / {qidx[q]}: {q}\n{text}")
    return "\n\n".join(parts)


def validate_items(raw_items: list, batch: list[dict], qidx: dict[str, str]) -> tuple[list[dict], list[str]]:
    by_id = {r["id"]: r for r in batch}
    q_by_idx = {v: k for k, v in qidx.items()}
    kept, dropped = [], []
    for it in raw_items:
        if not isinstance(it, dict):
            dropped.append("non-object item")
            continue
        rid, q, quote = it.get("id"), it.get("q"), str(it.get("quote") or "")
        label = _norm_ws(str(it.get("label") or ""))
        if rid not in by_id or q not in q_by_idx or not label:
            dropped.append(f"{rid}/{q}: unknown id/question or empty label")
            continue
        source = by_id[rid]["answers"].get(q_by_idx[q], "")
        if not quote or _norm_ws(quote) not in _norm_ws(source):
            dropped.append(f"{rid}/{q}: quote not found in source: {quote[:60]!r}")
            continue
        kind = it.get("kind") if it.get("kind") in KINDS else "idea"
        stance = it.get("stance") if it.get("stance") in STANCES else "neutral"
        kept.append({"rid": rid, "q": q, "question": q_by_idx[q], "label": label,
                     "kind": kind, "stance": stance, "quote": _norm_ws(quote)})
    return kept, dropped


TAG_TOOL = [{"type": "function", "function": {
    "name": "submit_items", "description": "Submit the extracted atomic points.",
    "parameters": {"type": "object", "required": ["items"], "properties": {"items": {
        "type": "array", "items": {"type": "object",
            "required": ["id", "q", "label", "kind", "stance", "quote"],
            "properties": {"id": {"type": "string"}, "q": {"type": "string"},
                           "label": {"type": "string"}, "kind": {"type": "string", "enum": sorted(KINDS)},
                           "stance": {"type": "string", "enum": sorted(STANCES)}, "quote": {"type": "string"}}}}}}}}]


def stage_tag(gw: Gateway, loaded: dict, batch_size: int, use_tools: bool,
              max_tokens: int = 8000) -> tuple[list[dict], list[str]]:
    system = (HERE / "prompts" / "tag.md").read_text()
    qidx = q_index(loaded["questions"])
    responses = [r for r in loaded["responses"] if r["answers"]]
    items, dropped = [], []
    for b in range(0, len(responses), batch_size):
        batch = responses[b:b + batch_size]
        user = format_batch(batch, qidx)
        obj = call_json(gw, system, user, ["items"], tag=f"tag{b // batch_size + 1}",
                        tools=TAG_TOOL if use_tools else None, max_tokens=max_tokens)
        kept, drop = validate_items(obj.get("items") or [], batch, qidx)
        items.extend(kept)
        dropped.extend(drop)
    for n, it in enumerate(items, start=1):
        it["iid"] = f"i{n:03d}"
    return items, dropped


# --------------------------------------------------------------------------- #
# Stage B: cluster
# --------------------------------------------------------------------------- #

CLUSTER_TOOL = [{"type": "function", "function": {
    "name": "submit_clusters", "description": "Submit themes, novel items, dissent and open questions.",
    "parameters": {"type": "object", "required": ["themes", "novel", "dissent", "open_questions"],
                   "properties": {
                       "themes": {"type": "array", "items": {"type": "object", "required": ["name", "summary", "item_ids"],
                                  "properties": {"name": {"type": "string"}, "summary": {"type": "string"},
                                                 "item_ids": {"type": "array", "items": {"type": "string"}}}}},
                       "novel": {"type": "array", "items": {"type": "string"}},
                       "dissent": {"type": "array", "items": {"type": "object", "required": ["theme", "item_ids", "summary"],
                                   "properties": {"theme": {"type": "string"}, "summary": {"type": "string"},
                                                  "item_ids": {"type": "array", "items": {"type": "string"}}}}},
                       "open_questions": {"type": "array", "items": {"type": "object", "required": ["question", "item_ids"],
                                          "properties": {"question": {"type": "string"},
                                                         "item_ids": {"type": "array", "items": {"type": "string"}}}}}}}}}]


def format_items(items: list[dict]) -> str:
    return "\n".join(f"{it['iid']} [{it['rid']}/{it['q']} {it['kind']} {it['stance']}] {it['label']}" for it in items)


def validate_clusters(obj: dict, items: list[dict]) -> dict:
    known = {it["iid"] for it in items}
    seen: set[str] = set()
    themes = []
    for t in obj.get("themes") or []:
        if not isinstance(t, dict):
            continue
        ids = []
        for i in t.get("item_ids") or []:
            if i in known and i not in seen:
                ids.append(i)
                seen.add(i)
        if ids:
            themes.append({"name": _norm_ws(str(t.get("name") or "Untitled")),
                           "summary": _norm_ws(str(t.get("summary") or "")), "item_ids": ids})
    unthemed = [i for i in sorted(known) if i not in seen]
    novel = [i for i in (obj.get("novel") or []) if i in known][:10]
    dissent = []
    for d in obj.get("dissent") or []:
        if not isinstance(d, dict):
            continue
        ids = [i for i in (d.get("item_ids") or []) if i in known]
        if ids:
            dissent.append({"theme": _norm_ws(str(d.get("theme") or "")),
                            "summary": _norm_ws(str(d.get("summary") or "")), "item_ids": ids})
    oq = []
    for q in (obj.get("open_questions") or [])[:6]:
        if isinstance(q, dict) and q.get("question"):
            oq.append({"question": _norm_ws(str(q["question"])),
                       "item_ids": [i for i in (q.get("item_ids") or []) if i in known]})
        elif isinstance(q, str):
            oq.append({"question": _norm_ws(q), "item_ids": []})
    return {"themes": themes, "unthemed": unthemed, "novel": novel, "dissent": dissent, "open_questions": oq}


def stage_cluster(gw: Gateway, items: list[dict], use_tools: bool, max_tokens: int = 8000) -> dict:
    system = (HERE / "prompts" / "cluster.md").read_text()
    obj = call_json(gw, system, format_items(items), ["themes"], tag="cluster",
                    tools=CLUSTER_TOOL if use_tools else None, max_tokens=int(max_tokens * 1.5))
    return validate_clusters(obj, items)


# --------------------------------------------------------------------------- #
# Render
# --------------------------------------------------------------------------- #

def render(loaded: dict, items: list[dict], clusters: dict, meta: dict) -> str:
    by_iid = {it["iid"]: it for it in items}
    resp_by_id = {r["id"]: r for r in loaded["responses"]}
    n_resp = len(loaded["responses"])

    def quote_line(iid: str) -> str:
        it = by_iid[iid]
        return f'> "{it["quote"]}" ({it["rid"]}, {it["q"]})'

    out: list[str] = []
    out.append(f"# Braindump synthesis\n")
    out.append(f"_Generated {meta['date']} from {meta['source']} · {n_resp} responses · "
               f"{len(items)} points · {len(clusters['themes'])} themes · model `{meta['model']}`_\n")
    out.append("Counts below are computed from the extracted points, not estimated by the model. "
               "Quotes are verbatim from the responses. The full responses are in the appendix.\n")

    # Priorities (checkbox tally in Python)
    for cq in loaded["choice_questions"]:
        tally = Counter()
        voters = 0
        for r in loaded["responses"]:
            picks = r["choices"].get(cq)
            if picks:
                voters += 1
                tally.update(p for p in picks)
        if not tally:
            continue
        out.append(f"## {cq}\n")
        out.append(f"_{voters} of {n_resp} responses answered._\n")
        width = max(tally.values())
        for opt, c in tally.most_common():
            bar = "█" * round(12 * c / width)
            out.append(f"- {bar} **{c}** · {opt}")
        out.append("")

    # Themes ordered by distinct respondents
    def theme_key(t):
        rids = {by_iid[i]["rid"] for i in t["item_ids"]}
        return (-len(rids), -len(t["item_ids"]))

    out.append("## Themes\n")
    out.append("| # | Theme | People | Points | For / Against |")
    out.append("|---|---|---|---|---|")
    themes = sorted(clusters["themes"], key=theme_key)
    for n, t in enumerate(themes, start=1):
        rids = {by_iid[i]["rid"] for i in t["item_ids"]}
        st = Counter(by_iid[i]["stance"] for i in t["item_ids"])
        out.append(f"| {n} | {t['name']} | {len(rids)} | {len(t['item_ids'])} | {st['for']} / {st['against']} |")
    out.append("")
    for n, t in enumerate(themes, start=1):
        rids = {by_iid[i]["rid"] for i in t["item_ids"]}
        kinds = Counter(by_iid[i]["kind"] for i in t["item_ids"])
        out.append(f"### {n}. {t['name']}\n")
        out.append(f"_{len(rids)} of {n_resp} people · {', '.join(f'{v} {k}' for k, v in kinds.most_common())}_\n")
        if t["summary"]:
            out.append(t["summary"] + "\n")
        labels = Counter(by_iid[i]["label"] for i in t["item_ids"])
        out.append("Points raised:\n")
        for label, c in labels.most_common():
            out.append(f"- {label}" + (f" (×{c})" if c > 1 else ""))
        out.append("")
        # Up to 4 quotes from distinct respondents
        shown: set[str] = set()
        for i in t["item_ids"]:
            rid = by_iid[i]["rid"]
            if rid in shown:
                continue
            out.append(quote_line(i))
            shown.add(rid)
            if len(shown) >= 4:
                break
        out.append("")

    if clusters["novel"]:
        out.append("## Standout ideas\n")
        out.append("_Raised by one person, specific enough to be worth a look._\n")
        for i in clusters["novel"]:
            it = by_iid[i]
            out.append(f"- **{it['label']}** ({it['kind']})  ")
            out.append(f"  {quote_line(i)[2:]}")
        out.append("")

    if clusters["dissent"]:
        out.append("## Notable dissent\n")
        out.append("_Places where people pull in different directions. Worth time in the meeting._\n")
        for d in clusters["dissent"]:
            head = f"**{d['theme']}**" if d["theme"] else "**Standalone concern**"
            out.append(f"- {head}: {d['summary']}")
            for i in d["item_ids"][:3]:
                out.append(f"  {quote_line(i)}")
        out.append("")

    if clusters["open_questions"]:
        out.append("## Open questions for the group\n")
        for n, q in enumerate(clusters["open_questions"], start=1):
            refs = ", ".join(sorted({by_iid[i]["rid"] for i in q["item_ids"]}))
            out.append(f"{n}. {q['question']}" + (f" _(from {refs})_" if refs else ""))
        out.append("")

    if clusters["unthemed"]:
        out.append("## Not placed in a theme\n")
        for i in clusters["unthemed"]:
            it = by_iid[i]
            out.append(f"- {it['label']} ({it['rid']}, {it['q']})")
        out.append("")

    out.append("## Appendix: all responses, verbatim\n")
    for q in loaded["questions"]:
        answered = [(r["id"], r["answers"][q]) for r in loaded["responses"] if q in r["answers"]]
        if not answered:
            continue
        out.append(f"### {q}\n")
        for rid, text in answered:
            body = text.replace("\n", "\n  ")
            out.append(f"- **{rid}:** {body}")
        out.append("")
    if meta.get("dropped"):
        out.append("## Processing notes\n")
        out.append(f"{len(meta['dropped'])} extracted point(s) were discarded because their quote could not be "
                   "verified against the source text or they referenced an unknown response.\n")
    return "\n".join(out).rstrip() + "\n"


# --------------------------------------------------------------------------- #
# Estimation
# --------------------------------------------------------------------------- #

def est_tokens(text: str) -> int:
    return max(1, len(text) // 4)


def estimate(loaded: dict, model: str, batch_size: int) -> dict:
    tag_sys = est_tokens((HERE / "prompts" / "tag.md").read_text())
    cl_sys = est_tokens((HERE / "prompts" / "cluster.md").read_text())
    responses = [r for r in loaded["responses"] if r["answers"]]
    n_answers = sum(len(r["answers"]) for r in responses)
    batches = max(1, -(-len(responses) // batch_size))
    tag_in = batches * tag_sys + sum(est_tokens(format_batch([r], q_index(loaded["questions"]))) for r in responses)
    est_items = int(n_answers * 1.6)
    tag_out = est_items * 60
    cl_in = cl_sys + est_items * 22
    cl_out = 1500 + est_items * 6
    pin, pout = PRICES.get(model, (1.80, 5.75))
    cost = (tag_in + cl_in) / 1e6 * pin + (tag_out + cl_out) / 1e6 * pout
    return {"responses": len(responses), "answers": n_answers, "batches": batches,
            "est_items": est_items, "input_tokens": tag_in + cl_in, "output_tokens": tag_out + cl_out,
            "est_cost_usd": round(cost, 3), "priced": model in PRICES}


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #

def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("csv", type=Path, help="CryptPad Form CSV export")
    ap.add_argument("--out", type=Path, default=Path("out"), help="output directory (default: out/)")
    ap.add_argument("--model", default=DEFAULT_MODEL)
    ap.add_argument("--base", default=os.environ.get("INFERENCE_COOP_BASE", DEFAULT_BASE))
    ap.add_argument("--batch-size", type=int, default=15, help="responses per tagging call (default 15)")
    ap.add_argument("--limit", type=int, help="only the first N responses (smoke tests)")
    ap.add_argument("--drop-col", action="append", default=[], help="column header to ignore (repeatable)")
    ap.add_argument("--choice-col", action="append", default=[], help="column header to tally as checkbox (repeatable)")
    ap.add_argument("--dry-run", action="store_true", help="parse, estimate tokens and cost, make no calls")
    ap.add_argument("--no-cache", action="store_true")
    ap.add_argument("--cache-dir", type=Path, default=Path(".cache"))
    ap.add_argument("--save-raw", type=Path, help="directory to save raw request/response JSON (fixtures)")
    ap.add_argument("--use-tools", action="store_true", help="force a function call for structured output")
    ap.add_argument("--no-thinking", action="store_true", help="send chat_template_kwargs.enable_thinking=false")
    ap.add_argument("--reasoning-effort", choices=["default", "none", "minimal", "low", "medium", "high"], default="low",
                    help="reasoning_effort sent to the model (default: low; 'default' sends nothing). "
                         "Without it GLM-5.3 can spend the whole completion budget thinking")
    ap.add_argument("--max-tokens", type=int, default=8000, help="completion budget per call (default 8000)")
    ap.add_argument("--allow-non-tee", action="store_true", help="permit a model outside Tinfoil's TEE enclaves")
    ap.add_argument("-v", "--verbose", action="store_true")
    a = ap.parse_args(argv)

    if not a.model.startswith("tinfoil/") and not a.allow_non_tee:
        print(f"refusing: {a.model} is not a Tinfoil TEE model, so responses would be processed outside an "
              f"enclave under policy-based privacy. Pass --allow-non-tee to do it anyway.", file=sys.stderr)
        return 2

    loaded = load_responses(a.csv, drop_cols=a.drop_col, choice_cols=a.choice_col, limit=a.limit)
    print(f"loaded {len(loaded['responses'])} responses · {len(loaded['questions'])} open questions · "
          f"{len(loaded['choice_questions'])} choice question(s)", file=sys.stderr)
    for q in loaded["questions"]:
        print(f"  Q: {q}", file=sys.stderr)
    for q in loaded["choice_questions"]:
        print(f"  choice: {q}", file=sys.stderr)

    est = estimate(loaded, a.model, a.batch_size)
    print(f"estimate: ~{est['input_tokens']:,} in / ~{est['output_tokens']:,} out tokens, "
          f"~{est['batches']} tagging call(s) + 1 cluster call, ~${est['est_cost_usd']}"
          + ("" if est["priced"] else " (model not in price table; GLM-5.3 rates assumed)"), file=sys.stderr)
    if a.dry_run:
        return 0

    gw = Gateway(a.base, a.model, None if a.no_cache else a.cache_dir, use_cache=not a.no_cache,
                 no_thinking=a.no_thinking,
                 reasoning_effort=None if a.reasoning_effort == "default" else a.reasoning_effort,
                 raw_dir=a.save_raw, verbose=a.verbose)
    items, dropped = stage_tag(gw, loaded, a.batch_size, a.use_tools, a.max_tokens)
    print(f"tagged: {len(items)} points kept, {len(dropped)} dropped", file=sys.stderr)
    for d in dropped[:10]:
        print(f"  dropped: {d}", file=sys.stderr)
    if not items:
        print("no points extracted; nothing to cluster", file=sys.stderr)
        return 1
    clusters = stage_cluster(gw, items, a.use_tools, a.max_tokens)
    print(f"clustered: {len(clusters['themes'])} themes · {len(clusters['novel'])} novel · "
          f"{len(clusters['dissent'])} dissent · {len(clusters['unthemed'])} unthemed", file=sys.stderr)

    meta = {"date": date.today().isoformat(), "source": a.csv.name, "model": a.model, "dropped": dropped}
    a.out.mkdir(parents=True, exist_ok=True)
    (a.out / "review.md").write_text(render(loaded, items, clusters, meta))
    (a.out / "synthesis.json").write_text(json.dumps(
        {"meta": meta, "questions": loaded["questions"], "choice_questions": loaded["choice_questions"],
         "items": items, "clusters": clusters, "usage": dict(gw.usage),
         "calls": gw.calls, "cache_hits": gw.cache_hits}, indent=1, ensure_ascii=False))
    pin, pout = PRICES.get(a.model, (1.80, 5.75))
    spent = gw.usage["prompt_tokens"] / 1e6 * pin + gw.usage["completion_tokens"] / 1e6 * pout
    print(f"wrote {a.out / 'review.md'} and {a.out / 'synthesis.json'} · {gw.calls} call(s), "
          f"{gw.cache_hits} cache hit(s) · {gw.usage['prompt_tokens']:,} in / {gw.usage['completion_tokens']:,} out "
          f"· ≈${spent:.3f}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())