"""Vulnerability registry — single source of truth for vulnerability definitions and framework mappings.

This module defines all vulnerability entries, builds inverted indexes for
category-to-vulnerability resolution, and provides resolver functions used
throughout the pipeline.
"""

from __future__ import annotations

import types
from collections.abc import Mapping

from evaluatorq.redteam.contracts import (
    AttackTechnique,
    Vulnerability,
    VulnerabilityDef,
    VulnerabilityDomain,
)

__all__ = ["VULNERABILITY_DEFS", "Vulnerability", "VulnerabilityDef"]

# ---------------------------------------------------------------------------
# Vulnerability definitions
# ---------------------------------------------------------------------------

VULNERABILITY_DEFS: Mapping[Vulnerability, VulnerabilityDef] = {
    # ── Agentic Security (OWASP ASI) ─────────────────────────────────
    Vulnerability.GOAL_HIJACKING: VulnerabilityDef(
        id=Vulnerability.GOAL_HIJACKING,
        name='Agent Goal Hijacking',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.INDIRECT_INJECTION,
        framework_mappings={'OWASP-ASI': ['ASI01']},
    ),
    Vulnerability.TOOL_MISUSE: VulnerabilityDef(
        id=Vulnerability.TOOL_MISUSE,
        name='Tool Misuse & Exploitation',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.TOOL_ABUSE,
        framework_mappings={'OWASP-ASI': ['ASI02']},
    ),
    Vulnerability.IDENTITY_PRIVILEGE_ABUSE: VulnerabilityDef(
        id=Vulnerability.IDENTITY_PRIVILEGE_ABUSE,
        name='Identity & Privilege Abuse',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.PRIVILEGE_ESCALATION,
        framework_mappings={'OWASP-ASI': ['ASI03']},
    ),
    Vulnerability.SUPPLY_CHAIN: VulnerabilityDef(
        id=Vulnerability.SUPPLY_CHAIN,
        name='Supply Chain Vulnerabilities',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.SUPPLY_CHAIN,
        framework_mappings={'OWASP-ASI': ['ASI04'], 'OWASP-LLM': ['LLM03']},
    ),
    Vulnerability.CODE_EXECUTION: VulnerabilityDef(
        id=Vulnerability.CODE_EXECUTION,
        name='Unexpected Code Execution',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.RCE_EXPLOIT,
        framework_mappings={'OWASP-ASI': ['ASI05']},
    ),
    Vulnerability.MEMORY_POISONING: VulnerabilityDef(
        id=Vulnerability.MEMORY_POISONING,
        name='Memory & Context Poisoning',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.CONTEXT_POISONING,
        framework_mappings={'OWASP-ASI': ['ASI06']},
    ),
    Vulnerability.INTER_AGENT_COMMS: VulnerabilityDef(
        id=Vulnerability.INTER_AGENT_COMMS,
        name='Insecure Inter-Agent Communication',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.MESSAGE_SPOOFING,
        framework_mappings={'OWASP-ASI': ['ASI07']},
    ),
    Vulnerability.CASCADING_FAILURES: VulnerabilityDef(
        id=Vulnerability.CASCADING_FAILURES,
        name='Cascading Failures',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.CASCADE_TRIGGER,
        framework_mappings={'OWASP-ASI': ['ASI08']},
    ),
    Vulnerability.TRUST_EXPLOITATION: VulnerabilityDef(
        id=Vulnerability.TRUST_EXPLOITATION,
        name='Human-Agent Trust Exploitation',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.SOCIAL_ENGINEERING,
        framework_mappings={'OWASP-ASI': ['ASI09']},
    ),
    Vulnerability.ROGUE_AGENTS: VulnerabilityDef(
        id=Vulnerability.ROGUE_AGENTS,
        name='Rogue Agents',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.BOUNDARY_VIOLATION,
        framework_mappings={'OWASP-ASI': ['ASI10']},
    ),
    # ── LLM Security (OWASP LLM Top 10) ──────────────────────────────
    Vulnerability.PROMPT_INJECTION: VulnerabilityDef(
        id=Vulnerability.PROMPT_INJECTION,
        name='Prompt Injection',
        domain=VulnerabilityDomain.MODEL,
        default_attack_technique=AttackTechnique.DIRECT_INJECTION,
        framework_mappings={'OWASP-LLM': ['LLM01']},
    ),
    Vulnerability.SENSITIVE_INFO_DISCLOSURE: VulnerabilityDef(
        id=Vulnerability.SENSITIVE_INFO_DISCLOSURE,
        name='Sensitive Information Disclosure',
        domain=VulnerabilityDomain.MODEL,
        default_attack_technique=AttackTechnique.SOCIAL_ENGINEERING,
        framework_mappings={'OWASP-LLM': ['LLM02']},
    ),
    Vulnerability.DATA_POISONING: VulnerabilityDef(
        id=Vulnerability.DATA_POISONING,
        name='Data and Model Poisoning',
        domain=VulnerabilityDomain.DATA,
        default_attack_technique=AttackTechnique.CONTEXT_POISONING,
        framework_mappings={'OWASP-LLM': ['LLM04']},
    ),
    Vulnerability.IMPROPER_OUTPUT: VulnerabilityDef(
        id=Vulnerability.IMPROPER_OUTPUT,
        name='Improper Output Handling',
        domain=VulnerabilityDomain.MODEL,
        default_attack_technique=AttackTechnique.DIRECT_INJECTION,
        framework_mappings={'OWASP-LLM': ['LLM05']},
    ),
    Vulnerability.EXCESSIVE_AGENCY: VulnerabilityDef(
        id=Vulnerability.EXCESSIVE_AGENCY,
        name='Excessive Agency',
        domain=VulnerabilityDomain.AGENT,
        default_attack_technique=AttackTechnique.TOOL_ABUSE,
        framework_mappings={'OWASP-LLM': ['LLM06']},
    ),
    Vulnerability.SYSTEM_PROMPT_LEAKAGE: VulnerabilityDef(
        id=Vulnerability.SYSTEM_PROMPT_LEAKAGE,
        name='System Prompt Leakage',
        domain=VulnerabilityDomain.MODEL,
        default_attack_technique=AttackTechnique.SOCIAL_ENGINEERING,
        framework_mappings={'OWASP-LLM': ['LLM07']},
    ),
    Vulnerability.VECTOR_EMBEDDING_WEAKNESS: VulnerabilityDef(
        id=Vulnerability.VECTOR_EMBEDDING_WEAKNESS,
        name='Vector and Embedding Weaknesses',
        domain=VulnerabilityDomain.DATA,
        default_attack_technique=AttackTechnique.CONTEXT_POISONING,
        framework_mappings={'OWASP-LLM': ['LLM08']},
    ),
    Vulnerability.MISINFORMATION: VulnerabilityDef(
        id=Vulnerability.MISINFORMATION,
        name='Misinformation',
        domain=VulnerabilityDomain.MODEL,
        default_attack_technique=AttackTechnique.SOCIAL_ENGINEERING,
        framework_mappings={'OWASP-LLM': ['LLM09']},
    ),
}

