"""
Smart schema inference and validation system with ML-aware type detection.
"""

import re
import json
from enum import Enum
from typing import Dict, List, Optional, Any, Callable
from dataclasses import dataclass, field
from collections import defaultdict

import pandas as pd
import numpy as np

try:
    import yaml
except ImportError:
    yaml = None  # type: ignore


class ColumnType(Enum):
    """ML-aware column type enumeration."""

    INTEGER = "integer"
    FLOAT = "float"
    BOOLEAN = "boolean"
    STRING = "string"
    CATEGORICAL = "categorical"
    DATETIME = "datetime"
    TIMEDELTA = "timedelta"
    COMPLEX = "complex"
    MISSING = "missing"
    MIXED = "mixed"
    IDENTIFIER = "identifier"
    TARGET = "target"
    FEATURE = "feature"
    TEXT = "text"


@dataclass
class ColumnSchema:
    """Detailed schema for a single column."""

    name: str
    dtype: ColumnType
    python_type: str
    nullable: bool = True
    unique: bool = False
    unique_ratio: float = 0.0
    null_count: int = 0
    null_ratio: float = 0.0
    min_value: Optional[Any] = None
    max_value: Optional[Any] = None
    mean: Optional[float] = None
    std: Optional[float] = None
    cardinality: int = 0
    cardinality_ratio: float = 0.0
    sample_values: List[Any] = field(default_factory=list)
    regex_pattern: Optional[str] = None
    semantic_type: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary with proper type handling."""
        result = {
            "name": self.name,
            "dtype": self.dtype.value,
            "python_type": self.python_type,
            "nullable": bool(self.nullable),
            "unique": bool(self.unique),
            "unique_ratio": float(self.unique_ratio),
            "null_count": int(self.null_count),
            "null_ratio": float(self.null_ratio),
            "cardinality": int(self.cardinality),
            "cardinality_ratio": float(self.cardinality_ratio),
            "sample_values": self._clean_values(self.sample_values),
            "regex_pattern": self.regex_pattern,
            "semantic_type": self.semantic_type,
        }

        if self.min_value is not None:
            result["min_value"] = float(self.min_value)
        if self.max_value is not None:
            result["max_value"] = float(self.max_value)
        if self.mean is not None:
            result["mean"] = float(self.mean)
        if self.std is not None:
            result["std"] = float(self.std)

        return result

    def _clean_values(self, values: List[Any]) -> List[Any]:
        """Clean sample values for serialization."""
        cleaned = []
        for v in values:
            if isinstance(v, (np.integer, np.floating)):
                cleaned.append(float(v) if isinstance(v, np.floating) else int(v))
            elif isinstance(v, np.bool_):
                cleaned.append(bool(v))
            elif pd.isna(v):
                cleaned.append(None)  # type: ignore
            else:
                cleaned.append(v)
        return cleaned

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ColumnSchema":
        """Create from dictionary."""
        data["dtype"] = ColumnType(data["dtype"])
        return cls(**data)


class Schema:
    """
    Intelligent schema for ML datasets with validation capabilities.
    """

    def __init__(self, columns: Dict[str, ColumnSchema], metadata: Optional[Dict] = None):
        self.columns = columns
        self.metadata = metadata or {}
        self._validation_rules: Dict[str, List[Callable]] = defaultdict(list)

    def __repr__(self) -> str:
        return (
            f"Schema(columns={len(self.columns)}, "
            f"features={self.feature_names}, target={self.target_name})"
        )

    @property
    def feature_names(self) -> List[str]:
        """Return likely feature column names."""
        return [
            name
            for name, col in self.columns.items()
            if col.dtype not in [ColumnType.TARGET, ColumnType.IDENTIFIER, ColumnType.MISSING]
        ]

    @property
    def target_name(self) -> Optional[str]:
        """Infer target column name."""
        targets = [name for name, col in self.columns.items() if col.dtype == ColumnType.TARGET]
        return targets[0] if targets else None

    @property
    def numeric_columns(self) -> List[str]:
        """Return numeric column names."""
        numeric_types = {ColumnType.INTEGER, ColumnType.FLOAT, ColumnType.COMPLEX}
        return [name for name, col in self.columns.items() if col.dtype in numeric_types]

    @property
    def categorical_columns(self) -> List[str]:
        """Return categorical column names."""
        cat_types = {ColumnType.CATEGORICAL, ColumnType.BOOLEAN}
        return [name for name, col in self.columns.items() if col.dtype in cat_types]

    def validate(
        self, df: pd.DataFrame, strict: bool = False, raise_error: bool = True
    ) -> Dict[str, List[str]]:
        """
        Validate a DataFrame against this schema.
        """
        violations = {}

        # Check for missing columns
        missing_cols = set(self.columns.keys()) - set(df.columns)
        if missing_cols:
            violations["__missing__"] = [f"Missing required columns: {missing_cols}"]

        # Check for extra columns in strict mode
        if strict:
            extra_cols = set(df.columns) - set(self.columns.keys())
            if extra_cols:
                violations["__extra__"] = [f"Extra columns not in schema: {extra_cols}"]

        # Validate each column
        for col_name, col_schema in self.columns.items():
            if col_name not in df.columns:
                continue

            col_violations = self._validate_column(df[col_name], col_schema)
            if col_violations:
                violations[col_name] = col_violations

        # Run custom validation rules
        for col_name, rules in self._validation_rules.items():
            if col_name in df.columns:
                for rule in rules:
                    try:
                        if not rule(df[col_name]):
                            msg = f"Custom rule failed for {col_name}"
                            violations.setdefault(col_name, []).append(msg)
                    except Exception as e:
                        msg = f"Custom rule error: {str(e)}"
                        violations.setdefault(col_name, []).append(msg)

        if raise_error and violations:
            raise SchemaValidationError(violations)

        return violations

    def _validate_column(self, series: pd.Series, schema: ColumnSchema) -> List[str]:
        """Validate a single column against its schema."""
        violations = []

        # Check nullability
        if not schema.nullable and series.isnull().any():
            violations.append(f"Column '{schema.name}' contains nulls but is non-nullable")

        # Check type compatibility (skip for TARGET type)
        if schema.dtype != ColumnType.TARGET:
            actual_dtype = self._infer_pandas_type(series)
            if not self._types_compatible(actual_dtype, schema.dtype):
                msg = f"Type mismatch: expected {schema.dtype.value}, " f"got {actual_dtype.value}"
                violations.append(msg)

        # Check range constraints for numeric (only if schema has numeric bounds)
        if schema.dtype in [ColumnType.INTEGER, ColumnType.FLOAT]:
            if schema.min_value is not None or schema.max_value is not None:
                try:
                    # Convert to numeric for comparison
                    numeric_series = pd.to_numeric(series, errors="coerce").dropna()
                    if len(numeric_series) > 0:
                        current_min = float(numeric_series.min())
                        current_max = float(numeric_series.max())

                        # Use 10% tolerance
                        if schema.min_value is not None:
                            tolerance_min = (
                                float(schema.min_value) - abs(float(schema.min_value)) * 0.1
                            )
                            if current_min < tolerance_min:
                                violations.append(f"Values below minimum {schema.min_value}")

                        if schema.max_value is not None:
                            tolerance_max = (
                                float(schema.max_value) + abs(float(schema.max_value)) * 0.1
                            )
                            if current_max > tolerance_max:
                                violations.append(f"Values above maximum {schema.max_value}")
                except (ValueError, TypeError):
                    # Skip range validation if conversion fails
                    pass

        # Check cardinality for categorical
        if schema.dtype == ColumnType.CATEGORICAL:
            actual_cardinality = series.nunique()
            if actual_cardinality > schema.cardinality * 1.5:
                msg = f"Cardinality increased: {schema.cardinality} -> " f"{actual_cardinality}"
                violations.append(msg)

        # Check regex pattern if defined
        if schema.regex_pattern and schema.dtype == ColumnType.STRING:
            try:
                pattern = re.compile(schema.regex_pattern)
                non_matching = series.dropna().astype(str).apply(lambda x: not pattern.match(x))
                if non_matching.any():
                    msg = (
                        f"{non_matching.sum()} values don't match pattern "
                        f"{schema.regex_pattern}"
                    )
                    violations.append(msg)
            except re.error:
                # Invalid regex, skip
                pass

        return violations

    def _infer_pandas_type(self, series: pd.Series) -> ColumnType:
        """Infer ColumnType from pandas Series."""
        return SchemaBuilder._infer_column_type(series)

    def _types_compatible(self, actual: ColumnType, expected: ColumnType) -> bool:
        """Check if two types are compatible."""
        if actual == expected:
            return True
        if actual == ColumnType.INTEGER and expected == ColumnType.FLOAT:
            return True
        if actual == ColumnType.STRING and expected == ColumnType.CATEGORICAL:
            return True
        if expected == ColumnType.TARGET:
            return True
        return False

    def add_validation_rule(self, column: str, rule: Callable[[pd.Series], bool]):
        """Add custom validation rule for a column."""
        self._validation_rules[column].append(rule)

    def adapt(self, df: pd.DataFrame, inplace: bool = False) -> "Schema":
        """
        Adapt schema to new data (evolution).
        """
        new_schema = SchemaBuilder().infer_schema(df)

        if inplace:
            for col_name, new_col in new_schema.columns.items():
                if col_name in self.columns:
                    old_col = self.columns[col_name]
                    if new_col.min_value is not None and old_col.min_value is not None:
                        new_col.min_value = min(old_col.min_value, new_col.min_value)
                    if new_col.max_value is not None and old_col.max_value is not None:
                        new_col.max_value = max(old_col.max_value, new_col.max_value)
                    new_col.nullable = old_col.nullable or new_col.nullable
            self.columns.update(new_schema.columns)
            return self
        else:
            return new_schema

    def to_json(self, path: Optional[str] = None) -> str:
        """Serialize to JSON."""
        data = {
            "metadata": self.metadata,
            "columns": {name: col.to_dict() for name, col in self.columns.items()},
        }
        json_str = json.dumps(data, indent=2, default=str)
        if path:
            with open(path, "w") as f:
                f.write(json_str)
        return json_str

    def to_yaml(self, path: Optional[str] = None) -> str:
        """Serialize to YAML."""
        if yaml is None:
            raise ImportError("PyYAML is required for YAML serialization")
        data = {
            "metadata": self.metadata,
            "columns": {name: col.to_dict() for name, col in self.columns.items()},
        }
        yaml_str = yaml.safe_dump(data, default_flow_style=False, sort_keys=False)
        if path:
            with open(path, "w") as f:
                f.write(yaml_str)
        return yaml_str

    @classmethod
    def from_json(cls, json_str: str) -> "Schema":
        """Deserialize from JSON string."""
        data = json.loads(json_str)
        columns = {
            name: ColumnSchema.from_dict(col_data) for name, col_data in data["columns"].items()
        }
        return cls(columns, data.get("metadata", {}))

    @classmethod
    def from_yaml(cls, yaml_str: str) -> "Schema":
        """Deserialize from YAML string."""
        if yaml is None:
            raise ImportError("PyYAML is required for YAML deserialization")
        data = yaml.safe_load(yaml_str)
        columns = {
            name: ColumnSchema.from_dict(col_data) for name, col_data in data["columns"].items()
        }
        return cls(columns, data.get("metadata", {}))

    def generate_pandera_schema(self):
        """Generate Pandera schema (if pandera is installed)."""
        try:
            import pandera as pa

            columns = {}
            for name, col in self.columns.items():
                pa_type = {
                    ColumnType.INTEGER: pa.Int,
                    ColumnType.FLOAT: pa.Float,
                    ColumnType.STRING: pa.String,
                    ColumnType.BOOLEAN: pa.Bool,
                }.get(col.dtype, pa.Object)

                checks = []
                if col.min_value is not None:
                    checks.append(pa.Check.greater_than_or_equal_to(col.min_value))
                if col.max_value is not None:
                    checks.append(pa.Check.less_than_or_equal_to(col.max_value))

                columns[name] = pa.Column(pa_type, nullable=col.nullable, checks=checks)
            return pa.DataFrameSchema(columns)
        except ImportError:
            raise ImportError("pandera is required for this feature")


class SchemaBuilder:
    """Builds Schema objects from data with intelligent inference."""

    PATTERNS = {
        "email": r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$",
        "url": r"^https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+",
        "uuid": r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$",
        "phone": r"^[\d\s\-\+\(\)]{10,20}$",
        "credit_card": r"^\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}$",
        "ip_address": r"^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$",
    }

    def __init__(
        self,
        categorical_threshold: int = 50,
        identifier_threshold: float = 0.95,
        text_threshold: int = 100,
    ):
        self.categorical_threshold = categorical_threshold
        self.identifier_threshold = identifier_threshold
        self.text_threshold = text_threshold

    def infer_schema(self, df: pd.DataFrame, target_col: Optional[str] = None) -> Schema:
        """
        Infer comprehensive schema from DataFrame.
        """
        columns = {}
        target_detected = False

        for col_name in df.columns:
            series = df[col_name]
            col_schema = self._analyze_column(col_name, series)

            if col_name == target_col:
                col_schema.dtype = ColumnType.TARGET
                target_detected = True
            elif not target_detected and self._is_likely_target(series, col_name):
                col_schema.dtype = ColumnType.TARGET
                target_detected = True

            columns[col_name] = col_schema

        metadata = {
            "inferred_at": pd.Timestamp.now().isoformat(),
            "row_count": len(df),
            "column_count": len(df.columns),
        }

        return Schema(columns, metadata)

    def _analyze_column(self, name: str, series: pd.Series) -> ColumnSchema:
        """Deep analysis of a single column."""
        dtype = self._infer_column_type(series)

        null_count = series.isnull().sum()
        null_ratio = null_count / len(series)
        unique_count = series.nunique(dropna=True)
        unique_ratio = unique_count / len(series.dropna()) if len(series.dropna()) > 0 else 0

        sample_values = series.dropna().head(5).tolist()
        semantic_type = self._detect_semantic_type(series, dtype)

        schema = ColumnSchema(
            name=name,
            dtype=dtype,
            python_type=str(series.dtype),
            nullable=bool(null_count > 0),
            unique=unique_count == len(series.dropna()),
            unique_ratio=float(unique_ratio),
            null_count=int(null_count),
            null_ratio=float(null_ratio),
            cardinality=int(unique_count),
            cardinality_ratio=float(unique_count / len(series)),
            sample_values=sample_values,
            semantic_type=semantic_type,
        )

        if dtype in [ColumnType.INTEGER, ColumnType.FLOAT]:
            numeric_series = pd.to_numeric(series, errors="coerce").dropna()
            if len(numeric_series) > 0:
                schema.min_value = float(numeric_series.min())
                schema.max_value = float(numeric_series.max())
                schema.mean = float(numeric_series.mean())
                schema.std = float(numeric_series.std())

        if dtype == ColumnType.STRING:
            pattern = self._infer_regex_pattern(series.dropna().astype(str))
            if pattern:
                schema.regex_pattern = pattern

        return schema

    @classmethod
    def _infer_column_type(cls, series: pd.Series) -> ColumnType:
        """Infer ML-aware column type."""
        if series.isnull().all():
            return ColumnType.MISSING

        non_null = series.dropna()
        if len(non_null) == 0:
            return ColumnType.MISSING

        if pd.api.types.is_datetime64_any_dtype(series):
            return ColumnType.DATETIME

        if pd.api.types.is_timedelta64_dtype(series):
            return ColumnType.TIMEDELTA

        if pd.api.types.is_integer_dtype(series):
            return ColumnType.INTEGER

        if pd.api.types.is_float_dtype(series):
            return ColumnType.FLOAT

        if pd.api.types.is_complex_dtype(series):
            return ColumnType.COMPLEX

        if pd.api.types.is_bool_dtype(series):
            return ColumnType.BOOLEAN

        if pd.api.types.is_object_dtype(series) or pd.api.types.is_string_dtype(series):
            converted = pd.to_numeric(series, errors="coerce")
            if converted.notna().sum() / len(non_null) > 0.9:
                if (converted.dropna() % 1 == 0).all():
                    return ColumnType.INTEGER
                return ColumnType.FLOAT

            bool_values = {"true", "false", "yes", "no", "1", "0", "t", "f"}
            lower_values = set(non_null.astype(str).str.lower().unique())
            if lower_values.issubset(bool_values):
                return ColumnType.BOOLEAN

            unique_ratio = non_null.nunique() / len(non_null)
            if non_null.nunique() <= 50 or unique_ratio < 0.05:
                return ColumnType.CATEGORICAL

            avg_length = non_null.astype(str).str.len().mean()
            if avg_length > 100:
                return ColumnType.TEXT

            return ColumnType.STRING

        return ColumnType.MIXED

    def _detect_semantic_type(self, series: pd.Series, dtype: ColumnType) -> Optional[str]:
        """Detect semantic type (email, URL, etc.)."""
        if dtype not in [ColumnType.STRING, ColumnType.CATEGORICAL, ColumnType.TEXT]:
            return None

        sample = series.dropna().astype(str).head(100)

        for sem_type, pattern in self.PATTERNS.items():
            matches = sample.str.match(pattern, na=False).sum()
            if matches / len(sample) > 0.8:
                return sem_type

        return None

    def _infer_regex_pattern(self, series: pd.Series) -> Optional[str]:
        """Try to infer a regex pattern that matches most values."""
        if len(series) == 0:
            return None

        for sem_type, pattern in self.PATTERNS.items():
            matches = series.str.match(pattern, na=False).sum()
            if matches / len(series) > 0.9:
                return pattern

        return None

    def _is_likely_target(self, series: pd.Series, name: str) -> bool:
        """Heuristic to detect if column is likely a target variable."""
        name_lower = name.lower()

        strong_indicators = ["target", "label", "y"]
        if any(ind == name_lower for ind in strong_indicators):
            if series.nunique() / len(series.dropna()) < 0.95:
                return True

        if series.nunique() == 2:
            value_counts = series.value_counts(normalize=True)
            if 0.1 < value_counts.min() < 0.9:
                weak_indicators = [
                    "class",
                    "prediction",
                    "output",
                    "result",
                    "flag",
                    "churn",
                    "fraud",
                    "click",
                    "is_",
                    "has_",
                ]
                if any(ind in name_lower for ind in weak_indicators):
                    return True

        return False


class SchemaValidationError(Exception):
    """Exception raised when data validation fails."""

    def __init__(self, violations: Dict[str, List[str]]):
        self.violations = violations
        message = "Schema validation failed:\n"
        for col, errors in violations.items():
            message += f"  {col}: {', '.join(errors)}\n"
        super().__init__(message)
