"""SAFE-MCP Threat Taxonomy — maps findings to standardized threat categories.

Provides a taxonomy of threat categories based on the SAFE-MCP threat model,
mapping each mcpaudit rule_id to one or more threat categories with MITRE
ATT&CK references and risk levels.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .models import Finding


@dataclass(frozen=True)
class ThreatCategory:
    """A standardized threat category in the SAFE-MCP taxonomy."""

    id: str
    name: str
    description: str
    risk_level: int  # 1 (low) to 5 (critical)
    mitre_attack: str = ""  # MITRE ATT&CK technique ID

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        result: dict[str, Any] = {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "risk_level": self.risk_level,
        }
        if self.mitre_attack:
            result["mitre_attack"] = self.mitre_attack
        return result


# ---------------------------------------------------------------------------
# Threat category definitions
# ---------------------------------------------------------------------------

CONFUSED_DEPUTY = ThreatCategory(
    id="confused_deputy",
    name="Confused Deputy",
    description=(
        "An MCP tool is tricked into performing privileged actions on behalf "
        "of an attacker by exploiting the trust boundary between the LLM and "
        "the tool server."
    ),
    risk_level=5,
    mitre_attack="T1210",
)

PROMPT_INJECTION = ThreatCategory(
    id="prompt_injection",
    name="Prompt Injection",
    description=(
        "Malicious input manipulates the LLM's behavior through crafted "
        "tool descriptions, arguments, or return values that alter the "
        "model's instructions."
    ),
    risk_level=5,
    mitre_attack="T1059",
)

DATA_EXFILTRATION = ThreatCategory(
    id="data_exfiltration",
    name="Data Exfiltration",
    description=(
        "Sensitive data is extracted from the server environment through "
        "file access, network requests, or side channels exposed by MCP tools."
    ),
    risk_level=4,
    mitre_attack="T1041",
)

PRIVILEGE_ESCALATION = ThreatCategory(
    id="privilege_escalation",
    name="Privilege Escalation",
    description=(
        "An attacker gains elevated permissions by exploiting unsafe "
        "deserialization, command injection, or excessive tool permissions."
    ),
    risk_level=5,
    mitre_attack="T1068",
)

DENIAL_OF_SERVICE = ThreatCategory(
    id="denial_of_service",
    name="Denial of Service",
    description=(
        "The MCP server or its resources become unavailable due to unbounded "
        "resource consumption, ReDoS patterns, or missing rate limiting."
    ),
    risk_level=3,
    mitre_attack="T1499",
)

UNAUTHORIZED_ACCESS = ThreatCategory(
    id="unauthorized_access",
    name="Unauthorized Access",
    description=(
        "Resources or operations are accessible without proper "
        "authentication or authorization checks."
    ),
    risk_level=4,
    mitre_attack="T1078",
)

INFORMATION_DISCLOSURE = ThreatCategory(
    id="information_disclosure",
    name="Information Disclosure",
    description=(
        "Sensitive information such as secrets, stack traces, environment "
        "variables, or internal paths is exposed to unauthorized parties."
    ),
    risk_level=3,
    mitre_attack="T1005",
)

SUPPLY_CHAIN = ThreatCategory(
    id="supply_chain",
    name="Supply Chain",
    description=(
        "The integrity of MCP server dependencies or tool definitions is "
        "compromised, enabling code injection through the supply chain."
    ),
    risk_level=4,
    mitre_attack="T1195",
)

# All categories indexed by id
ALL_CATEGORIES: dict[str, ThreatCategory] = {
    c.id: c for c in [
        CONFUSED_DEPUTY,
        PROMPT_INJECTION,
        DATA_EXFILTRATION,
        PRIVILEGE_ESCALATION,
        DENIAL_OF_SERVICE,
        UNAUTHORIZED_ACCESS,
        INFORMATION_DISCLOSURE,
        SUPPLY_CHAIN,
    ]
}

# ---------------------------------------------------------------------------
# Rule-to-category mapping
# ---------------------------------------------------------------------------

RULE_TAXONOMY: dict[str, list[str]] = {
    "PATH-001": ["data_exfiltration", "unauthorized_access"],
    "CMD-001": ["privilege_escalation", "confused_deputy"],
    "CMD-002": ["privilege_escalation", "confused_deputy"],
    "SEC-001": ["information_disclosure", "unauthorized_access"],
    "PATH-002": ["information_disclosure"],
    "VAL-001": ["prompt_injection", "confused_deputy"],
    "SQL-001": ["data_exfiltration", "unauthorized_access"],
    "DESER-001": ["privilege_escalation"],
    "INFO-001": ["information_disclosure"],
    "AUTH-001": ["unauthorized_access"],
    "CORS-001": ["unauthorized_access"],
    "TEMP-001": ["information_disclosure"],
    "PERM-001": ["privilege_escalation"],
    "REDOS-001": ["denial_of_service"],
    "RATE-001": ["denial_of_service"],
    "SSRF-001": ["data_exfiltration", "confused_deputy"],
    "LOG-001": ["information_disclosure"],
    "FILE-001": ["data_exfiltration"],
    "ERR-001": ["information_disclosure", "denial_of_service"],
    "RES-001": ["denial_of_service"],
    "MCP-001": ["prompt_injection"],
    "MCP-002": ["information_disclosure", "privilege_escalation"],
    "MCP-003": ["confused_deputy", "privilege_escalation"],
    "MCP-004": ["prompt_injection"],
}


def get_categories_for_rule(rule_id: str) -> list[ThreatCategory]:
    """Return threat categories mapped to a rule_id.

    Args:
        rule_id: The rule identifier (e.g. "CMD-001").

    Returns:
        List of ThreatCategory objects. Empty if the rule is unmapped.
    """
    category_ids = RULE_TAXONOMY.get(rule_id, [])
    return [ALL_CATEGORIES[cid] for cid in category_ids if cid in ALL_CATEGORIES]


def group_findings_by_category(
    findings: list[Finding],
) -> dict[str, list[Finding]]:
    """Group findings by their threat category.

    A finding may appear in multiple categories if its rule maps to more
    than one. Categories with no findings are omitted.

    Args:
        findings: List of Finding objects to classify.

    Returns:
        Dictionary mapping category_id to list of findings.
    """
    grouped: dict[str, list[Finding]] = {}
    for finding in findings:
        category_ids = RULE_TAXONOMY.get(finding.rule_id, [])
        for cat_id in category_ids:
            if cat_id not in grouped:
                grouped[cat_id] = []
            grouped[cat_id].append(finding)
    return grouped


def format_taxonomy_text(findings: list[Finding]) -> str:
    """Format findings grouped by threat category as text output.

    Args:
        findings: List of Finding objects.

    Returns:
        Human-readable text with findings organized by threat category.
    """
    grouped = group_findings_by_category(findings)
    if not grouped:
        return "No findings to classify.\n"

    lines: list[str] = []
    lines.append("")
    lines.append("=" * 60)
    lines.append("  SAFE-MCP Threat Taxonomy Report")
    lines.append("=" * 60)

    # Sort by risk level descending
    sorted_cats = sorted(
        grouped.items(),
        key=lambda item: ALL_CATEGORIES[item[0]].risk_level,
        reverse=True,
    )

    for cat_id, cat_findings in sorted_cats:
        cat = ALL_CATEGORIES[cat_id]
        lines.append("")
        lines.append(f"  [{cat.risk_level}/5] {cat.name}")
        lines.append(f"  {cat.description}")
        if cat.mitre_attack:
            lines.append(f"  MITRE ATT&CK: {cat.mitre_attack}")
        lines.append(f"  Findings: {len(cat_findings)}")
        lines.append("")
        for f in cat_findings:
            lines.append(f"    [{f.rule_id}] {f.title}")
            lines.append(f"      {f.file}:{f.line} — {f.description}")
        lines.append("")

    lines.append("=" * 60)
    total = len(findings)
    categories_hit = len(grouped)
    lines.append(
        f"  {total} findings across {categories_hit} threat categories"
    )
    lines.append("=" * 60)
    return "\n".join(lines)
