"""
PromptShield — Data Models
===========================

Maintainer : Joseph Mutua
Version    : 1.0.0

All request / response models used throughout the PromptShield engine.
Using Pydantic v2 for strict runtime validation and automatic serialization.

These models are shared between:
  - The HTTP API (app.py routes)
  - The SDK middleware (middleware.py)
  - The detection engine (scanner.py)

Extending this file
-------------------
Add new fields to ScanResult or ThreatMatch as the engine evolves.
Pydantic handles backward compatibility via Optional fields with defaults.
"""

from __future__ import annotations

import logging
from enum import Enum
from typing import List, Literal, Optional

from pydantic import BaseModel, Field, field_validator

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Enumerations
# ---------------------------------------------------------------------------

class ThreatLevel(str, Enum):
    """
    Overall threat classification produced by the scanner.

    Maps directly to score ranges:
        SAFE     0.00 – 0.19
        LOW      0.20 – 0.39
        MEDIUM   0.40 – 0.59
        HIGH     0.60 – 0.79
        CRITICAL 0.80 – 1.00
    """
    SAFE     = "safe"
    LOW      = "low"
    MEDIUM   = "medium"
    HIGH     = "high"
    CRITICAL = "critical"


class ActionTaken(str, Enum):
    """
    What the scanner did with the content after evaluating it.

    PASSED      — content is clean, returned unchanged.
    ANNOTATED   — suspicious content passed through with threat metadata.
    SANITIZED   — injected instructions removed, cleaned version returned.
    QUARANTINED — content is fully blocked, a safe placeholder returned.
    """
    PASSED      = "passed"
    ANNOTATED   = "annotated"
    SANITIZED   = "sanitized"
    QUARANTINED = "quarantined"


class DetectionLayer(str, Enum):
    """
    The detection layer that produced a ThreatMatch.

    Layer 4 (linguistic) matches are mapped to SEMANTIC since the schema
    predates Layer 4. They are distinguishable by the "Layer 4 —" prefix
    in ThreatMatch.description.
    """
    PATTERN   = "pattern"    # Layer 1: regex signatures
    HEURISTIC = "heuristic"  # Layer 2: statistical / structural signals
    SEMANTIC  = "semantic"   # Layer 3: TF-IDF similarity (+ Layer 4)


# ---------------------------------------------------------------------------
# Core models
# ---------------------------------------------------------------------------

class ThreatMatch(BaseModel):
    """
    A single threat signal that fired during a scan.

    One ScanResult may contain many ThreatMatches from all four layers.
    Each match records which layer caught it, the attack category,
    a human-readable description, the evidence snippet, and confidence.
    """

    layer:       DetectionLayer
    category:    str              # e.g. "instruction_override", "role_hijack"
    description: str              # human-readable explanation for reports
    evidence:    str              # the exact text snippet that triggered it
    confidence:  float = Field(ge=0.0, le=1.0)


class ScanRequest(BaseModel):
    """
    Input to the scanner. Used by the HTTP API and internal engine.

    Fields:
        content       — The text to scan. Required. Max 50 000 chars.
        source_type   — What kind of content this is (informational).
        context       — Optional description of what the agent was doing
                        (helps with future contextual scoring).
        action_policy — Override default auto action:
                        "auto"       — choose action by threat level
                        "annotate"   — always annotate, never sanitize
                        "sanitize"   — always sanitize if threat found
                        "quarantine" — always quarantine if threat found
    """

    content:       str = Field(..., min_length=1, max_length=50_000)
    source_type:   Literal["text", "document", "url_content", "api_response"] = "text"
    context:       Optional[str] = None
    action_policy: Literal["annotate", "sanitize", "quarantine", "auto"] = "auto"

    @field_validator("content")
    @classmethod
    def strip_content(cls, v: str) -> str:
        """Strip leading/trailing whitespace from content before scanning."""
        return v.strip()


class ScanResult(BaseModel):
    """
    The complete output of a scan. Returned by scanner.scan() and the HTTP API.

    Fields
    ------
    threat_level    — Overall classification (SAFE / LOW / MEDIUM / HIGH / CRITICAL).
    action_taken    — What the scanner did with the content.
    is_safe         — Convenience bool (True if threat_level == SAFE).
    injection_score — Final fused score 0.0–1.0. Higher = more dangerous.
    pattern_score   — Layer 1 (pattern) contribution.
    heuristic_score — Layer 2 (heuristic) contribution.
    semantic_score  — Layer 3 + 4 combined contribution.
    original_content— The original, unmodified input text.
    safe_content    — The text to pass to the LLM. Cleaned if sanitized,
                      placeholder if quarantined, original otherwise.
    matches         — All signals that fired, from all four layers.
    scan_duration_ms— Wall-clock scan time in milliseconds.
    total_matches   — Number of signals that fired.
    primary_category— The dominant attack category detected, if any.
    summary         — Plain-English one-liner verdict.
    recommendation  — What the developer / agent should do next.
    """

    # Verdict
    threat_level:    ThreatLevel
    action_taken:    ActionTaken
    is_safe:         bool

    # Scores
    injection_score: float = Field(ge=0.0, le=1.0)
    pattern_score:   float = Field(ge=0.0, le=1.0)
    heuristic_score: float = Field(ge=0.0, le=1.0)
    semantic_score:  float = Field(ge=0.0, le=1.0)

    # Content
    original_content: str
    safe_content:     str
    matches:          List[ThreatMatch] = []

    # Meta
    scan_duration_ms: float
    total_matches:    int
    primary_category: Optional[str] = None

    # Human-readable explanations
    summary:        str
    recommendation: str


class BatchScanRequest(BaseModel):
    """Scan multiple texts in a single API call (max 20)."""
    items: List[ScanRequest] = Field(..., min_length=1, max_length=20)


class BatchScanResult(BaseModel):
    """Results for a batch scan."""
    results:       List[ScanResult]
    total_scanned: int
    threats_found: int
    safe_count:    int


class HealthResponse(BaseModel):
    """Response from the /api/health endpoint."""
    status:          str
    version:         str
    patterns_loaded: int
    semantic_corpus: int
    uptime_seconds:  float
