"""DataShield Types - Pydantic models and enums for data scrambling.

Defines the configuration schema for shield projects, table configs,
column rules, and scrambling strategies.
"""

from enum import Enum
from pydantic import BaseModel, Field
from typing import Optional, List, Literal
from datetime import datetime
import uuid


# =============================================================================
# Enumerations
# =============================================================================

class ScrambleStrategy(str, Enum):
    """Scrambling strategy for a column."""
    FORMAT_PRESERVING_HASH = "format_preserving_hash"
    NUMERIC_SCALING = "numeric_scaling"
    SYNTHETIC_SUBSTITUTION = "synthetic_substitution"
    DATE_SHIFT = "date_shift"
    PATTERN_PRESERVING = "pattern_preserving"
    PASSTHROUGH = "passthrough"


class ColumnClassification(str, Enum):
    """Classification of a column's data type and sensitivity."""
    MEASURE = "measure"
    FACT_DIMENSION = "fact_dimension"
    DESCRIPTIVE = "descriptive"
    GEOGRAPHIC = "geographic"
    TEMPORAL = "temporal"
    IDENTIFIER = "identifier"
    CODE = "code"
    SENSITIVE_PII = "sensitive_pii"
    SAFE = "safe"


class ShieldScope(str, Enum):
    """Scope of shield application."""
    TABLE = "table"
    COLUMN = "column"
    SCHEMA = "schema"


class TrustLane(str, Enum):
    """Execution lane for trust policy enforcement."""
    DATABRIDGE_AI_MASKED = "databridge_ai_masked"
    CLIENT_CORTEX_RAW = "client_cortex_raw"


# =============================================================================
# Default strategy mapping per classification
# =============================================================================

CLASSIFICATION_STRATEGY_MAP = {
    ColumnClassification.MEASURE: ScrambleStrategy.NUMERIC_SCALING,
    ColumnClassification.FACT_DIMENSION: ScrambleStrategy.FORMAT_PRESERVING_HASH,
    ColumnClassification.DESCRIPTIVE: ScrambleStrategy.SYNTHETIC_SUBSTITUTION,
    ColumnClassification.GEOGRAPHIC: ScrambleStrategy.SYNTHETIC_SUBSTITUTION,
    ColumnClassification.TEMPORAL: ScrambleStrategy.DATE_SHIFT,
    ColumnClassification.IDENTIFIER: ScrambleStrategy.FORMAT_PRESERVING_HASH,
    ColumnClassification.CODE: ScrambleStrategy.FORMAT_PRESERVING_HASH,
    ColumnClassification.SENSITIVE_PII: ScrambleStrategy.PATTERN_PRESERVING,
    ColumnClassification.SAFE: ScrambleStrategy.PASSTHROUGH,
}


# =============================================================================
# Core Models
# =============================================================================

# =============================================================================
# Schema Monitoring Models
# =============================================================================

class ColumnSnapshot(BaseModel):
    """Snapshot of a single column's metadata."""
    name: str
    data_type: str
    nullable: bool = True
    ordinal_position: int = 0


class SchemaSnapshot(BaseModel):
    """Point-in-time snapshot of a table's schema."""
    database: str
    schema_name: str
    table_name: str
    columns: List[ColumnSnapshot] = []
    captured_at: str = Field(default_factory=lambda: datetime.now().isoformat())


class SchemaChangeEvent(BaseModel):
    """A detected schema change."""
    detected_at: str = Field(default_factory=lambda: datetime.now().isoformat())
    change_type: Literal["column_added", "column_removed", "type_changed"]
    column_name: str
    old_value: Optional[str] = None
    new_value: Optional[str] = None


# =============================================================================
# Core Models
# =============================================================================

class ColumnRule(BaseModel):
    """Per-column scrambling configuration."""
    column_name: str
    classification: ColumnClassification
    strategy: ScrambleStrategy
    preserve_nulls: bool = True
    preserve_format: bool = True
    synthetic_pool: Optional[str] = None
    custom_regex: Optional[str] = None


class TableShieldConfig(BaseModel):
    """Per-table shield configuration with column rules."""
    database: str
    schema_name: str
    table_name: str
    table_type: str = "unknown"
    column_rules: List[ColumnRule] = []
    key_columns: List[str] = []
    skip_columns: List[str] = []
    schema_snapshot: Optional[SchemaSnapshot] = None
    change_history: List[SchemaChangeEvent] = []


class ShieldProject(BaseModel):
    """Top-level DataShield project."""
    id: str = Field(default_factory=lambda: f"proj_{uuid.uuid4().hex[:6]}")
    name: str
    description: Optional[str] = None
    key_alias: str
    tables: List[TableShieldConfig] = []
    created_at: str = Field(default_factory=lambda: datetime.now().isoformat())
    updated_at: Optional[str] = None
    active: bool = True


class DataAttestation(BaseModel):
    """Signed attestation proving trust-lane and masking policy for a dataset."""
    attestation_id: str = Field(default_factory=lambda: f"att_{uuid.uuid4().hex[:10]}")
    dataset_id: str
    lane: TrustLane
    masked: bool
    policy_version: str = "v1"
    input_hash: str
    signature: str
    public_key_id: str
    issued_at: str = Field(default_factory=lambda: datetime.now().isoformat())
    expires_at: Optional[str] = None
