"""
v2: Re-classify only the (control_id, doc_type) PAIRS that aren't covered yet.

V1 used PK=control_id, so cross-doc-type variants (same control assigned
to e.g. AGB AND Widerruf with different check_questions) overwrote each
other. v2 migrates to PK=(control_id, doc_type) and classifies only the
~262 missing pairs.

Run from container:
  docker exec bp-compliance-backend python /app/scripts/classify_mc_check_type_v2.py [--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

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 steht.
        Beispiele: "Enthaelt die Datenschutzerklaerung die Aufsichtsbehoerde?"
        Diese MCs koennen gegen den Dokument-Text gematched werden.

PROCESS — Der MC prueft eine technische oder organisatorische MASSNAHME ausserhalb des Dokumenttexts.
          Beispiele: "Ist sichergestellt, dass Cookies erst nach Einwilligung gespeichert werden?"

REVIEW — Der MC stellt eine abstrakte Vollstaendigkeitsfrage, die menschliche Beurteilung braucht.
         Beispiele: "Sind ALLE Verarbeitungszwecke vollstaendig erfasst?"

Wichtig: derselbe Control kann fuer DIFFERENT doc_types unterschiedlich klassifiziert sein —
mit doc_type-spezifisch angepasster check_question kann ein "text"-Check fuer ein Dok zum
"process"-Check fuer ein anderes werden.

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


def migrate_schema() -> None:
    """Migrate sidecar from PK=control_id to PK=(control_id, doc_type)."""
    Path(SIDECAR_DB).parent.mkdir(parents=True, exist_ok=True)
    with sqlite3.connect(SIDECAR_DB) as c:
        # Check if v2 schema already in place (composite PK)
        cols = c.execute("PRAGMA table_info(mc_classification)").fetchall()
        if not cols:
            # First run — create fresh
            c.executescript("""
                CREATE TABLE mc_classification (
                    control_id    TEXT,
                    doc_type      TEXT,
                    title         TEXT,
                    check_type    TEXT,
                    confidence    REAL,
                    rationale     TEXT,
                    classified_at TEXT,
                    PRIMARY KEY (control_id, doc_type)
                );
                CREATE INDEX idx_doctype ON mc_classification(doc_type);
                CREATE INDEX idx_type    ON mc_classification(check_type);
            """)
            return

        # Check whether the existing table already has composite PK
        pk_cols = [r[1] for r in cols if r[5] > 0]
        if set(pk_cols) == {"control_id", "doc_type"}:
            print("Schema already v2 (composite PK). Skipping migration.")
            return

        print("Migrating sidecar schema to PK(control_id, doc_type)...")
        c.executescript("""
            CREATE TABLE mc_classification_v2 (
                control_id    TEXT,
                doc_type      TEXT,
                title         TEXT,
                check_type    TEXT,
                confidence    REAL,
                rationale     TEXT,
                classified_at TEXT,
                PRIMARY KEY (control_id, doc_type)
            );
            INSERT INTO mc_classification_v2
              (control_id, doc_type, title, check_type, confidence, rationale, classified_at)
            SELECT control_id, COALESCE(doc_type, ''), title, check_type, confidence, rationale, classified_at
            FROM mc_classification;
            DROP TABLE mc_classification;
            ALTER TABLE mc_classification_v2 RENAME TO mc_classification;
            CREATE INDEX idx_doctype ON mc_classification(doc_type);
            CREATE INDEX idx_type ON mc_classification(check_type);
        """)
        n = c.execute("SELECT COUNT(*) FROM mc_classification").fetchone()[0]
        print(f"Migrated {n} existing rows.")


def fetch_unclassified_pairs(conn) -> list[dict]:
    """All (control_id, doc_type) pairs in PG that aren't yet in sidecar."""
    side_pairs: set[tuple[str, str]] = set()
    with sqlite3.connect(SIDECAR_DB) as side:
        for cid, dt in side.execute("SELECT control_id, doc_type FROM mc_classification"):
            side_pairs.add((cid, dt or ""))

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

    missing = [r for r in all_rows if (r["control_id"], r["doc_type"] or "") not in side_pairs]
    return missing


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()
    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_results(rows: list[dict], lookup: dict[tuple[str, 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"),
                    r.get("doc_type") or "",
                    lookup.get((r.get("control_id"), r.get("doc_type") or ""), {}).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("--sample", action="store_true")
    args = ap.parse_args()

    migrate_schema()
    api_key = os.environ["ANTHROPIC_API_KEY"]
    conn = psycopg2.connect(os.environ["DATABASE_URL"])
    missing = fetch_unclassified_pairs(conn)

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

    print(f"Klassifiziere {len(missing)} fehlende (control_id, doc_type)-Paare in Batches von {BATCH_SIZE}", flush=True)
    if not missing:
        print("Alles klassifiziert. Nichts zu tun.")
        return

    lookup = {(m["control_id"], m["doc_type"] or ""): m for m in missing}
    total = len(missing)
    done = 0
    failed_batches = 0
    t0 = time.time()
    for i in range(0, total, BATCH_SIZE):
        batch = missing[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:>4}/{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).")
    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=== Final-Verteilung 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()
