from __future__ import annotations

from contextlib import contextmanager
from typing import Any, Dict, Optional, Tuple, Type, List

import dask
import dask.dataframe as dd
import pandas as pd
from sqlalchemy import func, inspect, select, create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.sql import Select
from sqlalchemy.pool import NullPool

from sibi_flux.df_helper.core import FilterHandler
from sibi_flux.core import SecureResource
from sibi_flux.core.type_maps import SQLALCHEMY_TO_DASK_DTYPE

from ._db_gatekeeper import DBGatekeeper

import logging
import warnings

logger = logging.getLogger(__name__)

warnings.simplefilter(action="ignore", category=FutureWarning)

# Global cache for dynamic chunk sizes to avoid repeated expensive sampling
_CHUNK_SIZE_CACHE: Dict[str, int] = {}


class SQLAlchemyDask(SecureResource):
    """
    Production-grade DB -> Dask loader.
    Supports OFFSET (simple) and RANGE (fast) pagination.
    Enforces UTC Everywhere for datetime columns.
    """

    _SQLALCHEMY_TO_DASK_DTYPE: Dict[str, str] = SQLALCHEMY_TO_DASK_DTYPE

    logger_extra: Dict[str, Any] = {"sibi_flux_component": __name__}

    def __init__(
        self,
        model: Type[Any],
        *,
        engine: Engine,
        filters: Optional[Dict[str, Any]] = None,
        chunk_size: int = 50_000,
        pagination: str = "offset",
        db_index: Optional[str] = None,
        num_workers: int = 1,
        **kwargs: Any,
    ):
        super().__init__(**kwargs)
        if pagination == "range" and not db_index:
            self.logger.warning(
                "Pagination 'range' requested but no db_index provided. Falling back to 'offset'."
            )
            pagination = "offset"

        limit_val = kwargs.get("limit")
        self.limit = int(limit_val) if limit_val is not None else None
        self.execution_options = kwargs.get("execution_options", {})

        if self.limit and pagination == "range":
            self.logger.info("Limit specified; forcing pagination='offset'.")
            pagination = "offset"

        self.model = model
        self.engine = engine
        self.filters = filters or {}
        self.chunk_size = int(chunk_size)
        self.pagination = pagination
        self.db_index = db_index
        self.num_workers = int(num_workers)
        self.filter_handler_cls = FilterHandler
        self.total_records: int = -1

        # Use normalized engine key for caching
        self._gate_key = self._normalized_engine_key(self.engine)
        self._cache_key = f"{self._gate_key}:{self.model.__name__}"

        # Dynamic Partitioning Strategy
        if self.chunk_size == 50_000:
            if self._cache_key in _CHUNK_SIZE_CACHE:
                self.chunk_size = _CHUNK_SIZE_CACHE[self._cache_key]
            else:
                self.chunk_size = self._calculate_dynamic_chunk_size()
                _CHUNK_SIZE_CACHE[self._cache_key] = self.chunk_size

        pool_size, max_overflow = self._engine_pool_limits()
        pool_capacity = max(1, pool_size + max_overflow)
        self.per_proc_cap = max(1, pool_capacity // max(1, self.num_workers))

        self._meta_dtypes = self.infer_meta_from_model(self.model)
        self._meta_df = self._build_meta()

    @classmethod
    def infer_meta_from_model(cls, model: Type[Any]) -> Dict[str, str]:
        mapper = inspect(model)
        dtypes: Dict[str, str] = {}
        for column in mapper.columns:
            # Use SQLAlchemy's type name if possible, or split string as fallback
            try:
                type_name = column.type.__class__.__name__.upper()
            except AttributeError:
                type_name = str(column.type).upper().split("(")[0]
            
            dtypes[column.name] = cls._SQLALCHEMY_TO_DASK_DTYPE.get(type_name, "string")
        return dtypes

    def _calculate_dynamic_chunk_size(self) -> int:
        """
        Estimates the optimal chunk size (rows) to target ~100MB partitions.
        Fetches a small sample to calculate average row size.
        Clamps between 1,000 and 1,000,000 rows.
        """
        TARGET_PARTITION_SIZE_MB = 100
        SAMPLE_SIZE = 1000

        try:
            # We use a raw connection to avoid overhead
            query = select(self.model).limit(SAMPLE_SIZE)
            with self._conn() as conn:
                # We need to use the pandas alignment logic to get realistic memory usage
                # because the raw DB tuples are smaller than the final Pandas Series.
                from sqlalchemy import text

                # Re-use fetch logic for consistency or just simpler load
                # Simplified: load to dataframe
                df = pd.read_sql(query, conn)

                if df.empty:
                    return 50_000

                # Measure memory usage (deep=True traverses objects/strings)
                total_bytes = df.memory_usage(deep=True).sum()
                avg_row_bytes = total_bytes / len(df)

                if avg_row_bytes <= 0:
                    return 50_000

                target_bytes = TARGET_PARTITION_SIZE_MB * 1024 * 1024
                optimal_rows = int(target_bytes / avg_row_bytes)

                # Clamp safety limits
                optimal_rows = max(1_000, min(optimal_rows, 1_000_000))

                self.logger.info(
                    f"Dynamic Partitioning: Avg Row={avg_row_bytes:.1f}B. "
                    f"Target={TARGET_PARTITION_SIZE_MB}MB. "
                    f"Calculated Chunk={optimal_rows} rows."
                )
                return optimal_rows

        except Exception as e:
            self.logger.warning(
                f"Failed to calculate dynamic chunk size: {e}. Using default."
            )
            return 50_000

    def _build_meta(self) -> pd.DataFrame:
        return pd.DataFrame(
            {col: pd.Series(dtype=dtype) for col, dtype in self._meta_dtypes.items()}
        )

    @contextmanager
    def _conn(self):
        sem = DBGatekeeper.get(self._gate_key, max_concurrency=self.per_proc_cap)
        with sem:
            with self.engine.connect() as c:
                yield c

    @classmethod
    def _fetch_partition_static(
        cls,
        *,
        query_str: str,
        connection_url: str,
        gate_key: str,
        concurrency_cap: int,
        meta_dtypes: Dict[str, str],
        model: Optional[Type[Any]] = None,
        engine: Optional[Engine] = None,
        meta: Optional[Dict[str, str]] = None,
        filters: Optional[Dict[str, Any]] = None,
        limit: Optional[int] = None,
        execution_options: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> pd.DataFrame:
        sem = DBGatekeeper.get(gate_key, max_concurrency=concurrency_cap)

        try:
            worker_engine = create_engine(
                connection_url,
                poolclass=NullPool,
                execution_options=execution_options or {},
            )
            with sem:
                with worker_engine.connect() as conn:
                    from sqlalchemy import text

                    result = conn.execute(text(query_str))
                    # Efficiently build DataFrame from result keys and rows
                    keys = list(result.keys())
                    df = pd.DataFrame(result.fetchall(), columns=keys)
                    return cls._align_and_coerce_static(df, meta_dtypes)
        except Exception as e:
            logger.error(f"Partition fetch failed: {e}", exc_info=True)
            empty_cols = list(meta_dtypes.keys())
            return pd.DataFrame(columns=empty_cols)

    @staticmethod
    def _align_and_coerce_static(
        df: pd.DataFrame, meta_dtypes: Dict[str, str]
    ) -> pd.DataFrame:
        out = pd.DataFrame(index=df.index)

        for col, target_dtype in meta_dtypes.items():
            if col not in df.columns:
                out[col] = pd.Series(pd.NA, index=df.index, dtype=target_dtype)
                continue

            s = df[col]
            try:
                if "datetime" in target_dtype or "timestamp" in target_dtype:
                    # 'format="mixed"' is critical to handle columns containing
                    # mixed naive/aware strings without crashing or returning NaT.
                    s = pd.to_datetime(s, errors="coerce", utc=True, format="mixed")
                    out[col] = s

                elif target_dtype == "boolean":
                    if not pd.api.types.is_bool_dtype(s.dtype):
                        s = pd.to_numeric(s, errors="coerce").fillna(0).astype(bool)
                    out[col] = s.astype("boolean")

                else:
                    out[col] = s.astype(target_dtype)

            except Exception:
                if "int" in target_dtype.lower():
                    out[col] = pd.to_numeric(s, errors="coerce")
                else:
                    out[col] = s.astype(str)

        return out

    def read_frame(self) -> Tuple[int, dd.DataFrame]:
        base_query = select(self.model)
        if self.filters:
            handler = self.filter_handler_cls(backend="sqlalchemy", logger=self.logger)
            base_query = handler.apply_filters(base_query, self.model, self.filters)

        total = self._count_total(base_query)
        if self.limit is not None and total > self.limit:
            total = self.limit
            base_query = base_query.limit(self.limit)

        self.total_records = total

        if total <= 0:
            return 0, dd.from_pandas(self._meta_df, npartitions=1)

        partitions = []

        if self.pagination == "range" and self.db_index:
            min_id, max_id = self._get_global_min_max_id()

            if min_id is not None and max_id is not None:
                id_span = max_id - min_id
                estimated_parts = id_span / max(1, self.chunk_size)

                # Heuristic: If range implies >2000 partitions for a small dataset,
                # the index is likely extremely sparse (e.g. IDs 1, 50, ..., 1_000_000).
                # Fall back to OFFSET to avoid launching thousands of empty tasks.
                if estimated_parts > 2000 and total < 50000:
                    self.logger.warning("Range too sparse for data size. Using OFFSET.")
                    partitions = self._plan_offset_partitions(base_query, total)
                else:
                    self.logger.debug(
                        f"Planning range partitions (Global Min: {min_id}, Max: {max_id})"
                    )
                    partitions = self._generate_range_partitions(
                        base_query, min_id, max_id
                    )
            else:
                self.logger.warning(
                    "Could not determine Min/Max ID. Falling back to OFFSET."
                )

        if not partitions:
            partitions = self._plan_offset_partitions(base_query, total)

        compiled_sqls = []
        try:
            for q in partitions:
                compiled = q.compile(
                    dialect=self.engine.dialect, compile_kwargs={"literal_binds": True}
                )
                compiled_sqls.append(str(compiled))
        except Exception as e:
            self.logger.error(f"Failed to compile queries: {e}")
            raise

        full_conn_str = self._get_raw_connection_url(self.engine)
        gate_key_safe = self._normalized_engine_key(self.engine)

        delayed_parts = [
            dask.delayed(self._fetch_partition_static)(
                query_str=sql,
                connection_url=full_conn_str,
                gate_key=gate_key_safe,
                concurrency_cap=self.per_proc_cap,
                meta_dtypes=self._meta_dtypes,
                execution_options=self.execution_options,
            )
            for sql in compiled_sqls
        ]

        ddf = dd.from_delayed(delayed_parts, meta=self._meta_df, verify_meta=False)
        return total, ddf

    def _get_global_min_max_id(self) -> Tuple[Optional[int], Optional[int]]:
        col = getattr(self.model, self.db_index)
        try:
            with self._conn() as conn:
                stmt = select(func.min(col), func.max(col))
                min_id, max_id = conn.execute(stmt).one()
                return min_id, max_id
        except Exception:
            return None, None

    def _generate_range_partitions(
        self, base_query: Select, min_id: int, max_id: int
    ) -> List[Select]:
        col = getattr(self.model, self.db_index)
        partitions = []

        current_start = int(min_id)
        max_id_int = int(max_id)
        chunk = int(self.chunk_size)

        while current_start <= max_id_int:
            current_end = current_start + chunk
            part_query = base_query.where(col >= current_start).where(col < current_end)
            partitions.append(part_query)
            current_start = current_end

        return partitions

    def _plan_offset_partitions(self, base_query: Select, total: int) -> List[Select]:
        if self.db_index:
            col = getattr(self.model, self.db_index)
            base_query = base_query.order_by(col)

        return [
            base_query.limit(self.chunk_size).offset(offset)
            for offset in range(0, total, self.chunk_size)
        ]

    def _count_total(self, query: Select) -> int:
        try:
            with self._conn() as conn:
                count_q = query.with_only_columns(func.count()).order_by(None)
                return conn.execute(count_q).scalar_one()
        except Exception:
            self.logger.error("Failed to count total records.", exc_info=True)
            return -1

    def _engine_pool_limits(self) -> Tuple[int, int]:
        pool = getattr(self.engine, "pool", None)

        def to_int(val, default):
            try:
                return int(val() if callable(val) else val)
            except Exception:
                return default

        size = to_int(getattr(pool, "size", None), 5)
        overflow = to_int(
            getattr(pool, "max_overflow", None) or getattr(pool, "_max_overflow", None),
            10,
        )
        return size, overflow

    @staticmethod
    def _normalized_engine_key(engine: Engine) -> str:
        try:
            return str(engine.url.set(query=None).set(password=None))
        except Exception:
            return str(engine.url)

    @staticmethod
    def _get_raw_connection_url(engine: Engine) -> str:
        if hasattr(engine.url, "render_as_string"):
            return engine.url.render_as_string(hide_password=False)
        return str(engine.url)
