"""
LLM-Audit: prueft fuer jeden text-MC ob er fachlich wirklich zu seinem
zugeordneten doc_type passt. BMW Run zeigte z.B. dass die meisten
Impressum-MCs eigentlich DSA/E-Commerce-Pflichten sind, die in einem
§5-TMG-Impressum gar nicht stehen.

Output:
- doc_type passt → MC bleibt active (kein DB-Update)
- doc_type passt NICHT → check_type wird auf 'misclassified' gesetzt;
  rag_document_checker filtert die dann aus

Sidecar SQLite. KEINE Postgres-Aenderung (Schema frozen).
"""

from __future__ import annotations

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

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

DOC_TYPE_DESCRIPTIONS = {
    "agb": "Allgemeine Geschaeftsbedingungen — Vertragsbedingungen "
           "zwischen Anbieter und Kunde",
    "avv": "Auftragsverarbeitungsvertrag Art. 28 DSGVO — Vertrag zwischen "
           "Verantwortlichem und Auftragsverarbeiter",
    "cookie": "Cookie-Richtlinie — Information ueber Zwecke, Anbieter, "
              "Speicherdauer und Widerruf von Cookies/Tracking-Pixeln",
    "dse": "Datenschutzerklaerung Art. 13 DSGVO — Informationen ueber "
           "Verantwortlichen, Zwecke, Rechtsgrundlagen, Empfaenger, "
           "Speicherdauer, Betroffenenrechte, Aufsichtsbehoerde",
    "dsfa": "Datenschutz-Folgenabschaetzung Art. 35 DSGVO — Risikobewertung "
            "von Verarbeitungen mit hohem Risiko",
    "impressum": "Impressum §5 TMG / §18 MStV — Anbieterkennzeichnung mit "
                 "Name, Anschrift, Vertretungsberechtigten, Handelsregister, "
                 "USt-IdNr., berufsrechtliche Angaben, Aufsicht",
    "loeschkonzept": "Loeschkonzept DIN 66398 — Definition Aufbewahrungs- "
                     "und Loeschfristen pro Datenkategorie + Prozess",
    "widerruf": "Widerrufsbelehrung §312 BGB — Verbraucher-Widerrufsrecht "
                "bei Fernabsatz, Frist, Folgen, Muster",
}

SYSTEM_PROMPT = """Du pruefst Master-Controls (MCs) der Compliance-Pruefung.

Fuer jeden MC bekommst du:
- den ihm zugeordneten doc_type (z.B. 'impressum', 'cookie', 'dse')
- den Titel und die check_question

Deine Aufgabe: entscheide ob der MC fachlich wirklich zu diesem doc_type passt.

Beispiele:
- MC "Wird die USt-IdNr im Impressum genannt?" + doc_type=impressum → PASST
- MC "Monatlich aktive Endnutzer fuer Online-Vermittlungsdienste messen" + doc_type=impressum → PASST NICHT
  (DSA-Pflicht, gehoert nicht in ein klassisches §5-TMG-Impressum)
- MC "Mithoeren von Funknachrichten verhindern" + doc_type=cookie → PASST NICHT
  (TKG-Spezialthema, nicht Cookie-Richtlinie)

Antworte als JSON-Array, eine Zeile pro MC:
[{"control_id": "<wie input>", "doc_type": "<wie input>", "fits": true|false,
  "rationale": "ein kurzer satz"}, ...]
Kein Markdown."""


def fetch_pairs_to_audit(conn) -> list[dict]:
    """All text-MCs that haven't been audited yet (no 'fits' column)."""
    with sqlite3.connect(SIDECAR_DB) as side:
        cols = [r[1] for r in side.execute("PRAGMA table_info(mc_classification)")]
        if "fits_doc_type" not in cols:
            side.execute("ALTER TABLE mc_classification ADD COLUMN fits_doc_type INTEGER")
            side.commit()
        already = set()
        for cid, dt in side.execute(
            "SELECT control_id, doc_type FROM mc_classification "
            "WHERE fits_doc_type IS NOT NULL"
        ):
            already.add((cid, dt or ""))

    with conn.cursor(cursor_factory=RealDictCursor) as c:
        c.execute("""SELECT dc.control_id, dc.doc_type, dc.title, dc.check_question
                     FROM compliance.doc_check_controls dc
                     WHERE dc.control_id IN (
                       SELECT control_id FROM compliance.doc_check_controls
                     )""")
        all_rows = list(c.fetchall())

    # Audit only those classified as 'text' in sidecar — process/review
    # never run through doc_check anyway
    with sqlite3.connect(SIDECAR_DB) as side:
        text_pairs = set()
        for cid, dt in side.execute(
            "SELECT control_id, doc_type FROM mc_classification "
            "WHERE check_type = 'text'"
        ):
            text_pairs.add((cid, dt or ""))

    target = [r for r in all_rows
              if (r["control_id"], r["doc_type"] or "") in text_pairs
              and (r["control_id"], r["doc_type"] or "") not in already]
    return target


