"""Rule Coverage Analyzer — Analyze framework coverage by mcpaudit rules.

Maps mcpaudit rules against OWASP Agentic Top 10, CWE categories, and
MITRE ATT&CK for AI to identify coverage and gaps.
Zero external dependencies — pure Python 3.11+ stdlib.
"""

from __future__ import annotations

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


@dataclass
class CoverageGap:
    """A gap in framework coverage."""

    category_id: str
    category_name: str
    description: str
    priority: str  # HIGH, MEDIUM, LOW

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


@dataclass
class CoverageMatrix:
    """Coverage analysis for a security framework."""

    framework_name: str
    total_categories: int
    covered_categories: int
    coverage_pct: float
    gaps: list[CoverageGap] = field(default_factory=list)
    covered: list[dict[str, Any]] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "framework_name": self.framework_name,
            "total_categories": self.total_categories,
            "covered_categories": self.covered_categories,
            "coverage_pct": round(self.coverage_pct, 1),
            "gaps": [g.to_dict() for g in self.gaps],
            "covered": self.covered,
        }


# ---------------------------------------------------------------------------
# OWASP Agentic Top 10 categories and rule mapping
# ---------------------------------------------------------------------------

_OWASP_AGENTIC_TOP10: list[dict[str, Any]] = [
    {
        "id": "OWASP-A01",
        "name": "Prompt Injection",
        "description": "Direct or indirect prompt injection attacks against AI agents",
        "rule_ids": ["PI-001", "TS-PI-001", "MCP-PI-001"],
    },
    {
        "id": "OWASP-A02",
        "name": "Tool Misuse & Confused Deputy",
        "description": "Agents tricked into misusing tools or acting as confused deputies",
        "rule_ids": ["CD-001", "MCP-TOOL-001", "MCP-TOOL-002", "TOOLZ-001", "TS-MTA-001"],
    },
    {
        "id": "OWASP-A03",
        "name": "Excessive Permissions",
        "description": "Agents granted overly broad permissions beyond minimum required",
        "rule_ids": ["PERM-001", "TS-PERM-001"],
    },
    {
        "id": "OWASP-A04",
        "name": "Missing Authentication & Authorization",
        "description": "Lack of authentication or authorization checks in agent endpoints",
        "rule_ids": ["AUTH-001", "TS-AUTH-001", "GO-AUTH-001", "MTA-001", "TS-MTA-001"],
    },
    {
        "id": "OWASP-A05",
        "name": "Information Disclosure",
        "description": "Unintended exposure of sensitive data through agent responses",
        "rule_ids": ["INFO-001", "TS-INFO-001", "LOG-001"],
    },
    {
        "id": "OWASP-A06",
        "name": "Command Injection",
        "description": "Injection of OS commands through unvalidated agent inputs",
        "rule_ids": ["CMD-001", "CMD-002", "TS-CMD-001", "GO-CMD-001"],
    },
    {
        "id": "OWASP-A07",
        "name": "SSRF & Network Abuse",
        "description": "Server-side request forgery and network abuse via agent actions",
        "rule_ids": ["SSRF-001", "TS-SSRF-001", "GO-SSRF-001"],
    },
    {
        "id": "OWASP-A08",
        "name": "Path Traversal & File Access",
        "description": "Path traversal attacks to access unauthorized files",
        "rule_ids": ["PATH-001", "TS-PATH-001", "GO-PATH-001", "FILE-001", "ABS-001"],
    },
    {
        "id": "OWASP-A09",
        "name": "Unsafe Deserialization",
        "description": "Deserialization of untrusted data leading to code execution",
        "rule_ids": ["DESER-001", "TS-DESER-001", "GO-DESER-001"],
    },
    {
        "id": "OWASP-A10",
        "name": "Denial of Service & Resource Exhaustion",
        "description": "Resource exhaustion attacks against agent infrastructure",
        "rule_ids": ["RATE-001", "TS-RATE-001", "RES-001", "REDOS-001"],
    },
]


# ---------------------------------------------------------------------------
# CWE categories relevant to MCP security
# ---------------------------------------------------------------------------

