"""Risk Calculator for MCP Security Audit.

Calculates composite risk scores for MCP server configurations
based on multiple risk dimensions weighted by impact and likelihood.

Features:
- Multi-dimensional risk assessment (5 dimensions)
- Weighted composite score (0-100)
- Risk level classification: MINIMAL, LOW, MODERATE, HIGH, CRITICAL
- Per-dimension breakdown with recommendations
- Historical risk tracking
- Risk comparison between configs

CLI: mcpaudit risk --score | --breakdown | --compare

Zero external dependencies. Part of mcpaudit v2.1.0.
"""

from __future__ import annotations

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


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

@unique
class RiskLevel(Enum):
    """Risk classification levels."""
    MINIMAL = "MINIMAL"
    LOW = "LOW"
    MODERATE = "MODERATE"
    HIGH = "HIGH"
    CRITICAL = "CRITICAL"

    @property
    def color(self) -> str:
        return {
            RiskLevel.MINIMAL: "green",
            RiskLevel.LOW: "blue",
            RiskLevel.MODERATE: "yellow",
            RiskLevel.HIGH: "orange",
            RiskLevel.CRITICAL: "red",
        }[self]


@unique
class RiskDimension(Enum):
    """Dimensions of risk assessment."""
    ATTACK_SURFACE = "attack_surface"      # Number/type of exposed tools
    DATA_EXPOSURE = "data_exposure"        # Access to sensitive data paths
    PERMISSION_SCOPE = "permission_scope"  # Breadth of granted permissions
    CODE_EXECUTION = "code_execution"      # Ability to run arbitrary code
    NETWORK_ACCESS = "network_access"      # External network capabilities


# Weights for each dimension (must sum to 1.0)
_DIMENSION_WEIGHTS: dict[RiskDimension, float] = {
    RiskDimension.CODE_EXECUTION: 0.30,
    RiskDimension.PERMISSION_SCOPE: 0.25,
    RiskDimension.DATA_EXPOSURE: 0.20,
    RiskDimension.ATTACK_SURFACE: 0.15,
    RiskDimension.NETWORK_ACCESS: 0.10,
}


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class DimensionScore:
    """Score for a single risk dimension."""
    dimension: RiskDimension
    score: float  # 0-100
    factors: tuple[str, ...] = ()
    recommendation: str = ""

    @property
    def level(self) -> RiskLevel:
        return _score_to_level(self.score)

    @property
    def weighted_score(self) -> float:
        weight = _DIMENSION_WEIGHTS.get(self.dimension, 0.2)
        return self.score * weight

    def to_dict(self) -> dict[str, Any]:
        return {
            "dimension": self.dimension.value,
            "score": self.score,
            "level": self.level.value,
            "factors": list(self.factors),
            "recommendation": self.recommendation,
        }


@dataclass
class RiskAssessment:
    """Complete risk assessment result."""
    dimensions: list[DimensionScore] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def composite_score(self) -> float:
        """Weighted composite risk score (0-100)."""
        if not self.dimensions:
            return 0.0
        return sum(d.weighted_score for d in self.dimensions)

    @property
    def risk_level(self) -> RiskLevel:
        return _score_to_level(self.composite_score)

    @property
    def highest_risk(self) -> Optional[DimensionScore]:
        if not self.dimensions:
            return None
        return max(self.dimensions, key=lambda d: d.score)

    @property
    def lowest_risk(self) -> Optional[DimensionScore]:
        if not self.dimensions:
            return None
        return min(self.dimensions, key=lambda d: d.score)

    def to_dict(self) -> dict[str, Any]:
        return {
            "composite_score": round(self.composite_score, 1),
            "risk_level": self.risk_level.value,
            "dimensions": [d.to_dict() for d in self.dimensions],
            "metadata": self.metadata,
        }

    def to_json(self, indent: int = 2) -> str:
        return json.dumps(self.to_dict(), indent=indent)


# ---------------------------------------------------------------------------
# Calculator
# ---------------------------------------------------------------------------

