"""Vulnerability models for PromptSentry."""

from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field


class VulnerabilitySeverity(str, Enum):
    """Severity levels for vulnerabilities."""

    CRITICAL = "CRITICAL"
    HIGH = "HIGH"
    MEDIUM = "MEDIUM"
    LOW = "LOW"
    INFO = "INFO"

    @property
    def color(self) -> str:
        """Rich color for display."""
        return {
            "CRITICAL": "red bold",
            "HIGH": "red",
            "MEDIUM": "yellow",
            "LOW": "blue",
            "INFO": "dim",
        }[self.value]

    @property
    def emoji(self) -> str:
        """Emoji for display."""
        return {
            "CRITICAL": "🔴",
            "HIGH": "🟠",
            "MEDIUM": "🟡",
            "LOW": "🔵",
            "INFO": "⚪",
        }[self.value]

    @property
    def score(self) -> int:
        """Numeric score for severity."""
        return {
            "CRITICAL": 100,
            "HIGH": 80,
            "MEDIUM": 50,
            "LOW": 25,
            "INFO": 10,
        }[self.value]


class OWASPCategory(str, Enum):
    """OWASP LLM Top 10 categories."""

    LLM01 = "LLM01: Prompt Injection"
    LLM02 = "LLM02: Insecure Output Handling"
    LLM03 = "LLM03: Training Data Poisoning"
    LLM04 = "LLM04: Model Denial of Service"
    LLM05 = "LLM05: Supply Chain Vulnerabilities"
    LLM06 = "LLM06: Sensitive Information Disclosure"
    LLM07 = "LLM07: Insecure Plugin Design"
    LLM08 = "LLM08: Excessive Agency"
    LLM09 = "LLM09: Overreliance"
    LLM10 = "LLM10: Model Theft"


class Vulnerability(BaseModel):
    """A detected vulnerability in an AI prompt."""

    vuln_id: str = Field(..., description="Unique vulnerability identifier")
    vuln_type: str = Field(..., description="Type of vulnerability (e.g., DIRECT_CONCATENATION)")
    severity: VulnerabilitySeverity = Field(..., description="Severity level")
    owasp_category: Optional[OWASPCategory] = Field(None, description="OWASP LLM Top 10 category")
    location: str = Field(..., description="Location in code (file:line)")
    vulnerable_code: str = Field(..., description="The vulnerable code snippet")
    description: str = Field(..., description="Description of the vulnerability")
    fix: str = Field(..., description="Recommended fix")
    confidence: float = Field(1.0, ge=0.0, le=1.0, description="Confidence in the finding")

    @property
    def short_code(self) -> str:
        """Truncated vulnerable code for display."""
        code = self.vulnerable_code.replace("\n", " ").strip()
        if len(code) <= 60:
            return code
        return code[:57] + "..."


class PatternMatch(BaseModel):
    """Result from pattern-based detection (Stage 2)."""

    pattern_id: str = Field(..., description="Pattern identifier")
    pattern_name: str = Field(..., description="Human-readable pattern name")
    severity: VulnerabilitySeverity = Field(..., description="Severity level")
    owasp_category: OWASPCategory = Field(..., description="OWASP category")
    location: str = Field(..., description="Location in code")
    matched_code: str = Field(..., description="Code that matched the pattern")
    description: str = Field(..., description="Description of the issue")
    fix: str = Field(..., description="Recommended fix")


class SimilarMatch(BaseModel):
    """Result from vector similarity search (Stage 3)."""

    rule_id: str = Field(..., description="Rule identifier from vector DB")
    similarity: float = Field(..., ge=0.0, le=1.0, description="Similarity score")
    example: str = Field(..., description="Similar vulnerable example")
    vulnerability: str = Field(..., description="Type of vulnerability")
    description: str = Field(..., description="Description of vulnerability")
    fix: str = Field(..., description="Recommended fix")
    owasp_category: Optional[OWASPCategory] = Field(None, description="OWASP category")


class AnalysisResult(BaseModel):
    """Complete analysis result for a prompt."""

    file_path: str = Field(..., description="Path to analyzed file")
    prompt_location: str = Field(..., description="Location of the prompt")
    vulnerabilities: list[Vulnerability] = Field(
        default_factory=list,
        description="List of detected vulnerabilities"
    )
    overall_score: int = Field(
        0, ge=0, le=100,
        description="Overall vulnerability score (0=safe, 100=critical)"
    )
    is_vulnerable: bool = Field(False, description="Whether the prompt is vulnerable")
    pattern_matches: list[PatternMatch] = Field(
        default_factory=list,
        description="Pattern-based findings"
    )
    similar_matches: list[SimilarMatch] = Field(
        default_factory=list,
        description="Vector similarity findings"
    )
    llm_analysis: Optional[str] = Field(None, description="Raw LLM analysis output")

    @property
    def severity_counts(self) -> dict:
        """Count vulnerabilities by severity."""
        counts = {s: 0 for s in VulnerabilitySeverity}
        for vuln in self.vulnerabilities:
            counts[vuln.severity] += 1
        return counts

    @property
    def highest_severity(self) -> Optional[VulnerabilitySeverity]:
        """Get the highest severity found."""
        if not self.vulnerabilities:
            return None
        return max(self.vulnerabilities, key=lambda v: v.severity.score).severity
