"""
ML Data Contracts - enforceable agreements between data producers and consumers.
"""

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

import pandas as pd

from .schema import Schema, ColumnSchema


class ContractLevel(Enum):
    """Severity levels for contract violations."""

    WARNING = "warning"
    ERROR = "error"
    CRITICAL = "critical"


@dataclass
class ContractViolation:
    """Represents a single contract violation."""

    rule_name: str
    column: Optional[str]
    message: str
    level: ContractLevel
    details: Dict[str, Any] = field(default_factory=dict)


class DataContract:
    """
    Enforceable data contract for ML pipelines.
    """

    def __init__(
        self,
        name: str,
        schema: Schema,
        description: Optional[str] = None,
        version: str = "1.0.0",
    ):
        self.name = name
        self.schema = schema
        self.description = description or f"Data contract for {name}"
        self.version = version
        self.rules: List[Dict] = []
        self._custom_validators: List[Callable] = []

    def add_quality_rule(
        self,
        column: Optional[str],
        check_type: str,
        threshold: float,
        level: ContractLevel = ContractLevel.ERROR,
    ):
        """
        Add data quality rule.
        """
        self.rules.append(
            {
                "column": column,
                "check_type": check_type,
                "threshold": threshold,
                "level": level,
            }
        )
        return self

    def add_distribution_rule(
        self,
        column: str,
        reference_distribution: Dict[str, float],
        max_ks_distance: float = 0.1,
        level: ContractLevel = ContractLevel.ERROR,
    ):
        """
        Add statistical distribution constraint.
        """
        self.rules.append(
            {
                "column": column,
                "check_type": "distribution",
                "reference": reference_distribution,
                "max_ks_distance": max_ks_distance,
                "level": level,
            }
        )
        return self

    def add_custom_validator(self, validator: Callable[[pd.DataFrame], List[ContractViolation]]):
        """Add custom validation function."""
        self._custom_validators.append(validator)
        return self

    def validate(self, df: pd.DataFrame, context: Optional[Dict] = None) -> List[ContractViolation]:
        """
        Validate DataFrame against contract.
        """
        violations = []
        context = context or {}

        # 1. Schema validation
        try:
            schema_violations = self.schema.validate(df, strict=False, raise_error=False)
            for col, errors in schema_violations.items():
                for error in errors:
                    violations.append(
                        ContractViolation(
                            rule_name="schema",
                            column=col if not col.startswith("__") else None,
                            message=error,
                            level=ContractLevel.ERROR,
                        )
                    )
        except Exception as e:
            violations.append(
                ContractViolation(
                    rule_name="schema",
                    column=None,
                    message=f"Schema validation error: {str(e)}",
                    level=ContractLevel.CRITICAL,
                )
            )

        # 2. Quality rules
        for rule in self.rules:
            if rule["check_type"] == "completeness":
                viol = self._check_completeness(df, rule)
                if viol:
                    violations.append(viol)
            elif rule["check_type"] == "uniqueness":
                viol = self._check_uniqueness(df, rule)
                if viol:
                    violations.append(viol)
            elif rule["check_type"] == "distribution":
                viol = self._check_distribution(df, rule)
                if viol:
                    violations.append(viol)
            elif rule["check_type"] == "freshness":
                viol = self._check_freshness(df, rule, context)
                if viol:
                    violations.append(viol)

        # 3. Custom validators
        for validator in self._custom_validators:
            try:
                custom_violations = validator(df)
                violations.extend(custom_violations)
            except Exception as e:
                violations.append(
                    ContractViolation(
                        rule_name="custom",
                        column=None,
                        message=f"Custom validator error: {str(e)}",
                        level=ContractLevel.ERROR,
                    )
                )

        return violations

    def _check_completeness(self, df: pd.DataFrame, rule: Dict) -> Optional[ContractViolation]:
        """Check data completeness (non-null ratio)."""
        col = rule["column"]
        threshold = rule["threshold"]

        if col:
            completeness = 1 - df[col].isnull().mean()
        else:
            completeness = 1 - df.isnull().mean().mean()

        if completeness < threshold:
            return ContractViolation(
                rule_name="completeness",
                column=col,
                message=(f"Completeness {completeness:.2%} below " f"threshold {threshold:.2%}"),
                level=rule["level"],
                details={"actual": completeness, "threshold": threshold},
            )
        return None

    def _check_uniqueness(self, df: pd.DataFrame, rule: Dict) -> Optional[ContractViolation]:
        """Check uniqueness constraint."""
        col = rule["column"]
        threshold = rule["threshold"]

        if not col:
            return None

        uniqueness = df[col].nunique() / len(df[col].dropna())

        if uniqueness < threshold:
            return ContractViolation(
                rule_name="uniqueness",
                column=col,
                message=(f"Uniqueness {uniqueness:.2%} below " f"threshold {threshold:.2%}"),
                level=rule["level"],
                details={"actual": uniqueness, "threshold": threshold},
            )
        return None

    def _check_distribution(self, df: pd.DataFrame, rule: Dict) -> Optional[ContractViolation]:
        """Check distribution similarity using KS test or PSI."""
        from scipy import stats

        col = rule["column"]
        reference = rule["reference"]
        max_distance = rule["max_ks_distance"]

        if col not in df.columns:
            return ContractViolation(
                rule_name="distribution",
                column=col,
                message=f"Column {col} not found for distribution check",
                level=rule["level"],
            )

        current = df[col].value_counts(normalize=True).to_dict()

        all_categories = set(reference.keys()) | set(current.keys())
        ref_values = [reference.get(cat, 0) for cat in all_categories]
        curr_values = [current.get(cat, 0) for cat in all_categories]

        ks_stat, p_value = stats.ks_2samp(ref_values, curr_values)

        if ks_stat > max_distance:
            return ContractViolation(
                rule_name="distribution",
                column=col,
                message=f"Distribution changed (KS={ks_stat:.3f} > {max_distance})",
                level=rule["level"],
                details={"ks_statistic": ks_stat, "p_value": p_value},
            )
        return None

    def _check_freshness(
        self, df: pd.DataFrame, rule: Dict, context: Dict
    ) -> Optional[ContractViolation]:
        """Check data freshness."""
        col = rule["column"]
        threshold = rule["threshold"]

        current_time = context.get("current_time", pd.Timestamp.now())

        if col in df.columns and pd.api.types.is_datetime64_any_dtype(df[col]):
            max_time = df[col].max()
            age = (current_time - max_time).total_seconds() / 3600

            if age > threshold:
                return ContractViolation(
                    rule_name="freshness",
                    column=col,
                    message=f"Data is {age:.1f} hours old (max {threshold})",
                    level=rule["level"],
                    details={"age_hours": age, "threshold_hours": threshold},
                )
        return None

    def is_valid(self, df: pd.DataFrame, max_level: ContractLevel = ContractLevel.ERROR) -> bool:
        """Quick check if data meets contract."""
        violations = self.validate(df)
        critical_levels = {ContractLevel.ERROR, ContractLevel.CRITICAL}
        if max_level == ContractLevel.WARNING:
            critical_levels.add(ContractLevel.WARNING)

        return not any(v.level in critical_levels for v in violations)

    def to_dict(self) -> Dict:
        """Serialize contract to dictionary."""
        return {
            "name": self.name,
            "version": self.version,
            "description": self.description,
            "schema": json.loads(self.schema.to_json()),
            "rules": self.rules,
        }

    @classmethod
    def from_dict(cls, data: Dict) -> "DataContract":
        """Deserialize contract from dictionary."""
        schema = Schema.from_json(json.dumps(data["schema"]))
        contract = cls(
            name=data["name"],
            schema=schema,
            description=data.get("description"),
            version=data.get("version", "1.0.0"),
        )
        contract.rules = data.get("rules", [])
        return contract