# ---------------------------------------------------------------------------
# Inverted indexes (built once at import time)
# ---------------------------------------------------------------------------

# Maps category code -> Vulnerability, e.g. 'ASI01' -> Vulnerability.GOAL_HIJACKING
CATEGORY_TO_VULNERABILITY: Mapping[str, Vulnerability]

# Maps framework -> {category -> [vulnerabilities]}, e.g.
# {'OWASP-ASI': {'ASI01': [Vulnerability.GOAL_HIJACKING]}}
FRAMEWORK_TO_VULNERABILITIES: Mapping[str, Mapping[str, list[Vulnerability]]]

_category_to_vulnerability: dict[str, Vulnerability] = {}
_framework_to_vulnerabilities: dict[str, dict[str, list[Vulnerability]]] = {}

for _vuln, _vdef in VULNERABILITY_DEFS.items():
    for _fw, _codes in _vdef.framework_mappings.items():
        if _fw not in _framework_to_vulnerabilities:
            _framework_to_vulnerabilities[_fw] = {}
        for _code in _codes:
            _category_to_vulnerability[_code] = _vuln
            # Also support OWASP- prefixed codes
            _category_to_vulnerability[f'OWASP-{_code}'] = _vuln
            _framework_to_vulnerabilities[_fw].setdefault(_code, []).append(_vuln)

_missing = set(Vulnerability) - set(VULNERABILITY_DEFS)
if _missing:
    raise RuntimeError(f'Missing VulnerabilityDef entries for: {_missing}')

VULNERABILITY_DEFS = types.MappingProxyType(VULNERABILITY_DEFS)  # pyright: ignore[reportConstantRedefinition]
CATEGORY_TO_VULNERABILITY = types.MappingProxyType(_category_to_vulnerability)  # pyright: ignore[reportConstantRedefinition]
FRAMEWORK_TO_VULNERABILITIES = types.MappingProxyType(  # pyright: ignore[reportConstantRedefinition]
    {k: types.MappingProxyType(v) for k, v in _framework_to_vulnerabilities.items()}
)


# ---------------------------------------------------------------------------
# Canonical evaluator metadata (consolidated from evaluators.py)
# ---------------------------------------------------------------------------

