#!/usr/bin/env python3
"""
Fix generic tokens: Re-classify controls that were assigned to
action-based tokens (documentation, procedure, process, etc.)
instead of topic-based tokens.

Runs sequentially in 5 batches. NO retry on timeout.

Usage:
  python3 /app/scripts/gpre0_fix_generic_tokens.py --dry-run
  python3 /app/scripts/gpre0_fix_generic_tokens.py
"""

import argparse
import json
import logging
import os
import time
from collections import defaultdict
from pathlib import Path

import httpx
from sqlalchemy import create_engine, text

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s"
)
logger = logging.getLogger("gpre0-fix-generic")

DB_URL = os.getenv(
    "DATABASE_URL",
    "postgresql://breakpilot:breakpilot123@postgres:5432/breakpilot_db",
)
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY", "")
ANTHROPIC_MODEL = "claude-haiku-4-5-20251001"
ANTHROPIC_URL = "https://api.anthropic.com/v1/messages"
CHECKPOINT_DIR = Path("/tmp/gpre0_fix_checkpoints")

# Tokens that are ACTION-based, not TOPIC-based → must be re-classified
FORBIDDEN_TOKENS = {
    "documentation", "procedure", "process",
    "compliance_reporting", "records_management",
}

SYSTEM_PROMPT = """Du bist ein Compliance-Klassifizierer. Ordne jeden Control dem THEMA zu, nicht der Handlung.

KRITISCH: Die Tokens "documentation", "procedure", "process", "compliance_reporting",
"records_management" sind VERBOTEN. Klassifiziere nach dem INHALTLICHEN THEMA.

Beispiele:
- "Risikobewertung dokumentieren" → risk_management (NICHT documentation)
- "Incident-Verfahren definieren" → incident (NICHT procedure)
- "Verschlüsselungsprozess implementieren" → encryption (NICHT process)
- "Audit-Ergebnisse berichten" → compliance_audit (NICHT compliance_reporting)
- "Datenschutz-Unterlagen verwalten" → personal_data (NICHT records_management)
- "Löschkonzept dokumentieren" → data_retention (NICHT documentation)
- "Zertifizierungsverfahren definieren" → certification (NICHT procedure)
- "Schulungsprozess durchführen" → training (NICHT process)

ERLAUBTE TOKENS:

SECURITY: multi_factor_auth, password_policy, credentials, session_management,
privileged_access, access_control, encryption, transport_encryption,
key_management, certificate_management, network_security, network_segmentation,
firewall, vpn, remote_access, monitoring, audit_logging, siem, alerting,
compliance_audit, vulnerability, patch_management, backup, disaster_recovery,
physical_security, secure_development, api_security, input_validation,
container_security, logging_configuration

DATA_PROTECTION: personal_data, sensitive_data, health_data, consent,
data_subject_rights, data_retention, data_transfer, data_breach_notification,
dpia, data_processing_agreement, privacy_by_design, data_processing_register,
data_classification, cookie_consent, video_surveillance

GOVERNANCE: policy, training, awareness, incident, risk_management,
third_party_management, change_management, asset_management,
human_resources_security

REGULATORY: supervisory_authority, certification, product_safety, ai_system,
financial_reporting, aml, whistleblowing, consumer_protection, ecommerce,
telecommunications, medical_device, payment_services, critical_infrastructure,
supply_chain_due_diligence, sustainability_reporting

Antworte NUR als JSON-Array: [{"id":"...","token":"...","conf":0.9}, ...]"""