_CWE_CATEGORIES: list[dict[str, Any]] = [
    {
        "id": "CWE-78",
        "name": "OS Command Injection",
        "description": "Improper Neutralization of Special Elements used in an OS Command",
        "rule_ids": ["CMD-001", "CMD-002", "TS-CMD-001", "GO-CMD-001"],
    },
    {
        "id": "CWE-89",
        "name": "SQL Injection",
        "description": "Improper Neutralization of Special Elements used in an SQL Command",
        "rule_ids": ["SQL-001", "TS-SQL-001", "GO-SQL-001"],
    },
    {
        "id": "CWE-22",
        "name": "Path Traversal",
        "description": "Improper Limitation of a Pathname to a Restricted Directory",
        "rule_ids": ["PATH-001", "TS-PATH-001", "GO-PATH-001"],
    },
    {
        "id": "CWE-918",
        "name": "Server-Side Request Forgery (SSRF)",
        "description": "Server receives crafted request causing it to make unintended requests",
        "rule_ids": ["SSRF-001", "TS-SSRF-001", "GO-SSRF-001"],
    },
    {
        "id": "CWE-502",
        "name": "Deserialization of Untrusted Data",
        "description": "Application deserializes untrusted data without verification",
        "rule_ids": ["DESER-001", "TS-DESER-001", "GO-DESER-001"],
    },
    {
        "id": "CWE-862",
        "name": "Missing Authorization",
        "description": "Product does not perform authorization check before action",
        "rule_ids": ["AUTH-001", "TS-AUTH-001", "GO-AUTH-001", "MTA-001", "TS-MTA-001"],
    },
    {
        "id": "CWE-200",
        "name": "Exposure of Sensitive Information",
        "description": "Product exposes sensitive information to unauthorized actors",
        "rule_ids": ["INFO-001", "TS-INFO-001", "LOG-001"],
    },
    {
        "id": "CWE-327",
        "name": "Use of Broken Cryptographic Algorithm",
        "description": "Use of a broken or risky cryptographic algorithm",
        "rule_ids": ["CRYPTO-001"],
    },
    {
        "id": "CWE-732",
        "name": "Incorrect Permission Assignment",
        "description": "Product sets overly permissive access rights for critical resources",
        "rule_ids": ["PERM-001", "TS-PERM-001"],
    },
    {
        "id": "CWE-1333",
        "name": "Inefficient Regular Expression Complexity",
        "description": "Regular expression with excessive complexity (ReDoS)",
        "rule_ids": ["REDOS-001"],
    },
    {
        "id": "CWE-400",
        "name": "Uncontrolled Resource Consumption",
        "description": "Product does not properly control resource allocation/usage",
        "rule_ids": ["RES-001", "RATE-001", "TS-RATE-001"],
    },
    {
        "id": "CWE-377",
        "name": "Insecure Temporary File",
        "description": "Creating and using temporary files in an insecure manner",
        "rule_ids": ["TEMP-001"],
    },
    {
        "id": "CWE-942",
        "name": "Overly Permissive Cross-domain Whitelist",
        "description": "CORS misconfiguration allowing unauthorized cross-origin requests",
        "rule_ids": ["CORS-001"],
    },
    {
        "id": "CWE-20",
        "name": "Improper Input Validation",
        "description": "Product does not validate or incorrectly validates input",
        "rule_ids": ["INPUT-001", "TS-INPUT-001"],
    },
    {
        "id": "CWE-639",
        "name": "Authorization Bypass Through User-Controlled Key (IDOR)",
        "description": "System uses user-controlled key to access data without authorization checks",
        "rule_ids": ["IDOR-001", "TS-IDOR-001"],
    },
    {
        "id": "CWE-269",
        "name": "Improper Privilege Management",
        "description": "Product does not properly manage privilege escalation",
        "rule_ids": ["PRIV-001", "TS-PRIV-001"],
    },
    {
        "id": "CWE-755",
        "name": "Improper Handling of Exceptional Conditions",
        "description": "Product does not handle or incorrectly handles exceptional conditions",
        "rule_ids": ["ERR-001", "GO-ERR-001"],
    },
    {
        "id": "CWE-532",
        "name": "Insertion of Sensitive Information into Log File",
        "description": "Logging sensitive data that should be protected",
        "rule_ids": ["LOG-001"],
    },
    {
        "id": "CWE-798",
        "name": "Use of Hard-coded Credentials",
        "description": "Product contains hard-coded credentials for inbound/outbound auth",
        "rule_ids": ["SEC-001"],
    },
]