# Maps vulnerability -> (evaluator_id, evaluator_name) for ORQ platform evaluators.
# This is the single source of truth; evaluators.py derives category-keyed lookups from here.
CANONICAL_EVALUATOR_METADATA: dict[Vulnerability, dict[str, str]] = {
    Vulnerability.GOAL_HIJACKING: {
        'evaluator_id': '01KFTGH3A3ZEQBY4JH0RHTAQYN',
        'evaluator_name': 'owasp_asi01_goal_hijacking',
    },
    Vulnerability.CODE_EXECUTION: {
        'evaluator_id': '01KFTGH3TW3SHBNT0FF8ZC53KA',
        'evaluator_name': 'owasp_asi05_code_execution',
    },
    Vulnerability.MEMORY_POISONING: {
        'evaluator_id': '01KFTGH4AHJD0DGMVXEV0QPM21',
        'evaluator_name': 'owasp_asi06_memory_poisoning',
    },
    Vulnerability.PROMPT_INJECTION: {
        'evaluator_id': '01KFX0DEY1S57B1Q19AK10RAYF',
        'evaluator_name': 'owasp_llm01_prompt_injection',
    },
    Vulnerability.SENSITIVE_INFO_DISCLOSURE: {
        'evaluator_id': '01KFX0DEY1RWTR466RMS2FT3P7',
        'evaluator_name': 'owasp_llm02_sensitive_info_disclosure',
    },
    Vulnerability.SYSTEM_PROMPT_LEAKAGE: {
        'evaluator_id': '01KFX0DEY1DBEW2ZAQ0M3DN12T',
        'evaluator_name': 'owasp_llm07_system_prompt_leakage',
    },
}


# ---------------------------------------------------------------------------
# Resolvers
# ---------------------------------------------------------------------------


def resolve_category(category: str) -> Vulnerability:
    """Resolve an OWASP category code to a Vulnerability enum.

    Args:
        category: Category code like 'ASI01', 'LLM01', or 'OWASP-ASI01'.

    Returns:
        The corresponding Vulnerability enum value.

    Raises:
        KeyError: If the category code is not recognized.
    """
    normalized = category.removeprefix('OWASP-')
    vuln = CATEGORY_TO_VULNERABILITY.get(normalized)
    if vuln is None:
        msg = f'Unknown category code: {category!r}'
        raise KeyError(msg)
    return vuln


def resolve_category_safe(category: str) -> Vulnerability | None:
    """Like resolve_category but returns None instead of raising."""
    try:
        return resolve_category(category)
    except KeyError:
        return None


def resolve_vulnerabilities(inputs: list[str]) -> list[Vulnerability]:
    """Resolve a mix of category codes and vulnerability IDs to a list of Vulnerabilities.

    Each input is tried first as a Vulnerability enum value, then as a category code.

    Args:
        inputs: List of vulnerability IDs or category codes.

    Returns:
        Deduplicated list of Vulnerability enum values preserving order.
    """
    seen: set[Vulnerability] = set()
    result: list[Vulnerability] = []
    for item in inputs:
        # Try as Vulnerability enum value first
        try:
            vuln = Vulnerability(item)
        except ValueError:
            # Try as category code
            try:
                vuln = resolve_category(item)
            except KeyError:
                raise ValueError(
                    f'Unrecognized vulnerability or category code: {item!r}. '
                    f'Valid vulnerability IDs: {[v.value for v in Vulnerability]}. '
                    f'Valid category codes: {sorted(CATEGORY_TO_VULNERABILITY.keys())}.'
                ) from None
        if vuln not in seen:
            seen.add(vuln)
            result.append(vuln)
    return result


def get_vulnerability_name(vuln: Vulnerability) -> str:
    """Get the human-readable name for a vulnerability."""
    vdef = VULNERABILITY_DEFS.get(vuln)
    if vdef is not None:
        return vdef.name
    return vuln.value


def get_vulnerability_def(vuln: Vulnerability) -> VulnerabilityDef | None:
    """Get the full definition for a vulnerability."""
    return VULNERABILITY_DEFS.get(vuln)


def get_framework_categories(vuln: Vulnerability) -> dict[str, list[str]]:
    """Get all framework category codes for a vulnerability."""
    vdef = VULNERABILITY_DEFS.get(vuln)
    if vdef is not None:
        return vdef.framework_mappings
    return {}


def get_primary_category(vuln: Vulnerability) -> str:
    """Get the primary OWASP category code for a vulnerability.

    Returns the first category code from the first framework mapping,
    or the vulnerability value as fallback.
    """
    vdef = VULNERABILITY_DEFS.get(vuln)
    if vdef is not None:
        for codes in vdef.framework_mappings.values():
            if codes:
                return codes[0]
    return vuln.value


def list_available_vulnerabilities() -> list[Vulnerability]:
    """List all vulnerabilities that have registered definitions."""
    return list(VULNERABILITY_DEFS.keys())
