"""
Classify doc_check_controls (1874 MCs) into check_type:
  - text    : MC asks about the DOCUMENT'S WORDING ("enthaelt die DSE..", "wird im Impressum genannt..")
  - process : MC asks about a TECHNICAL OR ORGANISATIONAL CONTROL ("ist sichergestellt..", "ist implementiert..")
  - review  : MC asks an ABSTRACT META-QUESTION that needs human judgment ("sind alle Verarbeitungen vollstaendig?")

Output goes to sidecar SQLite at /data/mc_classification.db (no Postgres schema change
per CLAUDE.md guardrails). Schema:

  CREATE TABLE mc_classification (
    control_id TEXT PRIMARY KEY,
    doc_type   TEXT,
    title      TEXT,
    check_type TEXT,    -- text|process|review
    confidence REAL,    -- 0..1
    rationale  TEXT,
    classified_at TEXT
  );

Run from inside bp-compliance-backend container:
  docker exec bp-compliance-backend python /app/scripts/classify_mc_check_type.py [--limit N] [--sample]
"""

from __future__ import annotations

import argparse
import json
import os
import sqlite3
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

import httpx
import psycopg2
from psycopg2.extras import RealDictCursor

ANTHROPIC_URL = "https://api.anthropic.com/v1/messages"
MODEL = "claude-sonnet-4-6"
SIDECAR_DB = os.getenv("MC_CLASS_DB", "/data/mc_classification.db")
BATCH_SIZE = 25
SLEEP_BETWEEN_BATCHES = 0.5   # sec — keep gentle for the parallel Haiku batch

SYSTEM_PROMPT = """Du klassifizierst Master-Controls (MCs) der DSGVO/TMG/TDDDG-Compliance-Pruefung in genau drei Klassen:

TEXT  — Der MC prueft, ob ein bestimmter Inhalt im TEXT eines Dokuments (DSE, Impressum, Cookie-Richtlinie etc.) steht.
        Beispiele: "Enthaelt die Datenschutzerklaerung die Aufsichtsbehoerde?" / "Wird im Impressum die USt-IdNr genannt?"
        Diese MCs koennen gegen den Dokument-Text gematched werden.

PROCESS — Der MC prueft eine technische oder organisatorische MASSNAHME, die ausserhalb des Dokumenttexts liegt.
          Beispiele: "Ist sichergestellt, dass Cookies erst nach Einwilligung gespeichert werden?" /
                    "Ist ein Loeschkonzept implementiert?" / "Wird das Consent-Tool monatlich getestet?"
          Diese MCs lassen sich NICHT aus dem Dokumenttext beantworten — sie brauchen Evidence/TOM-Nachweis.

REVIEW — Der MC stellt eine abstrakte Vollstaendigkeits-/Konsistenzfrage, die menschliche Beurteilung braucht.
         Beispiele: "Sind ALLE Verarbeitungszwecke vollstaendig erfasst?" / "Stimmen Cookie-Richtlinie und Banner ueberein?"
         Diese MCs sind Checklisten-Items fuer den DSB, kein automatischer Check moeglich.

Antworte ausschliesslich als JSON-Array — keine Erklaerung davor/danach, kein Markdown-Codeblock. Format:
[{"control_id": "<wie input>", "check_type": "text|process|review", "confidence": 0.9, "rationale": "ein kurzer satz"}, ...]"""


def fetch_mcs(conn, limit: int | None = None, only_unclassified: bool = True) -> list[dict]:
    sql = """SELECT control_id, doc_type, title, check_question
             FROM compliance.doc_check_controls"""
    if only_unclassified:
        sql += " WHERE control_id NOT IN (SELECT control_id FROM _classified_already)"
    sql += " ORDER BY doc_type, title"
    if limit:
        sql += f" LIMIT {limit}"
    with conn.cursor(cursor_factory=RealDictCursor) as c:
        try:
            c.execute("CREATE TEMP TABLE _classified_already (control_id TEXT)")
            with sqlite3.connect(SIDECAR_DB) as side:
                rows = side.execute("SELECT control_id FROM mc_classification").fetchall()
                if rows:
                    c.executemany("INSERT INTO _classified_already VALUES (%s)", [(r[0],) for r in rows])
        except Exception:
            pass
        c.execute(sql)
        return list(c.fetchall())


