from __future__ import annotations

import asyncio
import threading
from typing import Any, Dict, Optional, TypeVar, cast

import dask.dataframe as dd
import pandas as pd
from fsspec import AbstractFileSystem
from pydantic import BaseModel, ValidationError

from sibi_flux.core import SecureResource
from sibi_flux.utils.clickhouse_writer import ClickHouseWriter
from sibi_flux.df_helper.core import QueryConfig, ParamsConfig
from sibi_flux.df_helper.backends.http import HttpConfig
from sibi_flux.df_helper.backends.parquet import ParquetConfig
from sibi_flux.df_helper.backends.sqlalchemy import SqlAlchemyConnectionConfig

from .backends import (
    BackendParamsBase,
    SqlAlchemyBackendParams,
    build_backend_params,
)
from .backends.utils import is_dask_df, maybe_compute, maybe_persist
from .backends import SqlAlchemyBackend, ParquetBackend, HttpBackend

T = TypeVar("T", bound=BaseModel)


class DfHelper(SecureResource):
    _BACKEND_STRATEGIES = {
        "sqlalchemy": SqlAlchemyBackend,
        "parquet": ParquetBackend,
        "http": HttpBackend,
    }

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

    def __init__(self, backend: str = "sqlalchemy", **kwargs):
        self._sql_lock_internal = None
        kwargs.setdefault("auto_sse", False)

        backend_in_kwargs = kwargs.get("backend")
        if backend_in_kwargs is not None:
            backend = backend_in_kwargs
        kwargs["backend"] = backend

        backend_params: BackendParamsBase = build_backend_params(
            kwargs, backend=backend
        )

        kwargs.setdefault("debug", bool(getattr(backend_params, "debug", False)))
        kwargs.setdefault("fs", getattr(backend_params, "fs", None))
        kwargs.setdefault("logger", backend_params.log)

        super().__init__(**kwargs)

        self.backend = backend_params.backend
        self.total_records = -1

        self._backend_query = self._get_config(QueryConfig, kwargs)
        self._backend_params = self._get_config(ParamsConfig, kwargs)

        self.backend_db_connection: Optional[SqlAlchemyConnectionConfig] = None
        self.backend_parquet: Optional[ParquetConfig] = None
        self.backend_http: Optional[HttpConfig] = None

        # SQLAlchemy connection is lazy, only for sqlalchemy backend.
        self._sql_lock: Optional[threading.RLock] = (
            threading.RLock() if self.backend == "sqlalchemy" else None
        )
        self._sql_conn_kwargs: Optional[Dict[str, Any]] = None

        if self.backend == "sqlalchemy":
            bp = cast(SqlAlchemyBackendParams, backend_params)
            self._sql_conn_kwargs = {
                "connection_url": bp.connection_url,
                "table": bp.table,
                "debug": self.debug,
                "logger_extra": self.logger_extra,
            }
        elif self.backend == "parquet":
            self.backend_parquet = self._get_config(ParquetConfig, kwargs)
        elif self.backend == "http":
            self.backend_http = self._get_config(HttpConfig, kwargs)

        strategy_cls = self._BACKEND_STRATEGIES.get(self.backend)
        if not strategy_cls:
            raise ValueError(f"Unsupported backend: {self.backend}")
        self.backend_strategy = strategy_cls(self)

    @property
    def sql_lock(self) -> threading.RLock:
        """Lazy loader for the lock to avoid serialization issues."""
        if not hasattr(self, "_sql_lock_internal") or self._sql_lock_internal is None:
            self._sql_lock_internal = threading.RLock()
        return self._sql_lock_internal

    def get_sql_connection(self) -> SqlAlchemyConnectionConfig:
        if self.backend != "sqlalchemy":
            raise RuntimeError("SQLAlchemy connection only for backend='sqlalchemy'.")

        # Use the property, not a fixed attribute
        with self.sql_lock:
            if self.backend_db_connection is None:
                if not self._sql_conn_kwargs:
                    raise RuntimeError("SQLAlchemy parameters not initialized.")
                self.backend_db_connection = SqlAlchemyConnectionConfig(
                    **self._sql_conn_kwargs
                )
            return self.backend_db_connection

    # ---------- ManagedResource hooks ----------
    def get_sse(self):
        return self._ensure_sse()

    def _emit_bg(self, event: str, **data: Any) -> None:
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            asyncio.run(self.emit(event, **data))
        else:
            task = loop.create_task(self.emit(event, **data))
            # Track task to avoid immediate GC and allow waiting on shutdown
            self._tracked_tasks.add(task)
            task.add_done_callback(self._tracked_tasks.discard)

    def _cleanup(self):
        if self.backend == "sqlalchemy" and self.backend_db_connection is not None:
            try:
                self.backend_db_connection.close()
            finally:
                self.backend_db_connection = None
            return

        if (
            self.backend == "parquet"
            and self.backend_parquet is not None
            and hasattr(self.backend_parquet, "close")
        ):
            self.backend_parquet.close()
            return

        if (
            self.backend == "http"
            and self.backend_http is not None
            and hasattr(self.backend_http, "close")
        ):
            self.backend_http.close()
            return

    async def _acleanup(self):
        active = None
        if self.backend == "sqlalchemy":
            active = self.backend_db_connection
        elif self.backend == "parquet":
            active = self.backend_parquet
        elif self.backend == "http":
            active = self.backend_http

        if not active:
            return

        if hasattr(active, "aclose"):
            await active.aclose()
        elif hasattr(active, "close"):
            await asyncio.to_thread(active.close)

    # ---------- config helpers ----------
    def _get_config(self, model: type[T], kwargs: Dict[str, Any]) -> T:
        """
        Safely extracts and validates configuration from kwargs using Pydantic.

        This method attempts to instantiate the Pydantic model by filtering
        kwargs that match the model's fields (including aliases).
        """
        try:
            # 1. Inspect the model to find valid field names and aliases
            valid_keys = set(model.model_fields.keys())

            # If the model uses aliases (e.g. field 'conn_url' aliased to 'connection_url'),
            # we need to check for those too.
            for field in model.model_fields.values():
                if field.alias:
                    valid_keys.add(field.alias)

            # 2. Filter the incoming kwargs to only those relevant for this model
            # This prevents "Extra fields not permitted" errors if the model is strict,
            # while still passing through the relevant data.
            filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_keys}

            # 3. Use model_validate to leverage Pydantic's full validation logic
            # We wrap the filtered dict in a constructed object or pass directly
            return model.model_validate(filtered_kwargs)

        except ValidationError as e:
            # Enriched error logging to help debug which config failed
            self.logger.error(
                f"Configuration validation failed for {model.__name__}: {e}",
                extra=self.logger_extra,
            )
            raise ValueError(f"Invalid configuration for {model.__name__}: {e}") from e

    def to_params(self) -> Dict[str, Any]:
        """
        Returns a dictionary representation of the configuration.
        """
        # Return merged config from instance state + backend params for completeness
        base_config = self.config.copy() if hasattr(self, "config") else {}
        backend_config = (
            self._backend_params.model_dump(exclude_defaults=True)
            if hasattr(self, "_backend_params")
            else {}
        )
        return {**base_config, **backend_config}

    # ---------- load/aload ----------
    def load(
        self, *, persist: bool = False, as_pandas: bool = False, **options
    ) -> pd.DataFrame | dd.DataFrame:
        self.logger.debug(
            f"Loading data from {self.backend} backend with options: {options}",
            extra=self.logger_extra,
        )
        self.total_records, df = self.backend_strategy.load(**options)
        df = self._process_loaded_data(df)
        df = self._post_process_df(df)
        df = maybe_persist(df, persist)
        return maybe_compute(df, as_pandas)

    async def aload(
        self,
        *,
        persist: bool = False,
        as_pandas: bool = False,
        timeout: Optional[float] = None,
        **options,
    ) -> pd.DataFrame | dd.DataFrame:
        await self.emit(
            f"{self.__class__.__name__} load:start",
            message=f"Pulling data from {self.backend} backend",
        )

        load_awaitable = self.backend_strategy.aload(**options)
        total, df = await (
            asyncio.wait_for(load_awaitable, timeout) if timeout else load_awaitable
        )
        self.total_records = total

        await self.emit(
            event=f"{self.__class__.__name__} load:progress",
            message="Post-processing records",
        )
        df = await asyncio.to_thread(self._process_loaded_data, df)
        df = await asyncio.to_thread(self._post_process_df, df)

        if persist and is_dask_df(df):
            df = await asyncio.to_thread(df.persist)
        if as_pandas and is_dask_df(df):
            compute_awaitable = asyncio.to_thread(df.compute)
            return await (
                asyncio.wait_for(compute_awaitable, timeout)
                if timeout
                else compute_awaitable
            )

        await self.emit(
            event=f"{self.__class__.__name__} load:progress",
            message="Returning records",
        )
        return df

    # ---------- dataframe post-processing ----------
    def _post_process_df(
        self, df: dd.DataFrame | pd.DataFrame
    ) -> dd.DataFrame | pd.DataFrame:
        df_params = getattr(self._backend_params, "df_params", None)
        if not df_params:
            return df

        fieldnames = df_params.fieldnames
        column_names = df_params.column_names
        index_col = df_params.index_col

        if fieldnames and hasattr(df, "columns"):
            valid = [f for f in fieldnames if f in df.columns]
            if len(valid) < len(fieldnames):
                self.logger.warning(
                    f"Missing columns for filtering: {set(fieldnames) - set(valid)}",
                    extra=self.logger_extra,
                )
            df = df[valid]

        # Renames columns when names are provided and valid
        if column_names:
            if len(df.columns) == 0:
                # Handle case where backend returns empty DF with no columns
                if isinstance(df, pd.DataFrame):
                    df = pd.DataFrame(columns=column_names)
                else:
                    df = dd.from_pandas(
                        pd.DataFrame(columns=column_names), npartitions=1
                    )
            elif len(df.columns) != len(column_names):
                raise ValueError(
                    f"Length mismatch: DataFrame has {len(df.columns)} columns, but {len(column_names)} names were provided."
                )
            else:
                df = df.rename(columns=dict(zip(df.columns, column_names)))

        if index_col:
            if index_col not in df.columns:
                raise ValueError(f"Index column '{index_col}' not found in DataFrame.")
            df = df.set_index(index_col)

        return df

    def _process_loaded_data(
        self, df: dd.DataFrame | pd.DataFrame
    ) -> dd.DataFrame | pd.DataFrame:
        field_map = getattr(self._backend_params, "field_map", None) or {}
        if not isinstance(field_map, dict) or not field_map:
            return df

        rename_map = {k: v for k, v in field_map.items() if k in df.columns}
        return df.rename(columns=rename_map) if rename_map else df

    # ---------- sinks ----------
    def save_to_parquet(self, df: dd.DataFrame, show_progress: bool = False, **kwargs):
        fs: AbstractFileSystem = kwargs.pop("fs", self.fs)
        path: str = kwargs.pop(
            "parquet_storage_path",
            self.backend_parquet.parquet_storage_path if self.backend_parquet else None,
        )
        parquet_filename = kwargs.pop(
            "parquet_filename",
            self.backend_parquet.parquet_filename if self.backend_parquet else None,
        )

        # if not parquet_filename:
        #    raise ValueError("A 'parquet_filename' keyword argument must be provided.")
        if not fs:
            raise ValueError(
                "A filesystem (fs) must be provided to save the parquet file."
            )
        if not path:
            raise ValueError(
                "A 'parquet_storage_path' keyword argument must be provided."
            )
        if not self._has_any_rows(df):
            self.logger.warning(
                "Skipping save: The provided DataFrame is empty.",
                extra=self.logger_extra,
            )
            return

        from sibi_flux.parquet import ParquetSaver

        with ParquetSaver(
            df_result=df,
            parquet_storage_path=path,
            fs=fs,
            debug=self.debug,
            logger=self.logger,
            verbose=self.verbose,
            show_progress=show_progress,
            **kwargs,
        ) as saver:
            saver.save_to_parquet(parquet_filename)

        self.logger.debug(
            f"Successfully saved '{parquet_filename}' to '{path}'.",
            extra=self.logger_extra,
        )

    async def asave_to_parquet(self, df: dd.DataFrame, show_progress: bool = False, **kwargs):
        await self.emit(
            event=f"{self.__class__.__name__} save:start", message="Saving to parquet"
        )
        await asyncio.to_thread(self.save_to_parquet, df, show_progress=show_progress, **kwargs)
        await self.emit(
            event=f"{self.__class__.__name__} save:end", message="Saved to parquet"
        )

    def save_to_clickhouse(self, df: dd.DataFrame, show_progress: bool = False, **credentials):
        if not self._has_any_rows(df):
            self.logger.warning(
                "Skipping save to ClickHouse: The provided DataFrame is empty.",
                extra=self.logger_extra,
            )
            return
        with ClickHouseWriter(
            debug=self.debug, logger=self.logger, verbose=self.verbose, 
            show_progress=show_progress, **credentials
        ) as writer:
            writer.save_to_clickhouse(df)
            self.logger.debug("Save to ClickHouse completed.", extra=self.logger_extra)

    async def asave_to_clickhouse(self, df: dd.DataFrame, show_progress: bool = False, **credentials):
        await self.emit(
            event=f"{self.__class__.__name__} save:start",
            message="Saving to ClickHouse",
        )
        await asyncio.to_thread(self.save_to_clickhouse, df, show_progress=show_progress, **credentials)
        await self.emit(
            event=f"{self.__class__.__name__} save:end", message="Saved to ClickHouse"
        )

    # ---------- period loaders ----------
    def load_period(self, dt_field: str, start: str, end: str, **kwargs):
        final_kwargs = self._prepare_period_filters(dt_field, start, end, **kwargs)
        return self.load(**final_kwargs)

    async def aload_period(self, dt_field: str, start: str, end: str, **kwargs):
        final_kwargs = self._prepare_period_filters(dt_field, start, end, **kwargs)
        return await self.aload(**final_kwargs)

    def _prepare_period_filters(
        self, dt_field: str, start: str, end: str, **kwargs
    ) -> dict:
        start_date, end_date = pd.to_datetime(start).date(), pd.to_datetime(end).date()
        if start_date > end_date:
            raise ValueError("'start' date cannot be later than 'end' date.")

        field_map = getattr(self._backend_params, "field_map", None) or {}
        reverse_map = {v: k for k, v in field_map.items()} if field_map else {}
        if len(reverse_map) != len(field_map):
            self.logger.warning(
                "field_map values are not unique; reverse mapping may be unreliable.",
                extra=self.logger_extra,
            )

        mapped_field = reverse_map.get(dt_field, dt_field)
        if start_date == end_date:
            kwargs[f"{mapped_field}__date"] = start_date
        else:
            kwargs[f"{mapped_field}__date__range"] = [start_date, end_date]

        self.logger.debug(
            f"Period load generated filters: {kwargs}", extra=self.logger_extra
        )
        return kwargs

    @staticmethod
    def _has_any_rows(ddf: dd.DataFrame) -> bool:
        try:
            # O(1) check by looking for the first row only
            return bool(ddf.head(1).shape[0])
        except Exception:
            return False
