#!/usr/bin/env python3
"""Inter-coder re-rating for the IETF AI/agent landscape survey.

Re-rates the clean IETF corpus (source='ietf', not false-positive) with TWO
models (Sonnet + Haiku) using the EXACT pinned production prompt
(``RATE_PROMPT_COMPACT``, abstract[:2000]) via the Anthropic Batch API (50% off).

Safety / reproducibility:
- Does NOT touch the production ``ratings`` table. Output goes to
  ``data/rerate/<model-alias>.jsonl`` (one JSON object per draft).
- Batch IDs are persisted to ``data/rerate/batches.json`` so the run resumes.
- Idempotent: drafts already present in the output JSONL are skipped on re-submit.

Usage:
    PYTHONPATH=src python3 scripts/rerate-intercoder.py --dry-run   # cost estimate, submit nothing
    PYTHONPATH=src python3 scripts/rerate-intercoder.py --submit    # create batches
    PYTHONPATH=src python3 scripts/rerate-intercoder.py --collect   # poll + write results
    PYTHONPATH=src python3 scripts/rerate-intercoder.py --run       # submit then collect (blocking)
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

from anthropic import Anthropic

from ietf_analyzer.analyzer import (
    CATEGORIES_SHORT,
    RATE_PROMPT_COMPACT,
    _doc_type_label,
)
from ietf_analyzer.config import Config
from ietf_analyzer.db import Database

OUT_DIR = Path("data/rerate")
BATCH_FILE = OUT_DIR / "batches.json"
MAX_TOKENS = 512

# Anthropic batch pricing is 50% of standard. Standard (USD per 1M tokens):
PRICING = {  # input, output
    "sonnet": (3.00, 15.00),
    "haiku": (1.00, 5.00),
}
MODELS = {
    "sonnet": lambda c: c.claude_model,
    "haiku": lambda c: c.claude_model_cheap,
}


def clean_ietf_drafts(db: Database):
    """The survey corpus: source='ietf', not flagged false-positive."""
    rows = db.conn.execute(
        """SELECT d.name FROM drafts d JOIN ratings r ON d.name = r.draft_name
           WHERE d.source = 'ietf'
             AND (r.false_positive = 0 OR r.false_positive IS NULL)
           ORDER BY d.name"""
    ).fetchall()
    return [r[0] for r in rows]


def build_prompt(db: Database, name: str) -> str | None:
    d = db.get_draft(name)
    if d is None:
        return None
    return RATE_PROMPT_COMPACT.format(
        doc_type=_doc_type_label(d.source),
        name=d.name, title=d.title, time=d.date,
        pages=d.pages or "?",
        abstract=d.abstract[:2000],
        categories=", ".join(CATEGORIES_SHORT),
    )


def already_done(alias: str) -> set[str]:
    p = OUT_DIR / f"{alias}.jsonl"
    if not p.exists():
        return set()
    done = set()
    for line in p.read_text().splitlines():
        if line.strip():
            done.add(json.loads(line)["draft_name"])
    return done


def load_batches() -> dict:
    if BATCH_FILE.exists():
        return json.loads(BATCH_FILE.read_text())
    return {}


def save_batches(b: dict):
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    BATCH_FILE.write_text(json.dumps(b, indent=2))


def cmd_dry_run(db: Database, cfg: Config, names: list[str]):
    total_in = 0
    sample = None
    for n in names:
        p = build_prompt(db, n)
        if p is None:
            continue
        total_in += len(p) // 4  # ~4 chars/token
        if sample is None:
            sample = p
    out_est = len(names) * 350  # observed compact-JSON output size
    print(f"corpus: {len(names)} clean IETF drafts")
    print(f"est input tokens/draft: ~{total_in // max(len(names),1)}")
    print(f"est total input tokens: ~{total_in:,} | output: ~{out_est:,}")
    print("\nestimated cost (Batch API = 50% of standard):")
    grand = 0.0
    for alias, (pin, pout) in PRICING.items():
        c = (total_in / 1e6 * pin + out_est / 1e6 * pout) * 0.5
        grand += c
        print(f"  {alias:7} ({MODELS[alias](cfg)}): ${c:.2f}")
    print(f"  {'BOTH':7}: ${grand:.2f}")
    print(f"\n--- sample prompt ({len(sample)} chars) ---\n{sample[:600]}\n...")


def cmd_submit(db: Database, cfg: Config, names: list[str], client: Anthropic):
    batches = load_batches()
    for alias in MODELS:
        if batches.get(alias, {}).get("id"):
            print(f"[{alias}] batch already submitted: {batches[alias]['id']} — skip")
            continue
        done = already_done(alias)
        todo = [n for n in names if n not in done]
        if not todo:
            print(f"[{alias}] all {len(names)} already collected — nothing to submit")
            continue
        model = MODELS[alias](cfg)
        requests = []
        for n in todo:
            p = build_prompt(db, n)
            if p is None:
                continue
            requests.append({
                "custom_id": n,
                "params": {
                    "model": model,
                    "max_tokens": MAX_TOKENS,
                    "messages": [{"role": "user", "content": p}],
                },
            })
        batch = client.messages.batches.create(requests=requests)
        batches[alias] = {"id": batch.id, "model": model, "count": len(requests)}
        save_batches(batches)
        print(f"[{alias}] submitted {len(requests)} requests → batch {batch.id}")


def cmd_collect(client: Anthropic, poll: bool):
    batches = load_batches()
    if not batches:
        print("no batches submitted yet (run --submit)")
        return False
    all_done = True
    for alias, info in batches.items():
        bid = info["id"]
        b = client.messages.batches.retrieve(bid)
        print(f"[{alias}] {bid}: {b.processing_status} "
              f"(succeeded={b.request_counts.succeeded}, errored={b.request_counts.errored}, "
              f"processing={b.request_counts.processing})")
        if b.processing_status != "ended":
            all_done = False
            continue
        out_path = OUT_DIR / f"{alias}.jsonl"
        done = already_done(alias)
        n_new = 0
        with out_path.open("a") as f:
            for result in client.messages.batches.results(bid):
                cid = result.custom_id
                if cid in done:
                    continue
                if result.result.type != "succeeded":
                    f.write(json.dumps({"draft_name": cid, "error": result.result.type}) + "\n")
                    continue
                msg = result.result.message
                raw = msg.content[0].text
                rec = {
                    "draft_name": cid,
                    "model": info["model"],
                    "raw": raw,
                    "in_tok": msg.usage.input_tokens,
                    "out_tok": msg.usage.output_tokens,
                }
                f.write(json.dumps(rec) + "\n")
                n_new += 1
        print(f"[{alias}] wrote {n_new} new results → {out_path}")
    return all_done


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--submit", action="store_true")
    ap.add_argument("--collect", action="store_true")
    ap.add_argument("--run", action="store_true", help="submit then poll-collect until done")
    args = ap.parse_args()

    cfg = Config.load()
    db = Database(cfg)
    names = clean_ietf_drafts(db)

    if args.dry_run:
        cmd_dry_run(db, cfg, names)
        return

    client = Anthropic()
    if args.submit or args.run:
        cmd_submit(db, cfg, names, client)
    if args.collect:
        cmd_collect(client, poll=False)
    if args.run:
        while True:
            if cmd_collect(client, poll=True):
                print("all batches ended.")
                break
            print("...waiting 60s for batches to finish")
            time.sleep(60)


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