def call_claude(api_key: str, batch: list[dict]) -> list[dict]:
    payload = {
        "model": MODEL,
        "max_tokens": 4000,
        "system": SYSTEM_PROMPT,
        "messages": [{
            "role": "user",
            "content": (
                "Doc-Typen-Beschreibungen:\n"
                + "\n".join(f"- {k}: {v}" for k, v in DOC_TYPE_DESCRIPTIONS.items())
                + "\n\nPruefe folgende MCs:\n\n"
                + json.dumps([
                    {"control_id": m["control_id"], "doc_type": m["doc_type"],
                     "title": m["title"], "check_question": (m["check_question"] or "")[:300]}
                    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()
    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 store_audit(rows: list[dict]) -> None:
    ts = datetime.now(timezone.utc).isoformat()
    with sqlite3.connect(SIDECAR_DB) as c:
        c.executemany(
            "UPDATE mc_classification SET fits_doc_type = ?, "
            "rationale = COALESCE(?, rationale), classified_at = ? "
            "WHERE control_id = ? AND doc_type = ?",
            [
                (
                    1 if r.get("fits") else 0,
                    (r.get("rationale") or "")[:500] or None,
                    ts,
                    r.get("control_id"),
                    r.get("doc_type") or "",
                )
                for r in rows
            ],
        )
        c.commit()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--sample", action="store_true")
    args = ap.parse_args()

    api_key = os.environ["ANTHROPIC_API_KEY"]
    conn = psycopg2.connect(os.environ["DATABASE_URL"])
    pairs = fetch_pairs_to_audit(conn)

    if args.sample:
        for m in pairs[:5]:
            print(json.dumps(m, ensure_ascii=False, indent=2))
        print(f"\nTotal pairs to audit: {len(pairs)}")
        return

    print(f"Audit {len(pairs)} text-MCs in Batches von {BATCH_SIZE}", flush=True)
    if not pairs:
        print("Alles auditiert.")
        return

    done = 0
    failed_batches = 0
    t0 = time.time()
    for i in range(0, len(pairs), BATCH_SIZE):
        batch = pairs[i:i + BATCH_SIZE]
        try:
            out = call_claude(api_key, batch)
            store_audit(out)
            done += len(out)
            elapsed = time.time() - t0
            rate = done / max(elapsed, 0.01)
            eta = (len(pairs) - done) / max(rate, 0.01)
            print(f"  [{done:>4}/{len(pairs)}] {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}/{len(pairs)} auditiert ({failed_batches} Fehlbatches).")
    with sqlite3.connect(SIDECAR_DB) as c:
        c.row_factory = sqlite3.Row
        rows = c.execute(
            "SELECT doc_type, "
            "  SUM(CASE WHEN fits_doc_type = 1 THEN 1 ELSE 0 END) AS fits, "
            "  SUM(CASE WHEN fits_doc_type = 0 THEN 1 ELSE 0 END) AS misfits, "
            "  COUNT(*) AS total "
            "FROM mc_classification "
            "WHERE check_type = 'text' AND fits_doc_type IS NOT NULL "
            "GROUP BY doc_type ORDER BY doc_type"
        ).fetchall()
        print("\n=== Audit-Verteilung doc_type x fits ===")
        for r in rows:
            print(f"  {r['doc_type']:<14}  fits={r['fits']:<4} misfits={r['misfits']:<4} total={r['total']}")


if __name__ == "__main__":
    main()
