"""
A3-v2: nachschärfen welche MCs in Wahrheit auf das CMP-Banner-UI / einen
Prozess zielen, nicht auf den Doc-TEXT.

BMW-Run zeigte ~10 'Sprache klar/einfach' + 'Button-Text formuliert' MCs,
die der erste Audit als doc_type-passend markierte. Sie sind ABER nicht
gegen den Cookie-Policy- oder DSE-Text pruefbar — die fragen nach der
Verstaendlichkeit der Einwilligungs-UI.

Schaltet 'scope_requires' bei FRT/biometrischen MCs und re-klassifiziert
diese UI-Items auf check_type='process' (filtert sie aus doc_check raus).

Plus: Setzt scope_requires fuer offensichtlich konditionale MCs:
  - 'biometric_processing' bei FRT/Gesichtserkennung
  - 'ai_decision_making' bei automatisierten Einzelentscheidungen
  - 'child_targeting' bei Kinder-Einwilligungs-MCs
"""

from __future__ import annotations

import argparse
import json
import os
import sqlite3
import sys
import time

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 = 20

SYSTEM_PROMPT = """Du pruefst Master-Controls (MCs) der Compliance-Pruefung
zweiter Ordnung: jeder MC wurde bereits als 'text' klassifiziert und einem
doc_type zugeordnet. Du entscheidest:

A) Pruefst der MC den TEXT des Dokuments (z.B. "Wird im Impressum die
   USt-IdNr genannt?") → fits_doc_type=true, ui_only=false
B) Pruefst der MC die UI/UX eines Banners oder Buttons (z.B.
   "Ist der Ablehnungs-Button gleich gross wie Akzeptieren?", "Ist die
   Einwilligungsanfrage in einfacher Sprache formuliert?") → ui_only=true
   (Diese MCs koennen NIE im Dokumenttext stehen, weil sie sich auf eine
   externe UI beziehen.)

Zusaetzlich kennzeichne SCOPE-Bedingungen, wenn der MC nur fuer bestimmte
Sites relevant ist:
  - 'biometric_processing' : nur bei Sites die biometrische Daten
    (Gesichtserkennung/FRT, Fingerabdruecke) verarbeiten
  - 'ai_decision_making'   : nur bei automatisierten Einzelentscheidungen
    (Art. 22 DSGVO)
  - 'child_targeting'      : nur bei Sites die sich an Kinder richten
  - 'ecommerce'            : nur bei Webshops
  - 'b2c'                  : nur fuer Verbraucher-Geschaeft (UWG/BGB-Pflichten)
Wenn der MC fuer ALLE Sites gilt, setze scope_requires=null.

Antworte als JSON-Array — keine Erklaerung davor/danach, kein Markdown.
Format:
[{"control_id": "<wie input>", "doc_type": "<wie input>",
  "ui_only": true|false, "scope_requires": "biometric_processing|ai_decision_making|child_targeting|ecommerce|b2c|null",
  "rationale": "ein kurzer satz"}, ...]"""


def fetch_pairs_to_audit(conn) -> list[dict]:
    """All text-MCs that haven't been v2-audited yet (no 'ui_only' column)."""
    with sqlite3.connect(SIDECAR_DB) as side:
        cols = [r[1] for r in side.execute("PRAGMA table_info(mc_classification)")]
        added = False
        if "ui_only" not in cols:
            side.execute("ALTER TABLE mc_classification ADD COLUMN ui_only INTEGER")
            added = True
        if "scope_requires" not in cols:
            side.execute("ALTER TABLE mc_classification ADD COLUMN scope_requires TEXT")
            added = True
        if added:
            side.commit()
        already = set()
        for cid, dt in side.execute(
            "SELECT control_id, doc_type FROM mc_classification "
            "WHERE ui_only 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""")
        all_rows = list(c.fetchall())

    # Audit only those already classified as text+fits in sidecar
    with sqlite3.connect(SIDECAR_DB) as side:
        eligible = set()
        for cid, dt in side.execute(
            "SELECT control_id, doc_type FROM mc_classification "
            "WHERE check_type = 'text' AND (fits_doc_type IS NULL OR fits_doc_type = 1)"
        ):
            eligible.add((cid, dt or ""))

    target = [r for r in all_rows
              if (r["control_id"], r["doc_type"] or "") in eligible
              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": "Pruefe 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(rows: list[dict]) -> None:
    with sqlite3.connect(SIDECAR_DB) as c:
        c.executemany(
            "UPDATE mc_classification SET ui_only = ?, scope_requires = ? "
            "WHERE control_id = ? AND doc_type = ?",
            [
                (
                    1 if r.get("ui_only") else 0,
                    (r.get("scope_requires") or "").strip() or None
                       if (r.get("scope_requires") or "").lower() not in ("", "null")
                       else None,
                    r.get("control_id"),
                    r.get("doc_type") or "",
                )
                for r in rows
            ],
        )
        # MCs flagged ui_only become check_type='process' so they're not in doc_check
        c.executemany(
            "UPDATE mc_classification SET check_type='process' "
            "WHERE ui_only=1 AND control_id=? AND doc_type=?",
            [(r.get("control_id"), r.get("doc_type") or "") for r in rows
             if r.get("ui_only")],
        )
        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: {len(pairs)}")
        return

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

    done = 0
    fail = 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(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:
            fail += 1
            print(f"  FAIL batch {i//BATCH_SIZE + 1}: {e}", file=sys.stderr, flush=True)
            if fail >= 5: break
        time.sleep(0.5)

    print(f"\nFertig: {done}/{len(pairs)} ({fail} Fehlbatches).")
    with sqlite3.connect(SIDECAR_DB) as c:
        ui = c.execute("SELECT COUNT(*) FROM mc_classification WHERE ui_only=1").fetchone()[0]
        scope = c.execute(
            "SELECT scope_requires, COUNT(*) FROM mc_classification "
            "WHERE scope_requires IS NOT NULL GROUP BY scope_requires"
        ).fetchall()
        print(f"\nui_only=1: {ui} MCs (jetzt 'process')")
        print("scope_requires Verteilung:")
        for s, n in scope:
            print(f"  {s}: {n}")


if __name__ == "__main__":
    main()