# ---------------------------------------------------------------------------
# MITRE ATT&CK for AI categories
# ---------------------------------------------------------------------------

_MITRE_ATTACK_AI: list[dict[str, Any]] = [
    {
        "id": "AML.T0043",
        "name": "Prompt Injection",
        "description": "Adversary crafts input to manipulate AI agent behavior",
        "rule_ids": ["PI-001", "TS-PI-001", "MCP-PI-001"],
    },
    {
        "id": "AML.T0040",
        "name": "Model Evasion",
        "description": "Adversary evades AI model detection or safety mechanisms",
        "rule_ids": [],
    },
    {
        "id": "AML.T0048",
        "name": "Data Exfiltration via Agent",
        "description": "Using agent capabilities to exfiltrate data",
        "rule_ids": ["INFO-001", "TS-INFO-001", "SSRF-001"],
    },
    {
        "id": "AML.T0044",
        "name": "System/Tool Manipulation",
        "description": "Adversary manipulates tools accessible to the AI agent",
        "rule_ids": ["CD-001", "MCP-TOOL-001", "MCP-TOOL-002", "CMD-001"],
    },
    {
        "id": "AML.T0047",
        "name": "Supply Chain Compromise",
        "description": "Compromise of dependencies or upstream components",
        "rule_ids": [],
    },
    {
        "id": "AML.T0042",
        "name": "Privilege Escalation via Agent",
        "description": "Escalating privileges through agent actions",
        "rule_ids": ["PRIV-001", "TS-PRIV-001", "PERM-001"],
    },
    {
        "id": "AML.T0046",
        "name": "Credential Theft",
        "description": "Stealing credentials through agent interactions",
        "rule_ids": ["SEC-001", "LOG-001"],
    },
    {
        "id": "AML.T0045",
        "name": "Denial of Service",
        "description": "Denying service through resource exhaustion of agent infrastructure",
        "rule_ids": ["RES-001", "RATE-001", "REDOS-001"],
    },
]


# ---------------------------------------------------------------------------
# Analysis functions
# ---------------------------------------------------------------------------

def _get_all_rule_ids() -> set[str]:
    """Collect all rule IDs from all mcpaudit rules."""
    from .rules import get_all_rules

    ids: set[str] = set()
    try:
        rules = get_all_rules()
        for r in rules:
            ids.add(r.rule_id)
    except Exception:
        pass

    try:
        from .ts_rules import get_all_ts_rules
        for r in get_all_ts_rules():
            ids.add(r.rule_id)
    except Exception:
        pass

    try:
        from .go_rules import get_all_go_rules
        for r in get_all_go_rules():
            ids.add(r.rule_id)
    except Exception:
        pass

    return ids


def _analyze_framework(
    framework_name: str,
    categories: list[dict[str, Any]],
) -> CoverageMatrix:
    """Analyze coverage for a given framework.

    Args:
        framework_name: Name of the framework.
        categories: List of category dicts with id, name, description, rule_ids.

    Returns:
        CoverageMatrix with analysis.
    """
    available_rules = _get_all_rule_ids()
    total = len(categories)
    covered_count = 0
    gaps: list[CoverageGap] = []
    covered_list: list[dict[str, Any]] = []

    for cat in categories:
        cat_id = cat["id"]
        cat_name = cat["name"]
        expected_rules = cat.get("rule_ids", [])

        # Check if any of the expected rules exist in our rule set
        matched = [r for r in expected_rules if r in available_rules]

        if matched:
            covered_count += 1
            covered_list.append({
                "category_id": cat_id,
                "category_name": cat_name,
                "rules_matched": matched,
                "total_expected": len(expected_rules),
            })
        else:
            # Determine gap priority based on the number of expected rules
            if not expected_rules:
                priority = "HIGH"  # No rules defined at all
            else:
                priority = "HIGH" if len(expected_rules) >= 2 else "MEDIUM"

            gaps.append(CoverageGap(
                category_id=cat_id,
                category_name=cat_name,
                description=cat["description"],
                priority=priority,
            ))

    coverage_pct = (covered_count / total * 100) if total > 0 else 0.0

    return CoverageMatrix(
        framework_name=framework_name,
        total_categories=total,
        covered_categories=covered_count,
        coverage_pct=coverage_pct,
        gaps=gaps,
        covered=covered_list,
    )


