"""
Backlog-scale triage via the Cerebras Batch API (Private Preview).

The live console is for interactive, real-time triage. This is the other half of the
enterprise story: point it at a *huge* backlog (up to 50k findings) and let Cerebras
process it asynchronously, guaranteed within 24h. Same Analyst agent + strict schema.

    python -m backend.batch_triage huge_scan.json

Notes:
  - Batch is Private Preview; your model must be batch-enabled for your org.
  - Minimum 10 requests per batch; this script auto-chunks at 50k.
  - Results are NOT ordered — we match on custom_id.
"""
from __future__ import annotations

import json
import os
import sys
import time

import httpx
from dotenv import load_dotenv

from .agents import ANALYST_SCHEMA, ANALYST_SYS, _analyst_user
from .models import normalize

load_dotenv()
BASE = os.getenv("CEREBRAS_BASE_URL", "https://api.cerebras.ai/v1").rstrip("/")
MODEL = os.getenv("CEREBRAS_MODEL", "gemma-4-31b")


def _headers() -> dict:
    key = os.environ.get("CEREBRAS_API_KEY")
    if not key:
        sys.exit("Set CEREBRAS_API_KEY to run a batch.")
    return {"Authorization": f"Bearer {key}"}


def build_jsonl(findings: list[dict], path: str) -> None:
    with open(path, "w", encoding="utf-8", newline="\n") as f:
        for i, fd in enumerate(findings):
            req = {
                "custom_id": fd["id"] or f"f-{i}",
                "method": "POST",
                "url": "/v1/chat/completions",
                "body": {
                    "model": MODEL,
                    "max_completion_tokens": 350,
                    "messages": [
                        {"role": "system", "content": ANALYST_SYS},
                        {"role": "user", "content": _analyst_user(fd)},
                    ],
                    "response_format": {"type": "json_schema", "json_schema":
                                        {"name": "out", "strict": True, "schema": ANALYST_SCHEMA}},
                },
            }
            f.write(json.dumps(req) + "\n")


def main(src: str) -> None:
    findings = normalize(open(src, encoding="utf-8").read())
    if len(findings) < 10:
        sys.exit("Batch needs >= 10 findings; use the live console for small scans.")
    print(f"{len(findings)} findings -> batch")
    build_jsonl(findings, "/tmp/batch_in.jsonl")

    with httpx.Client(timeout=120) as c:
        up = c.post(f"{BASE}/files", headers=_headers(),
                    files={"purpose": (None, "batch"),
                           "file": ("batch_in.jsonl", open("/tmp/batch_in.jsonl", "rb"))})
        up.raise_for_status()
        file_id = up.json()["id"]
        print("uploaded:", file_id)

        b = c.post(f"{BASE}/batches", headers={**_headers(), "Content-Type": "application/json"},
                   json={"input_file_id": file_id, "endpoint": "/v1/chat/completions",
                         "completion_window": "24h",
                         "metadata": {"app": "flashtriage"}})
        b.raise_for_status()
        bid = b.json()["id"]
        print("batch:", bid)

        while True:
            st = c.get(f"{BASE}/batches/{bid}", headers=_headers()).json()
            rc = st.get("request_counts", {})
            print(f"  {st['status']}  {rc.get('completed',0)}/{rc.get('total',0)}")
            if st["status"] in ("completed", "failed", "expired", "cancelled"):
                break
            time.sleep(10)

        out_id = st.get("output_file_id")
        if not out_id:
            sys.exit(f"no output ({st['status']})")
        res = c.get(f"{BASE}/files/{out_id}/content", headers=_headers())
        open("batch_results.jsonl", "wb").write(res.content)
        print("wrote batch_results.jsonl")


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "samples/sample_findings.json")