"""
Data Models

Pydantic models for NeuralDLP API requests and responses.
"""
from typing import List, Optional, Dict, Any
from datetime import datetime
from dataclasses import dataclass


@dataclass
class Entity:
    """Detected sensitive entity"""
    type: str
    value: str
    start: int
    end: int
    confidence: float
    
    @classmethod
    def from_dict(cls, data: dict) -> "Entity":
        return cls(**data)
    
    def to_dict(self) -> dict:
        return {
            "type": self.type,
            "value": self.value,
            "start": self.start,
            "end": self.end,
            "confidence": self.confidence
        }


@dataclass
class Threat:
    """Detected security threat"""
    type: str
    confidence: float
    severity: str
    description: str
    
    @classmethod
    def from_dict(cls, data: dict) -> "Threat":
        return cls(**data)


@dataclass
class PolicyMatch:
    """Matched security policy"""
    policy_id: str
    policy_name: str
    action: str
    priority: int
    matched_rules: List[str]
    
    @classmethod
    def from_dict(cls, data: dict) -> "PolicyMatch":
        return cls(**data)


@dataclass
class InspectionResult:
    """Result of input/output inspection"""
    request_id: str
    risk_score: int
    entities: List[Entity]
    threats: List[Threat]
    matched_policies: List[str]
    recommended_action: str  # ALLOW, TOKENIZE, BLOCK
    processing_time_ms: float
    
    # Output inspection specific
    leakage_detected: Optional[bool] = None
    placeholder_integrity: Optional[bool] = None
    unexpected_entities: Optional[List[Entity]] = None
    
    @classmethod
    def from_dict(cls, data: dict) -> "InspectionResult":
        return cls(
            request_id=data["request_id"],
            risk_score=data["risk_score"],
            entities=[Entity.from_dict(e) for e in data["entities"]],
            threats=[Threat.from_dict(t) for t in data["threats"]],
            matched_policies=data["matched_policies"],
            recommended_action=data["recommended_action"],
            processing_time_ms=data["processing_time_ms"],
            leakage_detected=data.get("leakage_detected"),
            placeholder_integrity=data.get("placeholder_integrity"),
            unexpected_entities=[Entity.from_dict(e) for e in data.get("unexpected_entities", [])]
        )
    
    def is_safe(self) -> bool:
        """Check if content is safe (ALLOW action)"""
        return self.recommended_action == "ALLOW"
    
    def should_tokenize(self) -> bool:
        """Check if content should be tokenized"""
        return self.recommended_action == "TOKENIZE"
    
    def is_blocked(self) -> bool:
        """Check if content is blocked"""
        return self.recommended_action == "BLOCK"


@dataclass
class TokenInfo:
    """Information about a generated token"""
    token: str
    original_type: str
    start: int
    end: int
    
    @classmethod
    def from_dict(cls, data: dict) -> "TokenInfo":
        return cls(**data)


@dataclass
class TokenizationResult:
    """Result of tokenization"""
    request_id: str
    transformed_text: str
    mapping_id: str
    expires_at: datetime
    tokens: List[TokenInfo]
    processing_time_ms: float
    
    @classmethod
    def from_dict(cls, data: dict) -> "TokenizationResult":
        return cls(
            request_id=data["request_id"],
            transformed_text=data["transformed_text"],
            mapping_id=data["mapping_id"],
            expires_at=datetime.fromisoformat(data["expires_at"].replace("Z", "+00:00")),
            tokens=[TokenInfo.from_dict(t) for t in data["tokens"]],
            processing_time_ms=data["processing_time_ms"]
        )


@dataclass
class RestoredEntity:
    """Information about a restored entity"""
    token: str
    restored_value: str
    confidence: float
    
    @classmethod
    def from_dict(cls, data: dict) -> "RestoredEntity":
        return cls(**data)


@dataclass
class RehydrationResult:
    """Result of rehydration"""
    request_id: str
    rehydrated_text: str
    confidence: float
    restored_entities: List[RestoredEntity]
    processing_time_ms: float
    
    @classmethod
    def from_dict(cls, data: dict) -> "RehydrationResult":
        return cls(
            request_id=data["request_id"],
            rehydrated_text=data["rehydrated_text"],
            confidence=data["confidence"],
            restored_entities=[RestoredEntity.from_dict(e) for e in data["restored_entities"]],
            processing_time_ms=data["processing_time_ms"]
        )


@dataclass
class MappingStatus:
    """Status of a token mapping"""
    mapping_id: str
    exists: bool
    ttl_seconds: int
    expires_at: datetime
    
    @classmethod
    def from_dict(cls, data: dict) -> "MappingStatus":
        return cls(
            mapping_id=data["mapping_id"],
            exists=data["exists"],
            ttl_seconds=data["ttl_seconds"],
            expires_at=datetime.fromisoformat(data["expires_at"].replace("Z", "+00:00"))
        )