def analyze_owasp_coverage() -> CoverageMatrix:
    """Analyze coverage of OWASP Agentic Top 10.

    Returns:
        CoverageMatrix for OWASP Agentic Top 10.
    """
    return _analyze_framework("OWASP Agentic Top 10", _OWASP_AGENTIC_TOP10)


def analyze_cwe_coverage() -> CoverageMatrix:
    """Analyze coverage of relevant CWE categories.

    Returns:
        CoverageMatrix for CWE categories.
    """
    return _analyze_framework("CWE Categories", _CWE_CATEGORIES)


def analyze_mitre_coverage() -> CoverageMatrix:
    """Analyze coverage of MITRE ATT&CK for AI.

    Returns:
        CoverageMatrix for MITRE ATT&CK for AI.
    """
    return _analyze_framework("MITRE ATT&CK for AI", _MITRE_ATTACK_AI)


def suggest_new_rules(gaps: list[CoverageGap]) -> list[dict[str, Any]]:
    """Suggest new rules to fill coverage gaps.

    Args:
        gaps: List of coverage gaps.

    Returns:
        List of rule suggestion dicts.
    """
    suggestions: list[dict[str, Any]] = []
    for gap in gaps:
        suggestions.append({
            "category_id": gap.category_id,
            "category_name": gap.category_name,
            "priority": gap.priority,
            "suggested_rule_id": f"NEW-{gap.category_id.replace('.', '-').replace(' ', '-')}",
            "description": f"Add rule for: {gap.description}",
            "effort": "MEDIUM" if gap.priority == "HIGH" else "LOW",
        })
    return suggestions


# ---------------------------------------------------------------------------
# Output formatters
# ---------------------------------------------------------------------------

