"""Type definitions for Aegis SDK"""

from enum import Enum
from typing import Any, Dict, List, Optional
from datetime import datetime
from pydantic import BaseModel, Field


class TrustLevel(str, Enum):
    """Trust level hierarchy"""
    TRUSTED_SYSTEM = "trusted.system"
    TRUSTED_DEVELOPER = "trusted.developer"
    UNTRUSTED_USER = "untrusted.user"
    UNTRUSTED_DOCUMENT = "untrusted.document"
    UNTRUSTED_EXTERNAL = "untrusted.external"


class SensitivityLevel(str, Enum):
    """Sensitivity level hierarchy"""
    PUBLIC = "public"
    INTERNAL = "internal"
    CONFIDENTIAL = "confidential"
    PRIVATE = "private"
    SECRET = "secret"


class RiskLevel(str, Enum):
    """Risk level for capabilities"""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


class RunMode(str, Enum):
    """Gateway run mode"""
    ENFORCE = "enforce"
    AUDIT = "audit"
    MONITOR = "monitor"
    DRY_RUN = "dry-run"


class SpanSource(str, Enum):
    """Span source type"""
    SYSTEM = "system"
    USER = "user"
    DOCUMENT = "document"
    TOOL_OUTPUT = "tool_output"
    MODEL_REASONING = "model_reasoning"


class Span(BaseModel):
    """Span representation"""
    id: str
    content: str
    source: SpanSource
    trust_level: TrustLevel
    sensitivity: SensitivityLevel
    hash: str
    timestamp: datetime
    metadata: Optional[Dict[str, Any]] = None


class Capability(BaseModel):
    """Capability definition"""
    full_name: str
    namespace: str
    action: str
    risk_level: RiskLevel
    requires_approval: bool
    description: Optional[str] = None


class ToolCallDeclaration(BaseModel):
    """Tool call declaration with source spans"""
    sources: List[Dict[str, Any]]


class ToolCall(BaseModel):
    """Tool call request"""
    capability: str
    arguments: Dict[str, Any]
    declaration: Optional[ToolCallDeclaration] = None


class ToolCallResult(BaseModel):
    """Tool call result"""
    allowed: bool
    reason: Optional[str] = None
    approval_id: Optional[str] = None
    requires_approval: bool = False


class PolicyDecision(BaseModel):
    """Policy decision"""
    timestamp: datetime
    capability: str
    result: str
    reason: str
    trust_level: TrustLevel
    sensitivity: SensitivityLevel


class ApprovalRequest(BaseModel):
    """Approval request"""
    id: str
    capability: str
    arguments: Dict[str, Any]
    requested_at: datetime
    expires_at: datetime
    status: str
    reason: Optional[str] = None


class TraceSession(BaseModel):
    """Trace session"""
    session_id: str
    timestamp_start: datetime
    timestamp_end: Optional[datetime] = None
    upstream_model: str
    policy_hash: str
    spans: List[Span]
    decisions: List[PolicyDecision]


class GatewayConfig(BaseModel):
    """Gateway configuration"""
    listen_addr: str
    upstream_provider: str
    upstream_endpoint: Optional[str] = None
    policy_path: str
    mode: RunMode = RunMode.ENFORCE
    trace_dir: Optional[str] = None
    api_key: Optional[str] = None
