from __future__ import annotations
import dask.dataframe as dd
import pandas as pd
import logging
import sqlalchemy.types as sa_types
from typing import Any, ClassVar, List, Literal, Optional, TYPE_CHECKING
from collections.abc import Mapping, MutableMapping
from pydantic import BaseModel, Field, ConfigDict, SecretStr

if TYPE_CHECKING:
    from dask.distributed import Client

from sibi_flux.df_helper import DfHelper
from sibi_flux.utils import DataUtils
from sibi_flux.dask_cluster import dask_is_empty
from sibi_flux.df_validator import DfValidator

# Define types for clarity
DataFrameType = dd.DataFrame | pd.DataFrame


class DatacubeConfig(BaseModel):
    """
    Pydantic model for strict validation of Datacube configuration.
    Allows extra fields to pass through to DfHelper/Backend.
    """

    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)

    table: Optional[str] = None
    backend: str = "sqlalchemy"
    connection_url: Optional[str | SecretStr] = None
    field_map: Mapping[str, str] = Field(default_factory=dict)
    legacy_filters: bool = True
    debug: bool = False

    # Validation Config
    validation_schema: Mapping[str, str] = Field(default_factory=dict)
    enforce_schema: bool = False

    # Allow optional logger instance
    logger: Optional[Any] = None
    dask_client: Optional[Any] = None


