# src/sibi_flux/df_validator/_df_validator.py
from __future__ import annotations

import warnings
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
)

import dask.dataframe as dd
import pandas as pd

try:
    from dask.distributed import Client, get_client
except ImportError:
    Client = object
    get_client = None

from sibi_flux.core.type_maps import DASK_TO_CLICKHOUSE_DTYPE
from sibi_flux.logger import Logger

if TYPE_CHECKING:
    from collections.abc import Mapping


class DfValidator:
    """
    Ensures DataFrames (Dask or Pandas) meet target schema and quality standards.
    Optimized for Dask graph efficiency through batched column transformations.
    """

    _cached_client: ClassVar[Client | None] = None

    def __init__(
        self,
        df: pd.DataFrame | dd.DataFrame,
        dask_client: Client | None = None,
        debug: bool = False,
        logger: Logger | None = None,
    ) -> None:
        self.df = df
        self.is_dask = isinstance(df, dd.DataFrame)
        self.dask_client = dask_client

        log_level = 10 if debug else 20
        self.logger = logger or Logger.default_logger(
            logger_name=self.__class__.__name__, log_level=log_level
        )
        self.logger.info("--- DfValidator Initialized ---")

        self._establish_dask_client()

    def _establish_dask_client(self) -> None:
        """Finds or confirms the Dask client connection."""
        if self.is_dask:
            if self.dask_client:
                DfValidator._cached_client = self.dask_client
                return

            if DfValidator._cached_client:
                self.dask_client = DfValidator._cached_client
                self.logger.info("✅ Using cached Dask Client.")
                return

            try:
                if get_client is None:
                    raise ImportError("Dask Distributed not installed.")
                self.dask_client = get_client()
                DfValidator._cached_client = self.dask_client
                self.logger.info(
                    f"✅ Dask Client detected. Dashboard: {self.dask_client.dashboard_link}"
                )
            except (ValueError, ImportError, RuntimeError):
                self.logger.warning("No Dask Client found or installed. Running in local mode.")
        else:
            self.logger.info("Initialized with Pandas DataFrame (local mode).")

    def _normalize_dtype_alias(self, dtype_str: str) -> str:
        """
        Normalizes Dask/PyArrow dtype aliases to canonical Pandas Extension types.
        """
        dtype_str = str(dtype_str).lower()
        mapping = {
            "int64[pyarrow]": "Int64[pyarrow]",
            "int64": "Int64[pyarrow]",
            "int32[pyarrow]": "Int32[pyarrow]",
            "int32": "Int32[pyarrow]",
            "double[pyarrow]": "Float64[pyarrow]",
            "float64[pyarrow]": "Float64[pyarrow]",
            "float64": "Float64[pyarrow]",
            "string": "string[pyarrow]",
            "string[pyarrow]": "string[pyarrow]",
            "timestamp[ns][pyarrow]": "datetime64[ns, UTC]",
            "datetime64[us, utc]": "datetime64[ns, UTC]",
            "datetime64[ns, utc]": "datetime64[ns, UTC]",
            "bool": "boolean[pyarrow]",
            "boolean[pyarrow]": "boolean[pyarrow]",
            "category": "category",
        }
        return mapping.get(dtype_str, dtype_str)

    def apply_schema_map(self, schema_map: dict[str, str]) -> DfValidator:
        """
        Enforces type correctness by casting the current DF to the target dtypes.
        """
        dtype_map = {
            col: dtype for col, dtype in schema_map.items() if col in self.df.columns
        }
        if not dtype_map:
            self.logger.warning("Schema map contained no matching columns. Skipping.")
            return self

        self.logger.info(f"Applying strict schema to {len(dtype_map)} columns...")
        self.df = self.df.astype(dtype_map)
        return self

    def validate_schema(self, expected_schema_map: dict[str, str]) -> DfValidator:
        """
        Checks if the current DF's dtypes match the expected map.
        Raises TypeError on mismatch.
        """
        self.logger.info("Validating current schema against expected standards...")
        errors = []
        for col, expected_dtype in expected_schema_map.items():
            if col in self.df.columns:
                current_dtype_raw = str(self.df.dtypes[col])
                current_dtype_normalized = self._normalize_dtype_alias(current_dtype_raw)

                if current_dtype_normalized != expected_dtype:
                    errors.append(
                        f"❌ Col '{col}': Expected '{expected_dtype}', Found '{current_dtype_raw}'"
                    )

        if errors:
            error_msg = "Schema Validation FAILED:\n" + "\n".join(errors)
            self.logger.error(error_msg)
            raise TypeError(error_msg)

        self.logger.info("✅ Schema Validation PASSED.")
        return self

    def standardize_data_quality(self) -> DfValidator:
        """
        Applies data cleanup fixes using batched transformations to optimize Dask graphs.
        """
        self.logger.info("Standardizing data quality (batched)...")

        # 1. Column Name Standardization
        new_cols = (
            self.df.columns.str.strip()
            .str.lower()
            .str.replace(" ", "_", regex=False)
            .str.replace(r"[^a-z0-9_]", "", regex=True)
        )
        self.df.columns = new_cols

        # 2. Batch Transformations
        transforms = {}

        # String Cleanup
        string_cols = self.df.select_dtypes(
            include=["string", "object", "category"]
        ).columns
        
        for col in string_cols:
            series = self.df[col]
            is_cat = isinstance(series.dtype, pd.CategoricalDtype)
            
            # Avoid redundant casts if already a string-like extension type
            needs_cast = not is_cat and not (
                pd.api.types.is_string_dtype(series.dtype) and 
                getattr(series.dtype, "storage", None) == "pyarrow"
            )
            
            if needs_cast:
                series = series.astype(str)
            
            # For categories, we only strip if it's safe and doesn't break dtype
            # In Dask/Pandas, series.str.strip() on a category returns object.
            # To preserve category, we'd need to modify categories.
            # For now, we skip string cleaning on categories to preserve the memory optimization.
            if is_cat:
                continue

            transforms[col] = (
                series.str.strip()
                .replace({"nan": None, "None": None, "": None})
            )

        # Datetime Normalization
        datetime_cols = self.df.select_dtypes(include=["datetime", "datetimetz"]).columns
        for col in datetime_cols:
            # Enforce naive UTC-referenced datetime (removes tz if present)
            transforms[col] = self.df[col].dt.tz_localize(None)

        if transforms:
            self.logger.debug(f"Applying {len(transforms)} batched transformations.")
            self.df = self.df.assign(**transforms)

        self.logger.info("Standardization complete.")
        return self

    def generate_clickhouse_ddl(
        self, 
        table_name: str, 
        type_map: Mapping[str, str] = DASK_TO_CLICKHOUSE_DTYPE
    ) -> str:
        """
        Generates a ClickHouse CREATE TABLE statement based on the validated schema.
        """
        self.logger.info(f"Generating ClickHouse DDL for '{table_name}'...")
        
        # Extended map to handle special cases consistently
        full_map = dict(type_map)
        # Force LowCardinality for categories regardless of base map
        full_map["category"] = "LowCardinality(String)"

        ddl_parts = []
        for col, dtype in self.df.dtypes.items():
            dtype_str = str(dtype)
            normalized_dtype = self._normalize_dtype_alias(dtype_str)
            ch_type = full_map.get(normalized_dtype)

            if ch_type is None:
                self.logger.warning(f"No map for {dtype_str} on {col}. Using String.")
                ch_type = "String"

            ddl_parts.append(f"    {col} {ch_type}")

        ddl = f"CREATE TABLE {table_name} (\n"
        ddl += ",\n".join(ddl_parts)
        ddl += "\n) ENGINE = MergeTree ORDER BY tuple()"

        self.logger.info("✅ DDL generation complete.")
        return ddl

    def get_df(self) -> pd.DataFrame | dd.DataFrame:
        """Returns the final, validated DataFrame."""
        return self.df