# from typing import Dict, List, Optional, Callable, Any, Union
# from dataclasses import dataclass, field
# from enum import Enum
# import json
# import pandas as pd
# import numpy as np

# from .schema import Schema, ColumnType, SchemaValidationError


# class ContractLevel(Enum):
#     """Severity levels for contract violations."""

#     WARNING = "warning"
#     ERROR = "error"
#     CRITICAL = "critical"


# @dataclass
# class ContractViolation:
#     """Represents a single contract violation."""

#     rule_name: str
#     column: Optional[str]
#     message: str
#     level: ContractLevel
#     details: Dict[str, Any] = field(default_factory=dict)


# class DataContract:
#     """
#     Enforceable data contract for ML pipelines.

#     Defines expectations about data quality, schema, and distributions
#     that must be met before data enters the pipeline.
#     """

#     def __init__(
#         self,
#         name: str,
#         schema: Schema,
#         description: Optional[str] = None,
#         version: str = "1.0.0",
#     ):
#         self.name = name
#         self.schema = schema
#         self.description = description or f"Data contract for {name}"
#         self.version = version
#         self.rules: List[Dict] = []
#         self._custom_validators: List[Callable] = []

#     def add_quality_rule(
#         self,
#         column: Optional[str],
#         check_type: str,
#         threshold: float,
#         level: ContractLevel = ContractLevel.ERROR,
#     ):
#         """
#         Add data quality rule.