class Datacube(DfHelper):
    """
    Enhanced Datacube that simplifies subclassing via declarative configuration
    and automated common transformations.

    Consolidates functionality from the legacy BaseDatacube.

    Attributes:
        table (str): Table name for database connection.
        backend (str): Backend type (e.g., 'sqlalchemy').
        connection_url (str): Database connection URL.
        field_map (Mapping): Mapping of field names.
        legacy_filters (bool): Use legacy filters.
        live (bool): Live data flag.

        boolean_columns (List[str]): Columns to convert to boolean.
        numeric_float_columns (List[str]): Columns to convert to float.
        numeric_int_columns (List[str]): Columns to convert to int.
        date_fields (List[str]): Columns to convert to datetime.
    """

    # Declarative Config Attributes
    table: str = ""
    backend: Literal["sqlalchemy", "parquet", "http"] = "sqlalchemy"
    connection_url: str = ""
    field_map: Mapping[str, str] = {}
    legacy_filters: bool = True
    pii_columns: List[str] = []

    # Validation
    validation_schema: ClassVar[Mapping[str, str]] = {}
    enforce_schema: ClassVar[bool] = False

    # Declarative Transformation Lists
    boolean_columns: List[str] = []
    numeric_float_columns: List[str] = []
    numeric_int_columns: List[str] = []
    date_fields: List[str] = []

    # Config container compatible with DfHelper
    config: MutableMapping[str, Any] = {}

    def __init__(self, **kwargs):
        # State container
        self.df: Optional[DataFrameType] = None
        self._schema_inferred = False

        # Build config dictionary from class attributes if not already present
        if not self.config:
            self.config = {
                "table": self.table,
                "backend": self.backend,
                "connection_url": self.connection_url,
                "field_map": self.field_map,
                "legacy_filters": self.legacy_filters,
                "validation_schema": self.validation_schema,
                "enforce_schema": self.enforce_schema,
            }

        # Merge kwargs into the preliminary config
        raw_config = {**self.config, **kwargs}

        # Validate configuration using Pydantic
        # This ensures types are correct and provides clear checks
        validated_model = DatacubeConfig.model_validate(raw_config)

        # Update self.config with the validated, normalized dictionary
        self.config = validated_model.model_dump(exclude_unset=True)

        # Copy class-level lists to instance to avoid mutation of shared class state
        # and to allow instance-specific inference
        self.boolean_columns = list(self.boolean_columns)
        self.numeric_float_columns = list(self.numeric_float_columns)
        self.numeric_int_columns = list(self.numeric_int_columns)
        self.date_fields = list(self.date_fields)

        # Initialize DfHelper with validated config
        # We model_dump() allowing extra fields so that backend-specific params
        # (like pool_size) are preserved and passed to DfHelper.
        super().__init__(**self.config)

        # Store dask_client if provided, for resilience checks and external access
        self.dask_client = self.config.get("dask_client")

    def _infer_schema(self):
        """
        Attempts to infer column types from the SQLAlchemy model reflection.
        Populates the transformation lists if they are missing metadata.
        """
        if self.config.get("backend") != "sqlalchemy":
            return

        try:
            # This triggers reflection matching the table name
            conn = self.get_sql_connection()
            model = conn.model
            if not model:
                return

            field_map = self.config.get("field_map", {})

            # Use sets for O(1) membership checks during inference
            bool_set = set(self.boolean_columns)
            int_set = set(self.numeric_int_columns)
            float_set = set(self.numeric_float_columns)
            date_set = set(self.date_fields)

            for col in model.__table__.columns:
                # Map DB column name to DataFrame column name (if renamed via field_map)
                df_col_name = field_map.get(col.name, col.name)

                # Check types and append if not already present
                if isinstance(col.type, sa_types.Boolean):
                    if df_col_name not in bool_set:
                        self.boolean_columns.append(df_col_name)
                        bool_set.add(df_col_name)

                elif isinstance(
                    col.type,
                    (sa_types.Integer, sa_types.BigInteger, sa_types.SmallInteger),
                ):
                    if (
                        df_col_name not in int_set
                        and df_col_name not in bool_set
                        and df_col_name not in float_set
                    ):
                        self.numeric_int_columns.append(df_col_name)
                        int_set.add(df_col_name)

                elif isinstance(col.type, (sa_types.Float, sa_types.Numeric)):
                    # Numeric/Decimal often treated as float in analysis unless strict
                    if df_col_name not in float_set:
                        self.numeric_float_columns.append(df_col_name)
                        float_set.add(df_col_name)

                elif isinstance(
                    col.type, (sa_types.Date, sa_types.DateTime, sa_types.TIMESTAMP)
                ):
                    if df_col_name not in date_set:
                        self.date_fields.append(df_col_name)
                        date_set.add(df_col_name)

        except Exception as e:
            # Inference is an enhancement; failure shouldn't block app start
            # but we log it for debugging.
            self.logger.debug(f"Schema inference skipped: {e}")

    def validate_data(self, df: DataFrameType) -> DataFrameType:
        """
        Runs DfValidator if a schema is configured.
        Overrides BaseDatacube hook.
        """
        schema = self.config.get("validation_schema")
        if not schema:
            return df

        validator = DfValidator(
            df, logger=self.logger, debug=self.config.get("debug", False)
        )

        # 1. Standardize (always good practice if validating)
        # Note: We skip full standardization for now to avoid side effects on legacy names,
        # but we could adhere to stricter contracts here.

        if self.config.get("enforce_schema", False):
            validator.apply_schema_map(schema)
        else:
            try:
                validator.validate_schema(schema)
            except TypeError as e:
                # Re-raise to block pipeline
                self.logger.error(f"Schema Validation Failed: {e}")
                raise e

        return validator.get_df()

    async def avalidate_data(self, df: DataFrameType) -> DataFrameType:
        """
        Asynchronous validation hook.
        Offloads synchronous validation (CPU bound) to a thread.
        """
        import asyncio

        return await asyncio.to_thread(self.validate_data, df)

    def get_ddl(self, table_name: Optional[str] = None) -> str:
        """
        Generates ClickHouse DDL for the current cube.
        Requires that self.df is loaded (or at least inferred).
        """
        if self.df is None:
            raise RuntimeError("Cannot generate DDL: Datacube has not been loaded.")

        name = table_name or self.config.get("table") or "datacube_table"
        validator = DfValidator(self.df, logger=self.logger)
        return validator.generate_clickhouse_ddl(name)

    # -----------------------------------------------------------------------
    # Lifecycle (Consolidated from BaseDatacube)
    # -----------------------------------------------------------------------
    def load(self, **kwargs) -> DataFrameType:
        """
        Synchronous load pipeline.
        """
        # 1. Load Raw Data (Delegate to DfHelper)
        df = super().load(**kwargs)

        # 2. Check Emptiness (Resilient)
        if not self._is_empty(df):
            # 3. Apply Transform Hook
            df = self.fix_data(df, **kwargs)
            # 4. Validate
            df = self.validate_data(df)
        else:
            self.logger.debug(f"No data loaded by {self.__class__.__name__}")

        # 4. Update State & Return
        self.df = df
        return df

    async def aload(self, **kwargs) -> DataFrameType:
        """
        Asynchronous load pipeline.
        """
        # 1. Load Raw Data
        df = await super().aload(**kwargs)

        # 2. Check Emptiness (Async Safe offload)
        import asyncio

        is_empty = await asyncio.to_thread(self._is_empty, df)

        if not is_empty:
            # 3. Apply Async Transform Hook
            df = await self.afix_data(df, **kwargs)
            # 4. Validate (CPU bound)
            df = await self.avalidate_data(df)
        else:
            self.logger.debug(f"No data loaded by {self.__class__.__name__}")

        # 4. Update State & Return
        self.df = df
        return df

    # -----------------------------------------------------------------------
    # Transformation Hooks
    # -----------------------------------------------------------------------
    def fix_data(self, df: DataFrameType, **kwargs) -> DataFrameType:
        """
        Applies declarative transformations automatically.
        Infers schema types on first run if using SQLAlchemy backend.
        Subclasses can override this to add custom logic, calling super().fix_data(df) first.
        """
        if not self._schema_inferred:
            self._infer_schema()
            self._schema_inferred = True

        utils = DataUtils(logger=self.logger)

        if self.pii_columns:
            df = utils.transform_pii_columns(df, self.pii_columns)

        if self.boolean_columns:
            df = utils.transform_boolean_columns(df, self.boolean_columns)

        if self.numeric_float_columns:
            df = utils.transform_numeric_columns(
                df, columns=self.numeric_float_columns, dtype=float
            )

        if self.numeric_int_columns:
            df = utils.transform_numeric_columns(
                df, columns=self.numeric_int_columns, dtype=int
            )

        if self.date_fields:
            if isinstance(df, dd.DataFrame):
                df = utils.convert_to_datetime_dask(df, self.date_fields)
            else:
                for col in self.date_fields:
                    if col in df.columns:
                        df[col] = pd.to_datetime(
                            df[col], errors="coerce", utc=True, format="mixed"
                        )

        return df

    async def afix_data(self, df: DataFrameType, **kwargs) -> DataFrameType:
        """
        Asynchronous transformation hook.
        Offloads synchronous fix_data (and schema inference) to a thread.
        """
        import asyncio

        return await asyncio.to_thread(self.fix_data, df, **kwargs)

    # -----------------------------------------------------------------------
    # Internals
    # -----------------------------------------------------------------------
    def _is_empty(self, df: Optional[DataFrameType]) -> bool:
        """
        Robust emptiness check using dask_cluster.
        """
        if df is None:
            return True

        if isinstance(df, pd.DataFrame):
            return df.empty

        # Use resilient checker
        client = getattr(self, "dask_client", None)
        return dask_is_empty(df, dask_client=client, logger=self.logger)
