#!/usr/bin/env python3
"""Fix bad L2 subtopics: stakeholder_*, escalation fragments, *_approval*, *_documentation."""

import json
import logging
import os
import time
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("fix-subtopics")

DB_URL = os.getenv("DATABASE_URL", "postgresql://breakpilot:breakpilot123@postgres:5432/breakpilot_db")
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY", "")
ANTHROPIC_URL = "https://api.anthropic.com/v1/messages"

SYSTEM_PROMPT = """Du klassifizierst Controls mit einem L1_L2 Token. Das L2 soll den KONKRETEN fachlichen Aspekt beschreiben.

VERBOTENE L2-Wörter (zu generisch):
- stakeholder (zu vage — WER sind die Stakeholder? WAS wird getan?)
- documentation (ist eine Handlung, kein Thema)
- approval (ist eine Handlung)
- communication (zu vage)

Stattdessen SPEZIFISCH:
- "stakeholder_notification" bei Behördenmeldung → "authority_reporting"
- "stakeholder_consultation" bei DSFA → "impact_consultation"
- "stakeholder_engagement" bei Training → "participant_selection"
- "escalation_procedure" → "severity_classification" oder "response_plan"
- "access_documentation" → "access_policy" oder "permission_matrix"
- "approval_process" → "authorization_workflow" oder "sign_off"

L2 = 1-3 Wörter, snake_case, FACHLICH SPEZIFISCH.

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


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

    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.release_state NOT IN ('deprecated', 'rejected')
              AND cc.generation_metadata->>'merge_group_hint' IS NOT NULL
              AND (
                cc.generation_metadata->>'merge_group_hint' LIKE '%stakeholder%'
                OR cc.generation_metadata->>'merge_group_hint' LIKE '%_escalation_%'
                OR cc.generation_metadata->>'merge_group_hint' LIKE '%_approval_%'
                OR cc.generation_metadata->>'merge_group_hint' LIKE '%response_time%'
                OR cc.generation_metadata->>'merge_group_hint' LIKE '%machine_re%'
                OR cc.generation_metadata->>'merge_group_hint' LIKE '%management_app%'
              )
        """)).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 "",
        })

    logger.info("Found %d controls with bad subtopics to fix", len(controls))

    headers = {
        "x-api-key": ANTHROPIC_API_KEY,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json",
    }

    corrections = []
    total_fixed = 0
    batch_size = 20

    for i in range(0, len(controls), batch_size):
        batch = controls[i:i + batch_size]
        items = [
            f'- id="{c["control_id"]}" cur="{c["current_object"]}" t="{c["title"]}" o="{c["objective"][:80]}"'
            for c in batch
        ]

        try:
            resp = httpx.post(ANTHROPIC_URL, headers=headers, json={
                "model": "claude-haiku-4-5-20251001",
                "max_tokens": 1500, "temperature": 0.0,
                "system": SYSTEM_PROMPT,
                "messages": [{"role": "user", "content": "Fix:\n" + "\n".join(items)}],
            }, timeout=45.0)
            resp.raise_for_status()
            content = resp.json().get("content", [{}])[0].get("text", "")
            start = content.find("[")
            end = content.rfind("]") + 1
            results = json.loads(content[start:end]) if start >= 0 else []
        except Exception as e:
            logger.error("Batch %d failed: %s", i, e)
            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 == ctrl["current_object"]:
                continue
            if "stakeholder" in new_token or "approval" in new_token:
                continue  # Still bad

            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}",
            })
            total_fixed += 1

        if (i + batch_size) % 200 < batch_size:
            logger.info("Progress: %d/%d (fixed=%d)", min(i + batch_size, len(controls)), len(controls), total_fixed)
        time.sleep(0.3)

    logger.info("Fixed: %d of %d controls", total_fixed, len(controls))

    # Save + apply
    Path("/tmp/corrections_bad_subtopics.json").write_text(json.dumps(corrections))

    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.")


if __name__ == "__main__":
    main()