#         Args:
#             column: Column name (None for dataset-wide)
#             check_type: 'completeness', 'uniqueness', 'consistency', 'freshness'
#             threshold: Minimum acceptable value (0.0-1.0)
#             level: Severity level
#         """
#         self.rules.append(
#             {
#                 "column": column,
#                 "check_type": check_type,
#                 "threshold": threshold,
#                 "level": level,
#             }
#         )
#         return self

#     def add_distribution_rule(
#         self,
#         column: str,
#         reference_distribution: Dict[str, float],
#         max_ks_distance: float = 0.1,
#         level: ContractLevel = ContractLevel.ERROR,
#     ):
#         """
#         Add statistical distribution constraint.

#         Args:
#             column: Column to check
#             reference_distribution: Reference value counts (normalized)
#             max_ks_distance: Maximum allowed Kolmogorov-Smirnov distance
#             level: Severity level
#         """
#         self.rules.append(
#             {
#                 "column": column,
#                 "check_type": "distribution",
#                 "reference": reference_distribution,
#                 "max_ks_distance": max_ks_distance,
#                 "level": level,
#             }
#         )
#         return self

#     def add_custom_validator(
#         self, validator: Callable[[pd.DataFrame], List[ContractViolation]]
#     ):
#         """Add custom validation function."""
#         self._custom_validators.append(validator)
#         return self

#     def validate(
#         self, df: pd.DataFrame, context: Optional[Dict] = None
#     ) -> List[ContractViolation]:
#         """
#         Validate DataFrame against contract.

#         Returns:
#             List of ContractViolation objects (empty if all pass)
#         """
#         violations = []
#         context = context or {}

#         # 1. Schema validation
#         try:
#             schema_violations = self.schema.validate(
#                 df, strict=False, raise_error=False
#             )
#             for col, errors in schema_violations.items():
#                 for error in errors:
#                     violations.append(
#                         ContractViolation(
#                             rule_name="schema",
#                             column=col if not col.startswith("__") else None,
#                             message=error,
#                             level=ContractLevel.ERROR,
#                         )
#                     )
#         except Exception as e:
#             violations.append(
#                 ContractViolation(
#                     rule_name="schema",
#                     column=None,
#                     message=f"Schema validation error: {str(e)}",
#                     level=ContractLevel.CRITICAL,
#                 )
#             )

#         # 2. Quality rules
#         for rule in self.rules:
#             if rule["check_type"] == "completeness":
#                 viol = self._check_completeness(df, rule)
#                 if viol:
#                     violations.append(viol)
#             elif rule["check_type"] == "uniqueness":
#                 viol = self._check_uniqueness(df, rule)
#                 if viol:
#                     violations.append(viol)
#             elif rule["check_type"] == "distribution":
#                 viol = self._check_distribution(df, rule)
#                 if viol:
#                     violations.append(viol)
#             elif rule["check_type"] == "freshness":
#                 viol = self._check_freshness(df, rule, context)
#                 if viol:
#                     violations.append(viol)

#         # 3. Custom validators
#         for validator in self._custom_validators:
#             try:
#                 custom_violations = validator(df)
#                 violations.extend(custom_violations)
#             except Exception as e:
#                 violations.append(
#                     ContractViolation(
#                         rule_name="custom",
#                         column=None,
#                         message=f"Custom validator error: {str(e)}",
#                         level=ContractLevel.ERROR,
#                     )
#                 )

#         return violations

#     def _check_completeness(
#         self, df: pd.DataFrame, rule: Dict
#     ) -> Optional[ContractViolation]:
#         """Check data completeness (non-null ratio)."""
#         col = rule["column"]
#         threshold = rule["threshold"]

#         if col:
#             completeness = 1 - df[col].isnull().mean()
#         else:
#             completeness = 1 - df.isnull().mean().mean()