def call_claude(controls_batch: list[dict]) -> tuple[list[dict], dict]:
    """Send batch to Claude. NO retry on timeout."""
    items = []
    for c in controls_batch:
        items.append(
            f'- id="{c["control_id"]}" '
            f'cur="{c["current_object"]}" '
            f't="{c["title"]}" '
            f'o="{c["objective"][:100]}"'
        )

    prompt = "Klassifiziere nach THEMA (nicht Handlung):\n" + "\n".join(items)

    headers = {
        "x-api-key": ANTHROPIC_API_KEY,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json",
    }
    payload = {
        "model": ANTHROPIC_MODEL,
        "max_tokens": 1500,
        "temperature": 0.0,
        "system": SYSTEM_PROMPT,
        "messages": [{"role": "user", "content": prompt}],
    }

    try:
        resp = httpx.post(
            ANTHROPIC_URL, headers=headers, json=payload, timeout=45.0
        )
        resp.raise_for_status()
        data = resp.json()
        content = data.get("content", [{}])[0].get("text", "")
        usage = data.get("usage", {})
        start = content.find("[")
        end = content.rfind("]") + 1
        if start >= 0 and end > start:
            return json.loads(content[start:end]), usage
        return [], usage
    except httpx.TimeoutException:
        logger.error("TIMEOUT — skipping batch")
        return [], {}
    except httpx.HTTPStatusError as e:
        if e.response.status_code == 429:
            logger.warning("Rate limited — waiting 60s")
            time.sleep(60)
        else:
            logger.error("API error %d", e.response.status_code)
        return [], {}
    except Exception as e:
        logger.error("Failed: %s", e)
        return [], {}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch-size", type=int, default=20)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    engine = create_engine(
        DB_URL, connect_args={"options": "-c search_path=compliance,public"}
    )

    # Load only controls with forbidden tokens
    forbidden_pattern = "|".join(
        f":{tok}:" for tok in FORBIDDEN_TOKENS
    )
    with engine.connect() as c:
        rows = c.execute(text("""
            SELECT cc.id, cc.control_id, cc.title,
                   COALESCE(cc.objective, '') as objective,
                   cc.generation_metadata->>'merge_group_hint' as hint
            FROM canonical_controls cc
            WHERE cc.generation_metadata->>'merge_group_hint' IS NOT NULL
              AND cc.release_state NOT IN ('deprecated', 'rejected')
              AND (
                cc.generation_metadata->>'merge_group_hint' LIKE '%:documentation:%'
                OR cc.generation_metadata->>'merge_group_hint' LIKE '%:procedure:%'
                OR cc.generation_metadata->>'merge_group_hint' LIKE '%:process:%'
                OR cc.generation_metadata->>'merge_group_hint' LIKE '%:compliance_reporting:%'
                OR cc.generation_metadata->>'merge_group_hint' LIKE '%:records_management:%'
              )
        """)).fetchall()

    controls = []
    for uuid, cid, title, objective, hint in rows:
        parts = hint.split(":", 2) if hint else []
        controls.append({
            "uuid": str(uuid), "control_id": cid,
            "title": title or "", "objective": objective or "",
            "current_hint": hint,
            "current_object": parts[1] if len(parts) > 1 else hint,
        })

    logger.info("Found %d controls with forbidden tokens to re-classify", len(controls))

    # Process
    total_fixed = 0
    total_kept = 0
    total_skipped = 0
    total_input_tokens = 0
    total_output_tokens = 0
    corrections = []
    change_stats: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))

    for i in range(0, len(controls), args.batch_size):
        batch = controls[i:i + args.batch_size]
        results, usage = call_claude(batch)

        total_input_tokens += usage.get("input_tokens", 0)
        total_output_tokens += usage.get("output_tokens", 0)

        if not results:
            total_skipped += len(batch)
            continue

        result_map = {r.get("id", ""): r for r in results}
        for ctrl in batch:
            r = result_map.get(ctrl["control_id"], {})
            new_token = r.get("token", "")
            if not new_token or new_token in FORBIDDEN_TOKENS:
                total_kept += 1
                continue

            old_obj = ctrl["current_object"]
            if new_token != old_obj:
                total_fixed += 1
                parts = ctrl["current_hint"].split(":", 2)
                action = parts[0] if parts else "implement"
                phase = parts[2] if len(parts) > 2 else "implementation"
                corrections.append({
                    "uuid": ctrl["uuid"],
                    "old_hint": ctrl["current_hint"],
                    "new_hint": f"{action}:{new_token}:{phase}",
                })
                change_stats[old_obj][new_token] += 1
            else:
                total_kept += 1

        processed = min(i + args.batch_size, len(controls))
        if processed % 2000 < args.batch_size or processed >= len(controls):
            logger.info(
                "Progress: %d/%d (fixed=%d kept=%d skip=%d)",
                processed, len(controls), total_fixed, total_kept, total_skipped,
            )

        time.sleep(0.3)

    # Report
    cost_in = total_input_tokens / 1_000_000 * 0.80
    cost_out = total_output_tokens / 1_000_000 * 4.00
    logger.info("\n" + "=" * 60)
    logger.info("GENERIC TOKEN FIX REPORT")
    logger.info("=" * 60)
    logger.info("Total:   %d controls", len(controls))
    logger.info("Fixed:   %d", total_fixed)
    logger.info("Kept:    %d (LLM also chose forbidden → kept as-is)", total_kept)
    logger.info("Skipped: %d", total_skipped)
    logger.info("Cost:    $%.2f (Haiku)", cost_in + cost_out)

    logger.info("\nTop changes:")
    flat = []
    for old, news in change_stats.items():
        for new, cnt in news.items():
            flat.append((cnt, old, new))
    for cnt, old, new in sorted(flat, reverse=True)[:30]:
        logger.info("  %4d × %s → %s", cnt, old, new)

    # Save corrections
    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
    corr_file = CHECKPOINT_DIR / "corrections_generic_fix.json"
    corr_file.write_text(json.dumps(corrections))
    logger.info("Saved %d corrections to %s", len(corrections), corr_file)

    if args.dry_run:
        logger.info("DRY RUN — not updating DB")
        return

    if corrections:
        logger.info("Applying %d corrections...", len(corrections))
        with engine.begin() as c:
            c.execute(text("SET search_path TO compliance, public"))
            for corr in corrections:
                c.execute(text("""
                    UPDATE canonical_controls
                    SET generation_metadata = jsonb_set(
                        generation_metadata,
                        '{merge_group_hint}',
                        to_jsonb(CAST(:new_hint AS text))
                    )
                    WHERE id = CAST(:uuid AS uuid)
                """), {"uuid": corr["uuid"], "new_hint": corr["new_hint"]})
        logger.info("Done. %d hints corrected.", len(corrections))


if __name__ == "__main__":
    main()
