from __future__ import annotations

from typing import Any, Tuple, Dict, Optional

import dask.dataframe as dd
import pandas as pd

from sibi_flux.core import SecureResource
from sibi_flux.df_helper.core import ParamsConfig, QueryConfig
from ._db_connection import SqlAlchemyConnectionConfig
from ._io_dask import SQLAlchemyDask


class SqlAlchemyLoadFromDb(SecureResource):
    """
    Orchestrates loading data from a database using SQLAlchemy into a Dask DataFrame.
    """

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

    def __init__(
        self,
        plugin_sqlalchemy: SqlAlchemyConnectionConfig,
        plugin_query: Optional[QueryConfig] = None,
        plugin_params: Optional[ParamsConfig] = None,
        **kwargs: Any,
    ):
        super().__init__(**kwargs)
        self.db_connection = plugin_sqlalchemy

        # Access properties safely (triggers lazy load if needed)
        self.model = self.db_connection.model
        self.engine = self.db_connection.engine

        self.query_config = plugin_query
        self.params_config = plugin_params

        # Safe extraction of tuning parameters
        df_params = self.params_config.df_params if self.params_config else None

        # 1. Chunk Size (Default: 50k)
        # Priority: kwargs > df_params > default
        self.chunk_size = int(
            kwargs.get("chunk_size", self._safe_get(df_params, "chunk_size", 50_000))
        )

        # 2. Index Column (Critical for Range Pagination)
        # Priority: kwargs > df_params > None
        self.db_index = kwargs.get(
            "db_index", self._safe_get(df_params, "db_index", None)
        )

        # 3. Safety Limit
        # Priority: kwargs > df_params > None
        self.limit = kwargs.get("limit", self._safe_get(df_params, "limit", None))
        if self.limit is not None:
            self.limit = int(self.limit)

        # 4. Engine Options
        self.execution_options = kwargs.get("execution_options", {})

        self.total_records = -1

    @staticmethod
    def _safe_get(obj: Any, key: str, default: Any) -> Any:
        """Helper to extract values from Dict, Pydantic, or None safely."""
        if obj is None:
            return default
        if isinstance(obj, dict):
            return obj.get(key, default)
        return getattr(obj, key, default)

    def build_and_load(self) -> Tuple[int, dd.DataFrame]:
        try:
            # Determine pagination strategy automatically
            # If we have an index_col, try 'range' (fast). Otherwise 'offset' (safe).
            pagination_mode = "range" if self.db_index else "offset"

            # Extract filters safely
            filters = self.params_config.filters if self.params_config else {}

            # Use Context Manager to ensure SQLAlchemyDask resources are cleaned up
            with SQLAlchemyDask(
                model=self.model,
                filters=filters,
                engine=self.engine,
                chunk_size=self.chunk_size,
                pagination=pagination_mode,
                db_index=self.db_index,
                limit=self.limit,
                execution_options=self.execution_options,
                logger=self.logger,
                verbose=self.verbose,
                debug=self.debug,
            ) as loader:

                self.logger.debug(
                    f"SQLAlchemyDask initialized for {self.model.__name__} "
                    f"(strategy={pagination_mode}, chunk={self.chunk_size})",
                    extra=self.logger_extra,
                )

                self.total_records, dask_df = loader.read_frame()

                return self.total_records, dask_df

        except Exception as e:
            self.total_records = -1
            self.logger.error(
                f"{self.model.__name__} failed to load: {e}",
                exc_info=True,
                extra=self.logger_extra,
            )
            # Return empty dataframe structure on failure to prevent crashes downstream
            try:
                # Try to inspect columns from model to return empty DF with correct schema
                columns = [c.name for c in self.model.__table__.columns]
            except Exception:
                columns = []

            return -1, dd.from_pandas(pd.DataFrame(columns=columns), npartitions=1)

    def _cleanup(self) -> None:
        """
        Clean up instance references to prevent memory leaks.
        Note: The engine is owned by DfHelper/ConnectionConfig, not us.
        """
        try:
            self.logger.debug(f"Cleaning up {self.__class__.__name__} refs")
            self.db_connection = None
            self.engine = None
            self.model = None
        except Exception as e:
            self.logger.warning(f"Error during cleanup: {e}", extra=self.logger_extra)
