"""Threat Model Generator for MCP servers.

Generates automatic STRIDE-based threat models from scan results.
Identifies assets, threat agents, attack vectors, and mitigations.
Produces a threat matrix (probability x impact) with HTML output.

Zero external dependencies. Python 3.11+ stdlib only.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from enum import Enum
from typing import Any


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------

class StrideCategory(Enum):
    """STRIDE threat model categories."""

    SPOOFING = "Spoofing"
    TAMPERING = "Tampering"
    REPUDIATION = "Repudiation"
    INFO_DISCLOSURE = "Information Disclosure"
    DENIAL_OF_SERVICE = "Denial of Service"
    ELEVATION = "Elevation of Privilege"

    @classmethod
    def from_string(cls, value: str) -> StrideCategory:
        """Parse a STRIDE category string (case-insensitive)."""
        mapping = {
            "spoofing": cls.SPOOFING,
            "tampering": cls.TAMPERING,
            "repudiation": cls.REPUDIATION,
            "info_disclosure": cls.INFO_DISCLOSURE,
            "information_disclosure": cls.INFO_DISCLOSURE,
            "information disclosure": cls.INFO_DISCLOSURE,
            "denial_of_service": cls.DENIAL_OF_SERVICE,
            "denial of service": cls.DENIAL_OF_SERVICE,
            "dos": cls.DENIAL_OF_SERVICE,
            "elevation": cls.ELEVATION,
            "elevation_of_privilege": cls.ELEVATION,
            "elevation of privilege": cls.ELEVATION,
        }
        key = value.lower().strip()
        if key in mapping:
            return mapping[key]
        valid = ", ".join(sorted(mapping.keys()))
        raise ValueError(f"Unknown STRIDE category '{value}'. Valid: {valid}")


class Likelihood(Enum):
    """Threat likelihood levels."""

    VERY_HIGH = "Very High"
    HIGH = "High"
    MEDIUM = "Medium"
    LOW = "Low"
    VERY_LOW = "Very Low"

    @property
    def score(self) -> int:
        return {
            Likelihood.VERY_HIGH: 5,
            Likelihood.HIGH: 4,
            Likelihood.MEDIUM: 3,
            Likelihood.LOW: 2,
            Likelihood.VERY_LOW: 1,
        }[self]


class Impact(Enum):
    """Threat impact levels."""

    CRITICAL = "Critical"
    HIGH = "High"
    MEDIUM = "Medium"
    LOW = "Low"
    NEGLIGIBLE = "Negligible"

    @property
    def score(self) -> int:
        return {
            Impact.CRITICAL: 5,
            Impact.HIGH: 4,
            Impact.MEDIUM: 3,
            Impact.LOW: 2,
            Impact.NEGLIGIBLE: 1,
        }[self]


class AssetType(Enum):
    """Types of assets in MCP servers."""

    TOOL = "Tool"
    DATA = "Data"
    CONFIG = "Configuration"
    CREDENTIAL = "Credential"
    RESOURCE = "Resource"
    PROMPT = "Prompt"
    ENDPOINT = "Endpoint"


# ---------------------------------------------------------------------------
# Data Models
# ---------------------------------------------------------------------------

@dataclass
class Asset:
    """An asset identified in the MCP server."""

    name: str
    asset_type: AssetType
    description: str
    sensitivity: str = "medium"  # low, medium, high, critical

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "type": self.asset_type.value,
            "description": self.description,
            "sensitivity": self.sensitivity,
        }


@dataclass
class ThreatAgent:
    """A potential threat agent."""

    name: str
    description: str
    capability: str = "medium"  # low, medium, high
    motivation: str = "medium"  # low, medium, high

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "capability": self.capability,
            "motivation": self.motivation,
        }


@dataclass
class AttackVector:
    """An attack vector associated with a threat."""

    name: str
    description: str
    technique: str = ""
    prerequisites: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "technique": self.technique,
            "prerequisites": self.prerequisites,
        }


@dataclass
class Mitigation:
    """A recommended mitigation for a threat."""

    name: str
    description: str
    priority: str = "medium"  # low, medium, high, critical
    status: str = "recommended"  # recommended, implemented, partial
    effort: str = "medium"  # low, medium, high

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "priority": self.priority,
            "status": self.status,
            "effort": self.effort,
        }


@dataclass
class Threat:
    """A single threat in the model."""

    threat_id: str
    name: str
    description: str
    stride_category: StrideCategory
    likelihood: Likelihood
    impact: Impact
    affected_assets: list[str] = field(default_factory=list)
    attack_vectors: list[AttackVector] = field(default_factory=list)
    mitigations: list[Mitigation] = field(default_factory=list)
    related_findings: list[str] = field(default_factory=list)

    @property
    def risk_score(self) -> int:
        """Risk = likelihood x impact (1-25 scale)."""
        return self.likelihood.score * self.impact.score

    @property
    def risk_level(self) -> str:
        """Qualitative risk level from score."""
        s = self.risk_score
        if s >= 20:
            return "critical"
        if s >= 12:
            return "high"
        if s >= 6:
            return "medium"
        return "low"

    def to_dict(self) -> dict[str, Any]:
        return {
            "threat_id": self.threat_id,
            "name": self.name,
            "description": self.description,
            "stride_category": self.stride_category.value,
            "likelihood": self.likelihood.value,
            "impact": self.impact.value,
            "risk_score": self.risk_score,
            "risk_level": self.risk_level,
            "affected_assets": self.affected_assets,
            "attack_vectors": [v.to_dict() for v in self.attack_vectors],
            "mitigations": [m.to_dict() for m in self.mitigations],
            "related_findings": self.related_findings,
        }


@dataclass
class ThreatModel:
    """Complete threat model for an MCP server."""

    name: str
    description: str
    assets: list[Asset] = field(default_factory=list)
    threat_agents: list[ThreatAgent] = field(default_factory=list)
    threats: list[Threat] = field(default_factory=list)
    scan_grade: str = "?"
    total_findings: int = 0

    @property
    def threat_count(self) -> int:
        return len(self.threats)

    @property
    def stride_summary(self) -> dict[str, int]:
        """Count threats per STRIDE category."""
        counts: dict[str, int] = {c.value: 0 for c in StrideCategory}
        for t in self.threats:
            counts[t.stride_category.value] += 1
        return counts

    @property
    def risk_matrix(self) -> list[dict[str, Any]]:
        """Build risk matrix: list of (threat_id, likelihood, impact, score)."""
        return [
            {
                "threat_id": t.threat_id,
                "name": t.name,
                "likelihood": t.likelihood.value,
                "impact": t.impact.value,
                "risk_score": t.risk_score,
                "risk_level": t.risk_level,
            }
            for t in sorted(self.threats, key=lambda x: x.risk_score, reverse=True)
        ]

    @property
    def critical_threats(self) -> list[Threat]:
        """Threats with critical risk level."""
        return [t for t in self.threats if t.risk_level == "critical"]

    @property
    def high_threats(self) -> list[Threat]:
        """Threats with high risk level."""
        return [t for t in self.threats if t.risk_level == "high"]

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "scan_grade": self.scan_grade,
            "total_findings": self.total_findings,
            "threat_count": self.threat_count,
            "stride_summary": self.stride_summary,
            "risk_matrix": self.risk_matrix,
            "assets": [a.to_dict() for a in self.assets],
            "threat_agents": [a.to_dict() for a in self.threat_agents],
            "threats": [t.to_dict() for t in self.threats],
        }


# ---------------------------------------------------------------------------
# Rule-to-STRIDE mapping
# ---------------------------------------------------------------------------

# Maps rule ID prefixes to STRIDE categories with likelihood/impact defaults
_RULE_STRIDE_MAP: dict[str, tuple[StrideCategory, Likelihood, Impact]] = {
    "CMD": (StrideCategory.ELEVATION, Likelihood.HIGH, Impact.CRITICAL),
    "AUTH": (StrideCategory.SPOOFING, Likelihood.HIGH, Impact.HIGH),
    "RATE": (StrideCategory.DENIAL_OF_SERVICE, Likelihood.MEDIUM, Impact.MEDIUM),
    "PRIV": (StrideCategory.ELEVATION, Likelihood.MEDIUM, Impact.HIGH),
    "LOG": (StrideCategory.REPUDIATION, Likelihood.MEDIUM, Impact.MEDIUM),
    "INJ": (StrideCategory.TAMPERING, Likelihood.HIGH, Impact.HIGH),
    "PROMPT": (StrideCategory.TAMPERING, Likelihood.HIGH, Impact.HIGH),
    "CRYPTO": (StrideCategory.INFO_DISCLOSURE, Likelihood.MEDIUM, Impact.HIGH),
    "SEC": (StrideCategory.INFO_DISCLOSURE, Likelihood.MEDIUM, Impact.MEDIUM),
    "LEAK": (StrideCategory.INFO_DISCLOSURE, Likelihood.HIGH, Impact.HIGH),
    "SSRF": (StrideCategory.TAMPERING, Likelihood.HIGH, Impact.HIGH),
    "XSS": (StrideCategory.TAMPERING, Likelihood.MEDIUM, Impact.MEDIUM),
    "PATH": (StrideCategory.INFO_DISCLOSURE, Likelihood.MEDIUM, Impact.HIGH),
    "DESER": (StrideCategory.ELEVATION, Likelihood.MEDIUM, Impact.CRITICAL),
    "CORS": (StrideCategory.SPOOFING, Likelihood.MEDIUM, Impact.MEDIUM),
    "ERR": (StrideCategory.INFO_DISCLOSURE, Likelihood.LOW, Impact.LOW),
    "VALID": (StrideCategory.TAMPERING, Likelihood.MEDIUM, Impact.MEDIUM),
    "TOOL": (StrideCategory.ELEVATION, Likelihood.HIGH, Impact.HIGH),
    "MCP": (StrideCategory.TAMPERING, Likelihood.MEDIUM, Impact.HIGH),
    "RESOURCE": (StrideCategory.INFO_DISCLOSURE, Likelihood.MEDIUM, Impact.MEDIUM),
    "FLOW": (StrideCategory.INFO_DISCLOSURE, Likelihood.HIGH, Impact.HIGH),
    "SCOPE": (StrideCategory.ELEVATION, Likelihood.MEDIUM, Impact.HIGH),
    "AGENT": (StrideCategory.ELEVATION, Likelihood.HIGH, Impact.CRITICAL),
    "CONFUSED": (StrideCategory.SPOOFING, Likelihood.HIGH, Impact.CRITICAL),
    "DEPUTY": (StrideCategory.SPOOFING, Likelihood.HIGH, Impact.CRITICAL),
}

# Severity string to likelihood/impact boost
_SEVERITY_BOOST: dict[str, tuple[Likelihood, Impact]] = {
    "CRITICAL": (Likelihood.VERY_HIGH, Impact.CRITICAL),
    "HIGH": (Likelihood.HIGH, Impact.HIGH),
    "MEDIUM": (Likelihood.MEDIUM, Impact.MEDIUM),
    "LOW": (Likelihood.LOW, Impact.LOW),
    "INFO": (Likelihood.VERY_LOW, Impact.NEGLIGIBLE),
}


def _default_threat_agents() -> list[ThreatAgent]:
    """Return standard MCP threat agents."""
    return [
        ThreatAgent(
            name="Malicious User",
            description="End user attempting to abuse MCP tools",
            capability="medium",
            motivation="high",
        ),
        ThreatAgent(
            name="Compromised AI Agent",
            description="AI agent executing prompt injection attacks",
            capability="high",
            motivation="high",
        ),
        ThreatAgent(
            name="Supply Chain Attacker",
            description="Attacker injecting malicious MCP packages",
            capability="high",
            motivation="medium",
        ),
        ThreatAgent(
            name="Insider Threat",
            description="Developer with access to MCP server code",
            capability="high",
            motivation="low",
        ),
        ThreatAgent(
            name="Network Attacker",
            description="Attacker intercepting MCP communications",
            capability="medium",
            motivation="medium",
        ),
    ]


def _classify_rule(rule_id: str) -> str:
    """Get the prefix category of a rule ID."""
    for prefix in _RULE_STRIDE_MAP:
        if rule_id.startswith(prefix):
            return prefix
    return ""


def _extract_assets_from_scan(scan_data: dict[str, Any]) -> list[Asset]:
    """Extract assets from scan results."""
    assets: list[Asset] = []
    seen: set[str] = set()

    results = scan_data.get("results", [])
    for r in results:
        fname = r.get("file", "")
        if fname and fname not in seen:
            seen.add(fname)
            # Determine asset type from file content
            findings = r.get("findings", [])
            has_tool = any("tool" in f.get("rule_id", "").lower() for f in findings)
            has_auth = any("auth" in f.get("rule_id", "").lower() for f in findings)
            has_config = any(
                kw in fname.lower()
                for kw in ("config", "settings", ".env", ".yml", ".yaml", ".json")
            )

            if has_config:
                assets.append(Asset(
                    name=fname,
                    asset_type=AssetType.CONFIG,
                    description=f"Configuration file: {fname}",
                    sensitivity="high",
                ))
            elif has_tool:
                assets.append(Asset(
                    name=fname,
                    asset_type=AssetType.TOOL,
                    description=f"MCP tool definition: {fname}",
                    sensitivity="high",
                ))
            elif has_auth:
                assets.append(Asset(
                    name=fname,
                    asset_type=AssetType.CREDENTIAL,
                    description=f"Authentication-related file: {fname}",
                    sensitivity="critical",
                ))
            else:
                assets.append(Asset(
                    name=fname,
                    asset_type=AssetType.RESOURCE,
                    description=f"Source file: {fname}",
                    sensitivity="medium",
                ))

    # Always add standard MCP assets
    if not any(a.name == "MCP Transport Layer" for a in assets):
        assets.append(Asset(
            name="MCP Transport Layer",
            asset_type=AssetType.ENDPOINT,
            description="Communication channel between AI agent and MCP server",
            sensitivity="high",
        ))

    return assets


def _build_mitigations(
    stride: StrideCategory,
    rule_id: str,
    recommendation: str,
) -> list[Mitigation]:
    """Build mitigations for a threat based on STRIDE category."""
    mitigations: list[Mitigation] = []

    if recommendation:
        mitigations.append(Mitigation(
            name=f"Fix {rule_id}",
            description=recommendation,
            priority="high",
            status="recommended",
            effort="medium",
        ))

    # Generic STRIDE mitigations
    stride_mitigations: dict[StrideCategory, list[Mitigation]] = {
        StrideCategory.SPOOFING: [
            Mitigation(
                name="Implement authentication",
                description="Add mutual authentication for MCP transport",
                priority="high",
                effort="medium",
            ),
            Mitigation(
                name="Token validation",
                description="Validate all authentication tokens server-side",
                priority="high",
                effort="low",
            ),
        ],
        StrideCategory.TAMPERING: [
            Mitigation(
                name="Input validation",
                description="Validate and sanitize all tool inputs",
                priority="high",
                effort="medium",
            ),
            Mitigation(
                name="Integrity checks",
                description="Use checksums/signatures for critical data",
                priority="medium",
                effort="medium",
            ),
        ],
        StrideCategory.REPUDIATION: [
            Mitigation(
                name="Audit logging",
                description="Log all tool invocations with timestamps",
                priority="medium",
                effort="low",
            ),
            Mitigation(
                name="Non-repudiation",
                description="Sign audit log entries",
                priority="low",
                effort="high",
            ),
        ],
        StrideCategory.INFO_DISCLOSURE: [
            Mitigation(
                name="Data classification",
                description="Classify and protect sensitive data in responses",
                priority="high",
                effort="medium",
            ),
            Mitigation(
                name="Encryption",
                description="Encrypt sensitive data at rest and in transit",
                priority="high",
                effort="medium",
            ),
        ],
        StrideCategory.DENIAL_OF_SERVICE: [
            Mitigation(
                name="Rate limiting",
                description="Implement rate limits on tool invocations",
                priority="high",
                effort="low",
            ),
            Mitigation(
                name="Resource quotas",
                description="Set resource quotas per session/user",
                priority="medium",
                effort="medium",
            ),
        ],
        StrideCategory.ELEVATION: [
            Mitigation(
                name="Least privilege",
                description="Apply least-privilege principle to tool permissions",
                priority="critical",
                effort="medium",
            ),
            Mitigation(
                name="Sandbox execution",
                description="Run tools in sandboxed environments",
                priority="high",
                effort="high",
            ),
        ],
    }

    for m in stride_mitigations.get(stride, []):
        if not any(existing.name == m.name for existing in mitigations):
            mitigations.append(m)

    return mitigations


# ---------------------------------------------------------------------------
# Core Functions
# ---------------------------------------------------------------------------

def generate_threat_model(scan_result: dict[str, Any]) -> ThreatModel:
    """Generate a STRIDE threat model from scan results.

    Args:
        scan_result: ScanSummary.to_dict() output or compatible dict.

    Returns:
        Complete ThreatModel with threats, assets, and mitigations.
    """
    summary = scan_result.get("summary", {})
    results = scan_result.get("results", [])

    model = ThreatModel(
        name="MCP Server Threat Model",
        description="Auto-generated STRIDE threat model from mcpaudit scan",
        scan_grade=summary.get("overall_grade", "?"),
        total_findings=summary.get("total_findings", 0),
    )

    # Extract assets
    model.assets = _extract_assets_from_scan(scan_result)

    # Add threat agents
    model.threat_agents = _default_threat_agents()

    # Build threats from findings
    seen_rules: dict[str, list[dict[str, Any]]] = {}
    for r in results:
        for f in r.get("findings", []):
            rule_id = f.get("rule_id", "")
            if rule_id not in seen_rules:
                seen_rules[rule_id] = []
            seen_rules[rule_id].append(f)

    threat_idx = 0
    for rule_id, findings in seen_rules.items():
        prefix = _classify_rule(rule_id)
        if prefix and prefix in _RULE_STRIDE_MAP:
            stride, default_lh, default_imp = _RULE_STRIDE_MAP[prefix]
        else:
            stride = StrideCategory.TAMPERING
            default_lh = Likelihood.MEDIUM
            default_imp = Impact.MEDIUM

        # Boost from severity
        severity = findings[0].get("severity", "MEDIUM")
        if severity in _SEVERITY_BOOST:
            boosted_lh, boosted_imp = _SEVERITY_BOOST[severity]
            # Use the higher of default and severity-boosted values
            lh = max(default_lh, boosted_lh, key=lambda x: x.score)
            imp = max(default_imp, boosted_imp, key=lambda x: x.score)
        else:
            lh, imp = default_lh, default_imp

        first = findings[0]
        recommendation = first.get("recommendation", "")
        affected_files = list({f.get("file", "") for f in findings if f.get("file")})

        attack_vectors = [AttackVector(
            name=f"Exploit {rule_id}",
            description=first.get("description", ""),
            technique=first.get("title", ""),
            prerequisites=["Access to MCP server"],
        )]

        mitigations = _build_mitigations(stride, rule_id, recommendation)

        threat_idx += 1
        model.threats.append(Threat(
            threat_id=f"T-{threat_idx:03d}",
            name=first.get("title", rule_id),
            description=first.get("description", ""),
            stride_category=stride,
            likelihood=lh,
            impact=imp,
            affected_assets=affected_files,
            attack_vectors=attack_vectors,
            mitigations=mitigations,
            related_findings=[rule_id],
        ))

    return model


def format_threat_model_text(model: ThreatModel) -> str:
    """Format threat model as human-readable text.

    Args:
        model: ThreatModel to format.

    Returns:
        Formatted text string.
    """
    lines: list[str] = []
    lines.append(f"\n{'=' * 70}")
    lines.append(f"  {model.name}")
    lines.append(f"{'=' * 70}")
    lines.append(f"  Grade: {model.scan_grade}")
    lines.append(f"  Total Findings: {model.total_findings}")
    lines.append(f"  Threats Identified: {model.threat_count}")
    lines.append(f"  Critical Threats: {len(model.critical_threats)}")
    lines.append(f"  High Threats: {len(model.high_threats)}")
    lines.append("")

    # STRIDE summary
    lines.append("  STRIDE Summary:")
    for cat, count in model.stride_summary.items():
        lines.append(f"    {cat}: {count}")
    lines.append("")

    # Risk Matrix
    lines.append(f"  {'Threat':<40} {'Likelihood':<12} {'Impact':<10} {'Risk':<6} {'Level'}")
    lines.append(f"  {'-' * 40} {'-' * 12} {'-' * 10} {'-' * 6} {'-' * 10}")
    for entry in model.risk_matrix:
        lines.append(
            f"  {entry['name'][:40]:<40} "
            f"{entry['likelihood']:<12} "
            f"{entry['impact']:<10} "
            f"{entry['risk_score']:<6} "
            f"{entry['risk_level']}"
        )
    lines.append("")

    # Detailed threats
    for t in model.threats:
        lines.append(f"  [{t.threat_id}] {t.name}")
        lines.append(f"    Category: {t.stride_category.value}")
        lines.append(f"    Risk: {t.risk_score} ({t.risk_level})")
        if t.mitigations:
            lines.append(f"    Mitigations: {len(t.mitigations)}")
            for m in t.mitigations[:3]:
                lines.append(f"      - {m.name}: {m.description[:60]}")
        lines.append("")

    lines.append(f"{'=' * 70}\n")
    return "\n".join(lines)


def format_threat_model_json(model: ThreatModel) -> str:
    """Format threat model as JSON.

    Args:
        model: ThreatModel to format.

    Returns:
        JSON string.
    """
    return json.dumps(model.to_dict(), indent=2)


def format_threat_model_html(model: ThreatModel) -> str:
    """Format threat model as HTML report.

    Args:
        model: ThreatModel to format.

    Returns:
        Complete HTML document string.
    """
    # Build STRIDE summary rows
    stride_rows = ""
    stride_colors = {
        "Spoofing": "#ef4444",
        "Tampering": "#f97316",
        "Repudiation": "#eab308",
        "Information Disclosure": "#3b82f6",
        "Denial of Service": "#8b5cf6",
        "Elevation of Privilege": "#ec4899",
    }
    for cat, count in model.stride_summary.items():
        color = stride_colors.get(cat, "#6b7280")
        stride_rows += (
            f'<div class="stride-card">'
            f'<div class="stride-label" style="color:{color}">{cat}</div>'
            f'<div class="stride-value">{count}</div>'
            f'</div>'
        )

    # Build risk matrix rows
    risk_rows = ""
    risk_colors = {
        "critical": "#ef4444",
        "high": "#f97316",
        "medium": "#eab308",
        "low": "#22c55e",
    }
    for entry in model.risk_matrix:
        rc = risk_colors.get(entry["risk_level"], "#6b7280")
        risk_rows += (
            f'<tr>'
            f'<td>{entry["threat_id"]}</td>'
            f'<td>{_escape_html(entry["name"])}</td>'
            f'<td>{entry["likelihood"]}</td>'
            f'<td>{entry["impact"]}</td>'
            f'<td style="color:{rc};font-weight:bold">{entry["risk_score"]}</td>'
            f'<td style="color:{rc};font-weight:bold">{entry["risk_level"].upper()}</td>'
            f'</tr>'
        )

    # Build threat detail cards
    threat_cards = ""
    for t in model.threats:
        rc = risk_colors.get(t.risk_level, "#6b7280")
        mitigation_list = ""
        for m in t.mitigations:
            mitigation_list += (
                f'<li><strong>{_escape_html(m.name)}</strong> '
                f'[{m.priority}]: {_escape_html(m.description)}</li>'
            )

        av_list = ""
        for av in t.attack_vectors:
            av_list += f'<li>{_escape_html(av.name)}: {_escape_html(av.description[:100])}</li>'

        threat_cards += f"""
        <div class="threat-card">
            <div class="threat-header">
                <span class="threat-id">{t.threat_id}</span>
                <span class="threat-name">{_escape_html(t.name)}</span>
                <span class="risk-badge" style="background:{rc}">{t.risk_level.upper()}</span>
            </div>
            <p>{_escape_html(t.description[:200])}</p>
            <div class="threat-meta">
                <span>Category: {t.stride_category.value}</span>
                <span>Likelihood: {t.likelihood.value}</span>
                <span>Impact: {t.impact.value}</span>
                <span>Score: {t.risk_score}/25</span>
            </div>
            <div class="threat-section">
                <h4>Attack Vectors</h4>
                <ul>{av_list}</ul>
            </div>
            <div class="threat-section">
                <h4>Mitigations</h4>
                <ul>{mitigation_list}</ul>
            </div>
        </div>"""

    # Asset table
    asset_rows = ""
    for a in model.assets:
        sens_color = {
            "critical": "#ef4444",
            "high": "#f97316",
            "medium": "#eab308",
            "low": "#22c55e",
        }.get(a.sensitivity, "#6b7280")
        asset_rows += (
            f'<tr>'
            f'<td>{_escape_html(a.name)}</td>'
            f'<td>{a.asset_type.value}</td>'
            f'<td style="color:{sens_color}">{a.sensitivity.upper()}</td>'
            f'<td>{_escape_html(a.description)}</td>'
            f'</tr>'
        )

    grade_color = {
        "A+": "#22c55e", "A": "#22c55e", "B": "#eab308",
        "C": "#f97316", "D": "#ef4444", "F": "#dc2626",
    }.get(model.scan_grade, "#6b7280")

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>MCP Threat Model — {_escape_html(model.name)}</title>
<style>
* {{ margin: 0; padding: 0; box-sizing: border-box; }}
body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
       background: #0f172a; color: #e2e8f0; padding: 24px; line-height: 1.6; }}
h1 {{ font-size: 1.5rem; margin-bottom: 4px; }}
h2 {{ font-size: 1.2rem; margin: 24px 0 12px; color: #94a3b8; }}
h3 {{ font-size: 1rem; margin: 16px 0 8px; }}
h4 {{ font-size: 0.85rem; margin: 8px 0 4px; color: #94a3b8; }}
.subtitle {{ color: #94a3b8; margin-bottom: 24px; }}
.cards {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr));
          gap: 12px; margin-bottom: 24px; }}
.card {{ background: #1e293b; border-radius: 10px; padding: 16px; text-align: center; }}
.card-label {{ font-size: 0.7rem; text-transform: uppercase; color: #94a3b8; }}
.card-value {{ font-size: 1.8rem; font-weight: 700; margin-top: 4px; }}
.stride-grid {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(160px, 1fr));
                gap: 8px; margin-bottom: 24px; }}
.stride-card {{ background: #1e293b; border-radius: 8px; padding: 12px; text-align: center; }}
.stride-label {{ font-size: 0.75rem; text-transform: uppercase; font-weight: 600; }}
.stride-value {{ font-size: 1.5rem; font-weight: 700; color: #e2e8f0; }}
table {{ width: 100%; border-collapse: collapse; background: #1e293b;
         border-radius: 10px; overflow: hidden; margin-bottom: 24px; }}
th {{ text-align: left; padding: 10px 14px; background: #334155;
      font-size: 0.75rem; text-transform: uppercase; color: #94a3b8; }}
td {{ padding: 8px 14px; border-bottom: 1px solid #334155; font-size: 0.85rem; }}
tr:hover {{ background: #1e3a5f; }}
.threat-card {{ background: #1e293b; border-radius: 10px; padding: 16px;
                margin-bottom: 12px; }}
.threat-header {{ display: flex; align-items: center; gap: 12px; margin-bottom: 8px; }}
.threat-id {{ font-weight: 700; color: #94a3b8; }}
.threat-name {{ font-weight: 600; flex: 1; }}
.risk-badge {{ padding: 2px 8px; border-radius: 4px; color: white;
               font-size: 0.7rem; font-weight: 700; text-transform: uppercase; }}
.threat-meta {{ display: flex; gap: 16px; font-size: 0.8rem; color: #94a3b8;
                margin: 8px 0; flex-wrap: wrap; }}
.threat-section {{ margin-top: 8px; }}
.threat-section ul {{ margin-left: 20px; font-size: 0.85rem; }}
.threat-section li {{ margin-bottom: 4px; }}
.footer {{ text-align: center; color: #475569; margin-top: 32px; font-size: 0.8rem; }}
</style>
</head>
<body>
<h1>{_escape_html(model.name)}</h1>
<p class="subtitle">STRIDE-based threat model generated from mcpaudit scan</p>

<div class="cards">
    <div class="card">
        <div class="card-label">Scan Grade</div>
        <div class="card-value" style="color:{grade_color}">{model.scan_grade}</div>
    </div>
    <div class="card">
        <div class="card-label">Total Findings</div>
        <div class="card-value">{model.total_findings}</div>
    </div>
    <div class="card">
        <div class="card-label">Threats</div>
        <div class="card-value">{model.threat_count}</div>
    </div>
    <div class="card">
        <div class="card-label">Critical</div>
        <div class="card-value" style="color:#ef4444">{len(model.critical_threats)}</div>
    </div>
    <div class="card">
        <div class="card-label">High</div>
        <div class="card-value" style="color:#f97316">{len(model.high_threats)}</div>
    </div>
    <div class="card">
        <div class="card-label">Assets</div>
        <div class="card-value">{len(model.assets)}</div>
    </div>
</div>

<h2>STRIDE Summary</h2>
<div class="stride-grid">{stride_rows}</div>

<h2>Risk Matrix</h2>
<table>
<thead><tr>
    <th>ID</th><th>Threat</th><th>Likelihood</th><th>Impact</th><th>Score</th><th>Risk</th>
</tr></thead>
<tbody>{risk_rows}</tbody>
</table>

<h2>Assets</h2>
<table>
<thead><tr>
    <th>Name</th><th>Type</th><th>Sensitivity</th><th>Description</th>
</tr></thead>
<tbody>{asset_rows}</tbody>
</table>

<h2>Threat Details</h2>
{threat_cards}

<p class="footer">Generated by mcpaudit threat model generator</p>
</body>
</html>"""


def _escape_html(text: str) -> str:
    """Escape HTML special characters."""
    return (
        text.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )
