"""
Core data models for MCP Security Auditor.

Defines the structure for findings, scan results, and related entities.
"""

from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional
from enum import Enum
from datetime import datetime
import hashlib
import json


class Severity(str, Enum):
    """Severity levels for security findings."""
    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    INFO = "info"
    
    @classmethod
    def from_string(cls, value: str) -> "Severity":
        """Parse severity from string."""
        return cls(value.lower())
    
    def to_score(self) -> int:
        """Convert severity to numeric score (0-10)."""
        scores = {
            Severity.CRITICAL: 10,
            Severity.HIGH: 8,
            Severity.MEDIUM: 5,
            Severity.LOW: 3,
            Severity.INFO: 1
        }
        return scores.get(self, 0)


class FindingCategory(str, Enum):
    """Categories of security findings."""
    INJECTION = "injection"
    PERMISSIONS = "permissions"
    NETWORK = "network"
    SECRETS = "secrets"
    DEPENDENCIES = "dependencies"
    CONFIGURATION = "configuration"
    CODE_QUALITY = "code_quality"
    AUTHENTICATION = "authentication"
    DATA_EXPOSURE = "data_exposure"
    CRYPTOGRAPHY = "cryptography"


@dataclass
class CodeLocation:
    """Location of a finding in source code."""
    file: str
    line_start: int
    line_end: int
    column_start: Optional[int] = None
    column_end: Optional[int] = None
    snippet: Optional[str] = None
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "file": self.file,
            "line_start": self.line_start,
            "line_end": self.line_end,
            "column_start": self.column_start,
            "column_end": self.column_end,
            "snippet": self.snippet
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "CodeLocation":
        return cls(
            file=data["file"],
            line_start=data["line_start"],
            line_end=data["line_end"],
            column_start=data.get("column_start"),
            column_end=data.get("column_end"),
            snippet=data.get("snippet")
        )


@dataclass
class Remediation:
    """Remediation guidance for a finding."""
    description: str
    code_example: Optional[str] = None
    references: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "description": self.description,
            "code_example": self.code_example,
            "references": self.references
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Remediation":
        return cls(
            description=data["description"],
            code_example=data.get("code_example"),
            references=data.get("references", [])
        )


@dataclass
class Finding:
    """A security finding from the scan."""
    title: str
    description: str
    severity: str
    category: str
    analyzer: str
    location: Optional[CodeLocation] = None
    remediation: Optional[Remediation] = None
    cwe_id: Optional[str] = None
    cvss_score: Optional[float] = None
    confidence: str = "high"  # high, medium, low
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def __post_init__(self):
        """Generate unique ID after initialization."""
        self.id = self._generate_id()
    
    def _generate_id(self) -> str:
        """Generate unique ID based on finding properties."""
        content = f"{self.title}{self.category}{self.analyzer}"
        if self.location:
            content += f"{self.location.file}{self.location.line_start}"
        return hashlib.sha256(content.encode()).hexdigest()[:16]
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "title": self.title,
            "description": self.description,
            "severity": self.severity,
            "category": self.category,
            "analyzer": self.analyzer,
            "location": self.location.to_dict() if self.location else None,
            "remediation": self.remediation.to_dict() if self.remediation else None,
            "cwe_id": self.cwe_id,
            "cvss_score": self.cvss_score,
            "confidence": self.confidence,
            "metadata": self.metadata
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Finding":
        location = CodeLocation.from_dict(data["location"]) if data.get("location") else None
        remediation = Remediation.from_dict(data["remediation"]) if data.get("remediation") else None
        return cls(
            title=data["title"],
            description=data["description"],
            severity=data["severity"],
            category=data["category"],
            analyzer=data["analyzer"],
            location=location,
            remediation=remediation,
            cwe_id=data.get("cwe_id"),
            cvss_score=data.get("cvss_score"),
            confidence=data.get("confidence", "high"),
            metadata=data.get("metadata", {})
        )


@dataclass
class ToolDefinition:
    """Extracted MCP tool definition."""
    name: str
    description: Optional[str] = None
    parameters: Dict[str, Any] = field(default_factory=dict)
    annotations: Dict[str, Any] = field(default_factory=dict)
    file_path: Optional[str] = None
    line_number: Optional[int] = None
    
    @property
    def is_read_only(self) -> bool:
        return self.annotations.get("readOnlyHint", False)
    
    @property
    def is_destructive(self) -> bool:
        return self.annotations.get("destructiveHint", False)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "parameters": self.parameters,
            "annotations": self.annotations,
            "file_path": self.file_path,
            "line_number": self.line_number,
            "is_read_only": self.is_read_only,
            "is_destructive": self.is_destructive
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ToolDefinition":
        return cls(
            name=data["name"],
            description=data.get("description"),
            parameters=data.get("parameters", {}),
            annotations=data.get("annotations", {}),
            file_path=data.get("file_path"),
            line_number=data.get("line_number")
        )


@dataclass
class PermissionScope:
    """Permission scope requested by MCP server."""
    name: str
    level: str  # read, write, execute, admin
    resource: Optional[str] = None
    justification: Optional[str] = None
    file_path: Optional[str] = None
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "level": self.level,
            "resource": self.resource,
            "justification": self.justification,
            "file_path": self.file_path
        }


@dataclass
class NetworkEndpoint:
    """Network endpoint accessed by MCP server."""
    url: str
    method: Optional[str] = None
    purpose: Optional[str] = None
    file_path: Optional[str] = None
    line_number: Optional[int] = None
    is_hardcoded: bool = False
    
    @property
    def domain(self) -> str:
        """Extract domain from URL."""
        from urllib.parse import urlparse
        try:
            parsed = urlparse(self.url)
            return parsed.netloc
        except:
            return self.url
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "url": self.url,
            "domain": self.domain,
            "method": self.method,
            "purpose": self.purpose,
            "file_path": self.file_path,
            "line_number": self.line_number,
            "is_hardcoded": self.is_hardcoded
        }


@dataclass
class Dependency:
    """A dependency of the MCP server."""
    name: str
    version: Optional[str] = None
    source: str = "unknown"  # npm, pypi, etc.
    vulnerabilities: List[Dict[str, Any]] = field(default_factory=list)
    
    @property
    def has_vulnerabilities(self) -> bool:
        return len(self.vulnerabilities) > 0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "version": self.version,
            "source": self.source,
            "vulnerabilities": self.vulnerabilities,
            "has_vulnerabilities": self.has_vulnerabilities
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Dependency":
        return cls(
            name=data["name"],
            version=data.get("version"),
            source=data.get("source", "unknown"),
            vulnerabilities=data.get("vulnerabilities", [])
        )


@dataclass
class MCPServerInfo:
    """Information about the scanned MCP server."""
    name: str
    path: str
    language: str  # python, typescript, javascript
    framework: Optional[str] = None  # fastmcp, mcp-sdk, etc.
    version: Optional[str] = None
    tools: List[ToolDefinition] = field(default_factory=list)
    resources: List[Dict[str, Any]] = field(default_factory=list)
    permissions: List[PermissionScope] = field(default_factory=list)
    endpoints: List[NetworkEndpoint] = field(default_factory=list)
    dependencies: List[Dependency] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "path": self.path,
            "language": self.language,
            "framework": self.framework,
            "version": self.version,
            "tools": [t.to_dict() for t in self.tools],
            "resources": self.resources,
            "permissions": [p.to_dict() for p in self.permissions],
            "endpoints": [e.to_dict() for e in self.endpoints],
            "dependencies": [d.to_dict() for d in self.dependencies]
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "MCPServerInfo":
        return cls(
            name=data["name"],
            path=data["path"],
            language=data["language"],
            framework=data.get("framework"),
            version=data.get("version"),
            tools=[ToolDefinition.from_dict(t) for t in data.get("tools", [])],
            resources=data.get("resources", []),
            dependencies=[Dependency.from_dict(d) for d in data.get("dependencies", [])]
        )


@dataclass
class ScanResult:
    """Complete results of a security scan."""
    scan_id: str
    scan_date: str
    target_path: str
    scanner_version: str
    server_info: MCPServerInfo
    findings: List[Finding]
    risk_score: int  # 0-100
    scan_duration_seconds: float
    analyzers_run: List[str]
    
    def __post_init__(self):
        """Calculate statistics after initialization."""
        self.findings_by_severity = self._count_by_severity()
        self.findings_by_category = self._count_by_category()
    
    def _count_by_severity(self) -> Dict[str, int]:
        """Count findings by severity."""
        counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
        for finding in self.findings:
            counts[finding.severity] = counts.get(finding.severity, 0) + 1
        return counts
    
    def _count_by_category(self) -> Dict[str, int]:
        """Count findings by category."""
        counts = {}
        for finding in self.findings:
            counts[finding.category] = counts.get(finding.category, 0) + 1
        return counts
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "scan_id": self.scan_id,
            "scan_date": self.scan_date,
            "target_path": self.target_path,
            "scanner_version": self.scanner_version,
            "server_info": self.server_info.to_dict(),
            "findings": [f.to_dict() for f in self.findings],
            "findings_by_severity": self.findings_by_severity,
            "findings_by_category": self.findings_by_category,
            "risk_score": self.risk_score,
            "scan_duration_seconds": self.scan_duration_seconds,
            "analyzers_run": self.analyzers_run
        }
    
    def to_json(self, indent: int = 2) -> str:
        """Serialize to JSON string."""
        return json.dumps(self.to_dict(), indent=indent)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ScanResult":
        """Deserialize from dictionary (e.g. loaded from JSON)."""
        server_info = MCPServerInfo.from_dict(data["server_info"])
        findings = [Finding.from_dict(f) for f in data.get("findings", [])]
        return cls(
            scan_id=data["scan_id"],
            scan_date=data["scan_date"],
            target_path=data["target_path"],
            scanner_version=data["scanner_version"],
            server_info=server_info,
            findings=findings,
            risk_score=data["risk_score"],
            scan_duration_seconds=data.get("scan_duration_seconds", 0.0),
            analyzers_run=data.get("analyzers_run", [])
        )
    
    @classmethod
    def from_json(cls, json_str: str) -> "ScanResult":
        """Deserialize from JSON string."""
        return cls.from_dict(json.loads(json_str))


@dataclass
class InjectionTestCase:
    """A test case for prompt injection detection."""
    name: str
    payload: str
    category: str
    expected_behavior: str
    severity: str
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "payload": self.payload,
            "category": self.category,
            "expected_behavior": self.expected_behavior,
            "severity": self.severity
        }