def call_claude(api_key: str, batch: list[dict]) -> list[dict]:
    payload = {
        "model": MODEL,
        "max_tokens": 4000,
        "system": SYSTEM_PROMPT,
        "messages": [{
            "role": "user",
            "content": "Klassifiziere die folgenden MCs:\n\n" + json.dumps(
                [{"control_id": m["control_id"],
                  "doc_type": m["doc_type"],
                  "title": m["title"],
                  "check_question": (m["check_question"] or "")[:400]}
                 for m in batch],
                ensure_ascii=False, indent=2),
        }],
    }
    headers = {
        "x-api-key": api_key,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json",
    }
    r = httpx.post(ANTHROPIC_URL, json=payload, headers=headers, timeout=120.0)
    r.raise_for_status()
    txt = r.json()["content"][0]["text"].strip()
    # Strip code fences if Sonnet adds them
    if txt.startswith("```"):
        txt = txt.split("\n", 1)[1].rsplit("```", 1)[0].strip()
        if txt.startswith("json"):
            txt = txt[4:].strip()
    return json.loads(txt)


def ensure_sidecar() -> None:
    Path(SIDECAR_DB).parent.mkdir(parents=True, exist_ok=True)
    with sqlite3.connect(SIDECAR_DB) as c:
        c.executescript("""
            CREATE TABLE IF NOT EXISTS mc_classification (
                control_id    TEXT PRIMARY KEY,
                doc_type      TEXT,
                title         TEXT,
                check_type    TEXT,
                confidence    REAL,
                rationale     TEXT,
                classified_at TEXT
            );
            CREATE INDEX IF NOT EXISTS idx_doctype ON mc_classification(doc_type);
            CREATE INDEX IF NOT EXISTS idx_type    ON mc_classification(check_type);
        """)


def store_results(rows: list[dict], lookup: dict[str, dict]) -> None:
    ts = datetime.now(timezone.utc).isoformat()
    with sqlite3.connect(SIDECAR_DB) as c:
        c.executemany(
            "INSERT OR REPLACE INTO mc_classification "
            "(control_id, doc_type, title, check_type, confidence, rationale, classified_at) "
            "VALUES (?, ?, ?, ?, ?, ?, ?)",
            [
                (
                    r.get("control_id"),
                    lookup.get(r.get("control_id"), {}).get("doc_type", ""),
                    lookup.get(r.get("control_id"), {}).get("title", ""),
                    (r.get("check_type") or "").lower(),
                    float(r.get("confidence") or 0),
                    (r.get("rationale") or "")[:500],
                    ts,
                )
                for r in rows
            ],
        )
        c.commit()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit", type=int, default=None, help="limit number of MCs (for testing)")
    ap.add_argument("--sample", action="store_true", help="just print first 5 MCs and exit")
    ap.add_argument("--all", action="store_true", help="re-classify all (otherwise skips already classified)")
    args = ap.parse_args()

    ensure_sidecar()
    api_key = os.environ["ANTHROPIC_API_KEY"]
    conn = psycopg2.connect(os.environ["DATABASE_URL"])
    mcs = fetch_mcs(conn, limit=args.limit, only_unclassified=not args.all)

    if args.sample:
        for m in mcs[:5]:
            print(json.dumps(m, ensure_ascii=False, indent=2))
        return

    print(f"Klassifiziere {len(mcs)} MCs in Batches von {BATCH_SIZE}", flush=True)
    if not mcs:
        print("Nichts zu tun.")
        return

    lookup = {m["control_id"]: m for m in mcs}
    total = len(mcs)
    done = 0
    failed_batches = 0
    t0 = time.time()
    for i in range(0, total, BATCH_SIZE):
        batch = mcs[i:i + BATCH_SIZE]
        try:
            out = call_claude(api_key, batch)
            store_results(out, lookup)
            done += len(out)
            elapsed = time.time() - t0
            rate = done / max(elapsed, 0.01)
            eta = (total - done) / max(rate, 0.01)
            print(f"  [{done:>5}/{total}] {rate:.1f} MC/s  ETA {eta/60:.1f}min",
                  flush=True)
        except Exception as e:
            failed_batches += 1
            print(f"  FAIL batch {i//BATCH_SIZE + 1}: {e}", flush=True, file=sys.stderr)
            if failed_batches >= 5:
                print("  Zu viele Fehler — abbrechen.", file=sys.stderr)
                break
        time.sleep(SLEEP_BETWEEN_BATCHES)

    print(f"\nFertig: {done}/{total} klassifiziert ({failed_batches} Fehlbatches).")
    # Summary
    with sqlite3.connect(SIDECAR_DB) as c:
        c.row_factory = sqlite3.Row
        rows = c.execute(
            "SELECT doc_type, check_type, COUNT(*) n FROM mc_classification "
            "GROUP BY doc_type, check_type ORDER BY doc_type, check_type"
        ).fetchall()
        print("\n=== Verteilung nach doc_type x check_type ===")
        prev = None
        for r in rows:
            if r["doc_type"] != prev:
                print(); print(f"[{r['doc_type']}]")
                prev = r["doc_type"]
            print(f"  {r['check_type']:<8} {r['n']}")


if __name__ == "__main__":
    main()
