from __future__ import annotations

from collections import Counter
from typing import Any

SEVERITY_MAP = {"low": 1, "medium": 3, "high": 5}

CANONICAL_ATTACK_CATEGORIES = (
    "instruction_override",
    "prompt_leakage",
    "retrieved_indirect",
    "evasion_obfuscation",
    "adaptive_multiturn",
    "benign_control",
)

SUCCESS_DEFINITION_IDS = (
    "s1_instruction_override",
    "s2_prompt_leakage",
    "s3_retrieved_indirect",
    "s4_evasion_obfuscation",
    "s5_adaptive_multiturn",
    "s6_benign_control",
)

MAPPING_TABLE: dict[str, tuple[str, str, str]] = {
    "retrieved_content_injection": ("retrieved_indirect", "s3_retrieved_indirect", "high"),
    "benign_control": ("benign_control", "s6_benign_control", "high"),
    "system_prompt_extraction": ("prompt_leakage", "s2_prompt_leakage", "high"),
    "prompt_leakage": ("prompt_leakage", "s2_prompt_leakage", "high"),
    "direct_instruction_override": ("instruction_override", "s1_instruction_override", "high"),
    "instruction_override": ("instruction_override", "s1_instruction_override", "high"),
    "compound_instruction_attack": ("instruction_override", "s1_instruction_override", "medium"),
    "sandwich_defense_bypass": ("evasion_obfuscation", "s4_evasion_obfuscation", "high"),
    "blacklist_evasion": ("evasion_obfuscation", "s4_evasion_obfuscation", "high"),
    "xml_escape_evasion": ("evasion_obfuscation", "s4_evasion_obfuscation", "high"),
    "restricted_character_bypass": ("evasion_obfuscation", "s4_evasion_obfuscation", "high"),
    "emoji_only_bypass": ("evasion_obfuscation", "s4_evasion_obfuscation", "high"),
    "adaptive_attack": ("adaptive_multiturn", "s5_adaptive_multiturn", "high"),
}

EVALUATION_MODE_BY_CATEGORY = {
    "instruction_override": "hybrid",
    "prompt_leakage": "rule-based",
    "retrieved_indirect": "hybrid",
    "evasion_obfuscation": "rule-based",
    "adaptive_multiturn": "hybrid",
    "benign_control": "hybrid",
}


def _map_category(case: dict[str, Any]) -> tuple[str, str, str]:
    family = str(case.get("attack_family", "")).lower()
    subtype = str(case.get("attack_subtype", "")).lower()
    surface = str(case.get("attack_surface", "")).lower()
    label = str(case.get("label", "")).lower()

    if family in MAPPING_TABLE:
        return MAPPING_TABLE[family]
    if subtype == "hijacking":
        return ("instruction_override", "s1_instruction_override", "high")
    if subtype == "extraction":
        return ("prompt_leakage", "s2_prompt_leakage", "high")
    if surface == "tool_mediated":
        return ("retrieved_indirect", "s3_retrieved_indirect", "medium")
    if label == "prompt_leakage":
        return ("prompt_leakage", "s2_prompt_leakage", "medium")
    if label == "instruction_override":
        return ("instruction_override", "s1_instruction_override", "medium")
    return ("instruction_override", "s1_instruction_override", "low")


def _map_numeric_severity(case: dict[str, Any], confidence: str) -> tuple[int | None, str]:
    severity = case.get("severity")
    if isinstance(severity, int):
        return severity, confidence

    severity_level = case.get("severity_level")
    if isinstance(severity_level, str):
        mapped = SEVERITY_MAP.get(severity_level.lower())
        if mapped is not None:
            return mapped, confidence
    return None, "low" if confidence != "low" else confidence


def map_case(case: dict[str, Any]) -> dict[str, Any]:
    enriched = dict(case)
    category, success_definition_id, confidence = _map_category(case)
    severity, confidence = _map_numeric_severity(case, confidence)

    enriched["attack_category"] = category
    enriched["success_definition_id"] = success_definition_id
    enriched["mapping_confidence"] = confidence
    enriched["severity"] = severity
    enriched["evaluation_mode"] = EVALUATION_MODE_BY_CATEGORY.get(category, "hybrid")

    provenance = dict(enriched.get("provenance") or {})
    if "source_file" not in provenance and enriched.get("source_reference"):
        provenance["source_file"] = enriched["source_reference"]
    if "dataset_split" not in provenance and enriched.get("benchmark_split"):
        provenance["dataset_split"] = enriched["benchmark_split"]
    if provenance:
        enriched["provenance"] = provenance

    if confidence in {"low", "medium"}:
        enriched["mapping_audit"] = {
            "source_field_used": (
                enriched.get("attack_family")
                or enriched.get("attack_surface")
                or enriched.get("label")
                or ""
            ),
            "notes": "auto-mapped, review recommended",
        }

    return enriched


def summarize_mappings(cases: list[dict[str, Any]]) -> dict[str, dict[str, int]]:
    categories = Counter()
    confidences = Counter()
    sources = Counter()

    for case in cases:
        enriched = map_case(case)
        categories[enriched["attack_category"]] += 1
        confidences[enriched["mapping_confidence"]] += 1
        sources[str(enriched.get("source_name", "unknown"))] += 1

    return {
        "attack_categories": dict(categories),
        "mapping_confidence": dict(confidences),
        "source_names": dict(sources),
    }