class RiskCalculator:
    """Calculates multi-dimensional risk scores."""

    def __init__(self, custom_weights: Optional[dict[str, float]] = None) -> None:
        self._weights = dict(_DIMENSION_WEIGHTS)
        if custom_weights:
            for key, val in custom_weights.items():
                try:
                    dim = RiskDimension(key)
                    self._weights[dim] = val
                except ValueError:
                    pass

    def assess_tools(self, tools: list[str]) -> DimensionScore:
        """Assess attack surface based on exposed tools."""
        score = 0.0
        factors = []

        high_risk_tools = {"Bash", "MCP", "NotebookEdit"}
        medium_risk_tools = {"Write", "Edit", "WebFetch", "WebSearch"}

        for tool in tools:
            if tool in high_risk_tools:
                score += 25
                factors.append(f"High-risk tool: {tool}")
            elif tool in medium_risk_tools:
                score += 10
                factors.append(f"Medium-risk tool: {tool}")
            else:
                score += 3

        score = min(score, 100)
        rec = ""
        if score > 60:
            rec = "Consider restricting high-risk tools (Bash, MCP)"
        elif score > 30:
            rec = "Review write-capable tools for necessity"

        return DimensionScore(
            dimension=RiskDimension.ATTACK_SURFACE,
            score=score,
            factors=tuple(factors),
            recommendation=rec,
        )

    def assess_permissions(self, sandbox_tier: str,
                           file_patterns: Optional[list[str]] = None) -> DimensionScore:
        """Assess permission scope risk."""
        factors = []
        base_scores = {
            "ReadOnly": 10,
            "WorkspaceWrite": 40,
            "FullAccess": 80,
        }
        score = float(base_scores.get(sandbox_tier, 50))
        factors.append(f"Sandbox tier: {sandbox_tier}")

        if file_patterns:
            wildcards = sum(1 for p in file_patterns if "*" in p or "**" in p)
            if wildcards > 0:
                score += wildcards * 5
                factors.append(f"{wildcards} wildcard file patterns")

        score = min(score, 100)
        rec = ""
        if sandbox_tier == "FullAccess":
            rec = "FullAccess should only be used for DevOps/deploy agents"

        return DimensionScore(
            dimension=RiskDimension.PERMISSION_SCOPE,
            score=score,
            factors=tuple(factors),
            recommendation=rec,
        )

    def assess_code_execution(self, has_bash: bool, has_mcp: bool,
                              max_steps: int = 15) -> DimensionScore:
        """Assess code execution risk."""
        score = 0.0
        factors = []

        if has_bash:
            score += 50
            factors.append("Bash execution enabled")
        if has_mcp:
            score += 30
            factors.append("MCP tool access enabled")
        if max_steps > 30:
            score += 15
            factors.append(f"Extended execution depth: {max_steps} steps")
        elif max_steps > 15:
            score += 5

        score = min(score, 100)
        rec = ""
        if has_bash and has_mcp:
            rec = "Both Bash and MCP enabled — maximum execution risk"

        return DimensionScore(
            dimension=RiskDimension.CODE_EXECUTION,
            score=score,
            factors=tuple(factors),
            recommendation=rec,
        )

    def assess_data_exposure(self, sensitive_paths: list[str],
                             has_web_access: bool = False) -> DimensionScore:
        """Assess data exposure risk."""
        score = 0.0
        factors = []

        high_risk_patterns = {".env", "credentials", "secrets", "private", "key"}
        for path in sensitive_paths:
            path_lower = path.lower()
            if any(p in path_lower for p in high_risk_patterns):
                score += 30
                factors.append(f"Sensitive path access: {path}")
            else:
                score += 5

        if has_web_access:
            score += 20
            factors.append("Web access may exfiltrate data")

        score = min(score, 100)
        rec = ""
        if score > 50:
            rec = "Restrict access to sensitive paths; add file ownership boundaries"

        return DimensionScore(
            dimension=RiskDimension.DATA_EXPOSURE,
            score=score,
            factors=tuple(factors),
            recommendation=rec,
        )

    def assess_network(self, has_web_search: bool, has_web_fetch: bool,
                       has_mcp: bool = False) -> DimensionScore:
        """Assess network access risk."""
        score = 0.0
        factors = []

        if has_web_search:
            score += 20
            factors.append("WebSearch enabled")
        if has_web_fetch:
            score += 30
            factors.append("WebFetch enabled (can download content)")
        if has_mcp:
            score += 25
            factors.append("MCP may connect to external services")

        score = min(score, 100)
        rec = ""
        if score > 50:
            rec = "Consider URL allowlists for web tools"

        return DimensionScore(
            dimension=RiskDimension.NETWORK_ACCESS,
            score=score,
            factors=tuple(factors),
            recommendation=rec,
        )

    def full_assessment(
        self,
        tools: list[str],
        sandbox_tier: str = "WorkspaceWrite",
        file_patterns: Optional[list[str]] = None,
        sensitive_paths: Optional[list[str]] = None,
        max_steps: int = 15,
    ) -> RiskAssessment:
        """Run a full risk assessment."""
        has_bash = "Bash" in tools
        has_mcp = "MCP" in tools
        has_web_search = "WebSearch" in tools
        has_web_fetch = "WebFetch" in tools

        dimensions = [
            self.assess_tools(tools),
            self.assess_permissions(sandbox_tier, file_patterns),
            self.assess_code_execution(has_bash, has_mcp, max_steps),
            self.assess_data_exposure(sensitive_paths or [], has_web_search or has_web_fetch),
            self.assess_network(has_web_search, has_web_fetch, has_mcp),
        ]

        return RiskAssessment(
            dimensions=dimensions,
            metadata={
                "tools_count": len(tools),
                "sandbox_tier": sandbox_tier,
                "max_steps": max_steps,
            },
        )

    @staticmethod
    def compare(a: RiskAssessment, b: RiskAssessment) -> dict[str, Any]:
        """Compare two risk assessments."""
        a_scores = {d.dimension: d.score for d in a.dimensions}
        b_scores = {d.dimension: d.score for d in b.dimensions}

        comparison = {
            "composite_delta": round(b.composite_score - a.composite_score, 1),
            "direction": "IMPROVED" if b.composite_score < a.composite_score else (
                "DEGRADED" if b.composite_score > a.composite_score else "UNCHANGED"
            ),
            "dimensions": {},
        }

        for dim in RiskDimension:
            a_val = a_scores.get(dim, 0)
            b_val = b_scores.get(dim, 0)
            comparison["dimensions"][dim.value] = {
                "before": a_val,
                "after": b_val,
                "delta": round(b_val - a_val, 1),
            }

        return comparison

    def markdown_report(self, assessment: RiskAssessment) -> str:
        """Generate Markdown risk report."""
        lines = ["# Risk Assessment Report", ""]

        lines.append("## Summary")
        lines.append(f"- **Composite Score**: {assessment.composite_score:.1f}/100")
        lines.append(f"- **Risk Level**: {assessment.risk_level.value}")
        lines.append("")

        lines.append("## Dimensions")
        lines.append("| Dimension | Score | Level | Recommendation |")
        lines.append("|-----------|-------|-------|----------------|")
        for d in sorted(assessment.dimensions, key=lambda x: -x.score):
            lines.append(
                f"| {d.dimension.value} | {d.score:.0f} | {d.level.value} "
                f"| {d.recommendation or '-'} |"
            )
        lines.append("")

        # Factors
        lines.append("## Risk Factors")
        for d in assessment.dimensions:
            if d.factors:
                lines.append(f"### {d.dimension.value}")
                for f in d.factors:
                    lines.append(f"- {f}")
                lines.append("")

        return "\n".join(lines)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _score_to_level(score: float) -> RiskLevel:
    """Convert numeric score to risk level."""
    if score <= 15:
        return RiskLevel.MINIMAL
    if score <= 35:
        return RiskLevel.LOW
    if score <= 55:
        return RiskLevel.MODERATE
    if score <= 75:
        return RiskLevel.HIGH
    return RiskLevel.CRITICAL
