"""
Schema Engine - Enterprise-grade schema management.

Features:
- Schema versioning (semantic)
- Schema diff
- Constraint validation
- Column lineage tracking
"""

from typing import List, Dict, Optional, Any, Set
from dataclasses import dataclass, field
from datetime import datetime
import pandas as pd
import json


@dataclass
class ColumnConstraint:
    """Constraints for a column."""
    nullable: bool = True
    unique: bool = False
    min_value: Optional[float] = None
    max_value: Optional[float] = None
    allowed_values: Optional[List[Any]] = None
    regex_pattern: Optional[str] = None


@dataclass
class ColumnSchema:
    """Schema for a single column."""
    name: str
    dtype: str
    nullable: bool = True
    constraints: Optional[ColumnConstraint] = None
    lineage: List[str] = field(default_factory=list)  # Source columns
    description: str = ""


@dataclass
class DataSchema:
    """Complete dataset schema."""
    name: str
    version: str  # Semantic version (e.g., "1.2.0")
    columns: Dict[str, ColumnSchema]
    created_at: datetime = field(default_factory=datetime.now)
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> Dict:
        """Convert to dictionary."""
        return {
            'name': self.name,
            'version': self.version,
            'created_at': self.created_at.isoformat(),
            'columns': {
                name: {
                    'dtype': col.dtype,
                    'nullable': col.nullable,
                    'lineage': col.lineage,
                    'description': col.description,
                    'constraints': {
                        'nullable': col.constraints.nullable if col.constraints else True,
                        'unique': col.constraints.unique if col.constraints else False,
                        'min_value': col.constraints.min_value if col.constraints else None,
                        'max_value': col.constraints.max_value if col.constraints else None,
                        'allowed_values': col.constraints.allowed_values if col.constraints else None
                    } if col.constraints else {}
                }
                for name, col in self.columns.items()
            },
            'metadata': self.metadata
        }
    
    def to_json(self, indent: int = 2) -> str:
        """Export as JSON."""
        return json.dumps(self.to_dict(), indent=indent)


@dataclass
class SchemaDiff:
    """Differences between two schemas."""
    columns_added: List[str]
    columns_removed: List[str]
    columns_modified: Dict[str, Dict[str, Any]]  # column -> {field: (old, new)}
    version_change: tuple  # (old_version, new_version)


class SchemaEngine:
    """Enterprise schema management."""
    
    def infer_schema(
        self,
        df: pd.DataFrame,
        name: str = "dataset",
        detect_constraints: bool = True
    ) -> DataSchema:
        """
        Infer schema from DataFrame.
        
        Args:
            df: DataFrame to analyze
            name: Schema name
            detect_constraints: Detect constraints (min/max, allowed values)
        """
        columns = {}
        
        for col_name in df.columns:
            col = df[col_name]
            
            # Basic info
            dtype = str(col.dtype)
            nullable = col.isna().any()
            
            # Detect constraints
            constraints = None
            if detect_constraints:
                constraints = ColumnConstraint(
                    nullable=nullable,
                    unique=(col.nunique() == len(col)),
                )
                
                # Numeric constraints
                if pd.api.types.is_numeric_dtype(col):
                    constraints.min_value = float(col.min())
                    constraints.max_value = float(col.max())
                
                # Categorical constraints
                if col.nunique() < 50 and col.nunique() / len(col) < 0.1:
                    constraints.allowed_values = col.dropna().unique().tolist()
            
            columns[col_name] = ColumnSchema(
                name=col_name,
                dtype=dtype,
                nullable=nullable,
                constraints=constraints
            )
        
        return DataSchema(
            name=name,
            version="1.0.0",
            columns=columns
        )
    
    def diff(
        self,
        schema1: DataSchema,
        schema2: DataSchema
    ) -> SchemaDiff:
        """
        Compare two schemas.
        
        Args:
            schema1: First schema (baseline)
            schema2: Second schema (new)
        """
        cols1 = set(schema1.columns.keys())
        cols2 = set(schema2.columns.keys())
        
        added = list(cols2 - cols1)
        removed = list(cols1 - cols2)
        
        # Check modified columns
        modified = {}
        for col in cols1 & cols2:
            col1 = schema1.columns[col]
            col2 = schema2.columns[col]
            
            changes = {}
            if col1.dtype != col2.dtype:
                changes['dtype'] = (col1.dtype, col2.dtype)
            if col1.nullable != col2.nullable:
                changes['nullable'] = (col1.nullable, col2.nullable)
            
            if changes:
                modified[col] = changes
        
        return SchemaDiff(
            columns_added=added,
            columns_removed=removed,
            columns_modified=modified,
            version_change=(schema1.version, schema2.version)
        )
    
    def validate(
        self,
        df: pd.DataFrame,
        schema: DataSchema,
        strict: bool = False
    ) -> List[str]:
        """
        Validate DataFrame against schema.
        
        Args:
            df: DataFrame to validate
            schema: Schema to validate against
            strict: Strict mode (fail on any violation)
        
        Returns:
            List of validation errors
        """
        errors = []
        
        # Check columns
        df_cols = set(df.columns)
        schema_cols = set(schema.columns.keys())
        
        if strict and df_cols != schema_cols:
            extra = df_cols - schema_cols
            missing = schema_cols - df_cols
            
            if extra:
                errors.append(f"Extra columns: {extra}")
            if missing:
                errors.append(f"Missing columns: {missing}")
        
        # Validate each column
        for col_name, col_schema in schema.columns.items():
            if col_name not in df.columns:
                continue
            
            col = df[col_name]
            
            # Type check
            if str(col.dtype) != col_schema.dtype:
                errors.append(f"Column '{col_name}' type mismatch: {col.dtype} != {col_schema.dtype}")
            
            # Nullable check
            if not col_schema.nullable and col.isna().any():
                errors.append(f"Column '{col_name}' contains nulls but should not")
            
            # Constraint checks
            if col_schema.constraints:
                c = col_schema.constraints
                
                # Min/max
                if c.min_value is not None and col.min() < c.min_value:
                    errors.append(f"Column '{col_name}' min value {col.min()} < {c.min_value}")
                if c.max_value is not None and col.max() > c.max_value:
                    errors.append(f"Column '{col_name}' max value {col.max()} > {c.max_value}")
                
                # Allowed values
                if c.allowed_values:
                    invalid = set(col.dropna().unique()) - set(c.allowed_values)
                    if invalid:
                        errors.append(f"Column '{col_name}' has invalid values: {invalid}")
        
        return errors
    
    def version_bump(
        self,
        schema: DataSchema,
        diff: SchemaDiff,
        bump_type: str = 'auto'
    ) -> str:
        """
        Calculate next version based on diff.
        
        Args:
            schema: Current schema
            diff: Schema diff
            bump_type: 'major', 'minor', 'patch', or 'auto'
        
        Returns:
            Next version string
        """
        major, minor, patch = map(int, schema.version.split('.'))
        
        if bump_type == 'auto':
            # Breaking changes -> major
            if diff.columns_removed or any('dtype' in changes for changes in diff.columns_modified.values()):
                major += 1
                minor = 0
                patch = 0
            # New columns -> minor
            elif diff.columns_added:
                minor += 1
                patch = 0
            # Other changes -> patch
            else:
                patch += 1
        elif bump_type == 'major':
            major += 1
            minor = 0
            patch = 0
        elif bump_type == 'minor':
            minor += 1
            patch = 0
        else:  # patch
            patch += 1
        
        return f"{major}.{minor}.{patch}"