def generate_coverage_html(matrices: list[CoverageMatrix]) -> str:
    """Generate an HTML heatmap dashboard for coverage analysis.

    Args:
        matrices: List of CoverageMatrix from analyze_* functions.

    Returns:
        Self-contained HTML string.
    """
    rows: list[str] = []
    for m in matrices:
        # Color based on coverage
        if m.coverage_pct >= 80:
            color = "#2ecc71"
            bg = "#e8f8f0"
        elif m.coverage_pct >= 60:
            color = "#f39c12"
            bg = "#fef5e7"
        else:
            color = "#e74c3c"
            bg = "#fdedec"

        covered_html: list[str] = []
        for c in m.covered:
            rules = ", ".join(c["rules_matched"])
            covered_html.append(
                f'<tr><td style="padding:4px 8px;border:1px solid #ddd;">{c["category_id"]}</td>'
                f'<td style="padding:4px 8px;border:1px solid #ddd;">{c["category_name"]}</td>'
                f'<td style="padding:4px 8px;border:1px solid #ddd;">{rules}</td></tr>'
            )

        gap_html: list[str] = []
        for g in m.gaps:
            p_color = {"HIGH": "#e74c3c", "MEDIUM": "#f39c12", "LOW": "#3498db"}.get(
                g.priority, "#999"
            )
            gap_html.append(
                f'<tr><td style="padding:4px 8px;border:1px solid #ddd;">{g.category_id}</td>'
                f'<td style="padding:4px 8px;border:1px solid #ddd;">{g.category_name}</td>'
                f'<td style="padding:4px 8px;border:1px solid #ddd;">'
                f'<span style="color:{p_color};font-weight:bold;">{g.priority}</span></td></tr>'
            )

        rows.append(f"""
        <div style="margin-bottom:30px;border:1px solid #ddd;border-radius:8px;padding:16px;">
            <h2 style="margin-top:0;">{m.framework_name}</h2>
            <div style="display:flex;align-items:center;gap:20px;margin-bottom:12px;">
                <div style="background:{bg};border-left:4px solid {color};padding:10px 16px;border-radius:4px;">
                    <span style="font-size:2em;font-weight:bold;color:{color};">{m.coverage_pct:.0f}%</span>
                    <span style="color:#666;margin-left:8px;">{m.covered_categories}/{m.total_categories} categories covered</span>
                </div>
            </div>
            {"<h3>Covered</h3><table style='border-collapse:collapse;width:100%;'><tr><th style='padding:4px 8px;border:1px solid #ddd;text-align:left;'>ID</th><th style='padding:4px 8px;border:1px solid #ddd;text-align:left;'>Category</th><th style='padding:4px 8px;border:1px solid #ddd;text-align:left;'>Rules</th></tr>" + "".join(covered_html) + "</table>" if covered_html else ""}
            {"<h3>Gaps</h3><table style='border-collapse:collapse;width:100%;'><tr><th style='padding:4px 8px;border:1px solid #ddd;text-align:left;'>ID</th><th style='padding:4px 8px;border:1px solid #ddd;text-align:left;'>Category</th><th style='padding:4px 8px;border:1px solid #ddd;text-align:left;'>Priority</th></tr>" + "".join(gap_html) + "</table>" if gap_html else "<p style='color:#2ecc71;'>No gaps found!</p>"}
        </div>
        """)

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<title>mcpaudit — Rule Coverage Analysis</title>
<style>
body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; max-width: 900px; margin: 0 auto; padding: 20px; color: #333; }}
h1 {{ color: #2c3e50; border-bottom: 2px solid #3498db; padding-bottom: 10px; }}
</style>
</head>
<body>
<h1>mcpaudit — Rule Coverage Analysis</h1>
{"".join(rows)}
<footer style="margin-top:30px;color:#999;font-size:0.9em;">
Generated by mcpaudit v{__import__('mcpaudit').__version__}
</footer>
</body>
</html>"""


def format_coverage_text(matrices: list[CoverageMatrix]) -> str:
    """Format coverage matrices as human-readable text.

    Args:
        matrices: List of CoverageMatrix.

    Returns:
        Formatted text.
    """
    lines: list[str] = []
    lines.append("")
    lines.append("=" * 60)
    lines.append("  mcpaudit — Rule Coverage Analysis")
    lines.append("=" * 60)

    for m in matrices:
        lines.append("")
        lines.append(f"  {m.framework_name}")
        lines.append(f"  {'─' * 50}")
        pct_bar = "#" * int(m.coverage_pct / 5) + "." * (20 - int(m.coverage_pct / 5))
        lines.append(
            f"  Coverage: [{pct_bar}] {m.coverage_pct:.0f}% "
            f"({m.covered_categories}/{m.total_categories})"
        )

        if m.covered:
            lines.append("")
            lines.append("  Covered:")
            for c in m.covered:
                rules = ", ".join(c["rules_matched"])
                lines.append(f"    [+] {c['category_id']}: {c['category_name']} ({rules})")

        if m.gaps:
            lines.append("")
            lines.append("  Gaps:")
            for g in m.gaps:
                lines.append(
                    f"    [-] {g.category_id}: {g.category_name} "
                    f"[{g.priority}]"
                )

    lines.append("")
    lines.append("=" * 60)
    return "\n".join(lines)


def format_coverage_json(matrices: list[CoverageMatrix]) -> str:
    """Format coverage matrices as JSON.

    Args:
        matrices: List of CoverageMatrix.

    Returns:
        JSON string.
    """
    data = {
        "coverage_analysis": [m.to_dict() for m in matrices],
    }
    return json.dumps(data, indent=2, ensure_ascii=False)