#         if completeness < threshold:
#             return ContractViolation(
#                 rule_name="completeness",
#                 column=col,
#                 message=f"Completeness {completeness:.2%} below threshold {threshold:.2%}",
#                 level=rule["level"],
#                 details={"actual": completeness, "threshold": threshold},
#             )
#         return None

#     def _check_uniqueness(
#         self, df: pd.DataFrame, rule: Dict
#     ) -> Optional[ContractViolation]:
#         """Check uniqueness constraint."""
#         col = rule["column"]
#         threshold = rule["threshold"]

#         if not col:
#             return None

#         uniqueness = df[col].nunique() / len(df[col].dropna())

#         if uniqueness < threshold:
#             return ContractViolation(
#                 rule_name="uniqueness",
#                 column=col,
#                 message=f"Uniqueness {uniqueness:.2%} below threshold {threshold:.2%}",
#                 level=rule["level"],
#                 details={"actual": uniqueness, "threshold": threshold},
#             )
#         return None

#     def _check_distribution(
#         self, df: pd.DataFrame, rule: Dict
#     ) -> Optional[ContractViolation]:
#         """Check distribution similarity using KS test or PSI."""
#         from scipy import stats

#         col = rule["column"]
#         reference = rule["reference"]
#         max_distance = rule["max_ks_distance"]

#         if col not in df.columns:
#             return ContractViolation(
#                 rule_name="distribution",
#                 column=col,
#                 message=f"Column {col} not found for distribution check",
#                 level=rule["level"],
#             )

#         # Calculate current distribution
#         current = df[col].value_counts(normalize=True).to_dict()

#         # Align distributions
#         all_categories = set(reference.keys()) | set(current.keys())
#         ref_values = [reference.get(cat, 0) for cat in all_categories]
#         curr_values = [current.get(cat, 0) for cat in all_categories]

#         # KS test
#         ks_stat, p_value = stats.ks_2samp(ref_values, curr_values)

#         if ks_stat > max_distance:
#             return ContractViolation(
#                 rule_name="distribution",
#                 column=col,
#                 message=f"Distribution changed (KS={ks_stat:.3f} > {max_distance})",
#                 level=rule["level"],
#                 details={"ks_statistic": ks_stat, "p_value": p_value},
#             )
#         return None

#     def _check_freshness(
#         self, df: pd.DataFrame, rule: Dict, context: Dict
#     ) -> Optional[ContractViolation]:
#         """Check data freshness."""
#         col = rule["column"]
#         threshold = rule["threshold"]

#         # Expect context to contain 'current_time' or use now
#         current_time = context.get("current_time", pd.Timestamp.now())

#         if col in df.columns and pd.api.types.is_datetime64_any_dtype(df[col]):
#             max_time = df[col].max()
#             age = (current_time - max_time).total_seconds() / 3600  # hours

#             # threshold interpreted as max age in hours
#             if age > threshold:
#                 return ContractViolation(
#                     rule_name="freshness",
#                     column=col,
#                     message=f"Data is {age:.1f} hours old (max {threshold})",
#                     level=rule["level"],
#                     details={"age_hours": age, "threshold_hours": threshold},
#                 )
#         return None

#     def is_valid(
#         self, df: pd.DataFrame, max_level: ContractLevel = ContractLevel.ERROR
#     ) -> bool:
#         """Quick check if data meets contract (no critical or error violations)."""
#         violations = self.validate(df)
#         critical_levels = {ContractLevel.ERROR, ContractLevel.CRITICAL}
#         if max_level == ContractLevel.WARNING:
#             critical_levels.add(ContractLevel.WARNING)

#         return not any(v.level in critical_levels for v in violations)

#     def to_dict(self) -> Dict:
#         """Serialize contract to dictionary."""
#         return {
#             "name": self.name,
#             "version": self.version,
#             "description": self.description,
#             "schema": json.loads(self.schema.to_json()),
#             "rules": self.rules,
#         }

#     @classmethod
#     def from_dict(cls, data: Dict) -> "DataContract":
#         """Deserialize contract from dictionary."""
#         schema = Schema.from_json(json.dumps(data["schema"]))
#         contract = cls(
#             name=data["name"],
#             schema=schema,
#             description=data.get("description"),
#             version=data.get("version", "1.0.0"),
#         )
#         contract.rules = data.get("rules", [])
#         return contract
