from __future__ import annotations

import asyncio
import copy
from typing import TYPE_CHECKING, Any, Optional, Dict

import dask.dataframe as dd
import pandas as pd
import pyarrow as pa

# Import the filter handler and SQL loader
from sibi_flux.df_helper.core import FilterHandler
from .sqlalchemy import SqlAlchemyLoadFromDb

if TYPE_CHECKING:
    from sibi_flux.df_helper import DfHelper


class BaseBackend:
    def __init__(self, helper: DfHelper):
        self.helper = helper
        self.logger = helper.logger
        self.debug = helper.debug
        self.total_records = -1

    def load(self, **options) -> tuple[int, dd.DataFrame | pd.DataFrame]:
        raise NotImplementedError

    async def aload(self, **options) -> tuple[int, dd.DataFrame | pd.DataFrame]:
        return await asyncio.to_thread(self.load, **options)


class SqlAlchemyBackend(BaseBackend):
    def load(self, **options) -> tuple[int, dd.DataFrame | pd.DataFrame]:
        try:
            # 1. Thread-Safe Params Update
            # Create a new config instance with the runtime filters applied.
            # This prevents the "Select *" hang by ensuring filters are respected.
            if hasattr(self.helper._backend_params, "with_updates"):
                # Immutable API
                local_params = self.helper._backend_params.with_updates(**options)
            else:
                # Fallback API (Legacy)
                local_params = self._clone_params(self.helper._backend_params)
                if options and hasattr(local_params, "parse_params"):
                    local_params.parse_params(options)

            # 2. Get Connection
            plugin_sqlalchemy = self.helper.get_sql_connection()

            # 3. Load
            with SqlAlchemyLoadFromDb(
                plugin_sqlalchemy=plugin_sqlalchemy,
                plugin_query=self.helper._backend_query,
                plugin_params=local_params,  # Pass the LOCAL copy
                logger=self.logger,
                debug=self.debug,
            ) as db_loader:
                self.total_records, result = db_loader.build_and_load()
                return self.total_records, result

        except Exception as e:
            self.logger.error(
                f"SQLAlchemy load failed: {e}",
                exc_info=True,
                extra=self.helper.logger_extra,
            )
            raise

    def _clone_params(self, params: Any) -> Any:
        if hasattr(params, "model_copy"):
            return params.model_copy(deep=True)
        if hasattr(params, "copy"):
            return params.copy(deep=True)
        return copy.deepcopy(params)


class ParquetBackend(BaseBackend):
    def load(self, **options) -> tuple[int, dd.DataFrame | pd.DataFrame]:
        base_config = self.helper.backend_parquet
        if not base_config:
            raise ValueError("Parquet backend is not configured.")

        try:
            # 1. Update Parameters (Centralized via ParamsConfig)
            # This handles field mapping, filter distribution, etc.
            local_params = self.helper._backend_params.with_updates(**options)
            
            # 2. Config Separation (for Parquet specific fields)
            # We still need to extract Parquet-specific options from 'options'
            config_fields = set()
            if hasattr(base_config, "model_fields"):
                config_fields = set(base_config.model_fields.keys())
            elif hasattr(base_config, "__fields__"):
                config_fields = set(base_config.__fields__.keys())

            config_updates = {k: v for k, v in options.items() if k in config_fields}
            
            # Extract Limit (if present in root or options)
            limit: Optional[int] = options.get("limit")

            # 3. Update Config (create local copy)
            if hasattr(base_config, "model_copy"):
                local_config = base_config.model_copy(update=config_updates)
            elif hasattr(base_config, "copy"):
                local_config = base_config.copy(update=config_updates)
            else:
                local_config = copy.deepcopy(base_config)
                for k, v in config_updates.items():
                    setattr(local_config, k, v)

            # 4. Use Centralized Filters (Already mapped if legacy_filters=True or auto-detected)
            final_filters = local_params.filters.copy()

            handler = FilterHandler(
                backend="dask", logger=self.logger, debug=self.debug
            )

            pushdown = None
            residual: Dict[str, Any] = {}

            # 5. Optimistic Split (Try to push down)
            if final_filters:
                pushdown, residual = handler.split_pushdown_and_residual(final_filters)
                if pushdown:
                    self.logger.debug(
                        f"Attempting parquet pushdown: {pushdown}",
                        extra=self.helper.logger_extra,
                    )

            # 5. Load with Fallback Mechanism
            try:
                # Attempt Optimized Load
                df = local_config.load_files(filters=pushdown)

            except Exception as e:
                # Check for PyArrow Schema Mismatch (String vs Timestamp)
                is_arrow_error = isinstance(
                    e, (pa.lib.ArrowNotImplementedError, pa.lib.ArrowInvalid)
                )
                error_msg = str(e).lower()
                is_schema_msg = (
                    "no kernel matching input types" in error_msg
                    or "arrownotimplementederror" in str(type(e)).lower()
                )

                if (is_arrow_error or is_schema_msg) and pushdown:
                    self.logger.warning(
                        "Parquet pushdown failed (Schema Mismatch: likely String column vs Timestamp filter). "
                        "Retrying with full scan and in-memory filtering.",
                        extra=self.helper.logger_extra,
                    )
                    # REVERT OPTIMIZATION:
                    # 1. Clear pushdown (load everything)
                    # 2. Move ALL filters to residual (filter in memory)
                    pushdown = None
                    residual = final_filters

                    # Retry Load Raw
                    df = local_config.load_files(filters=None)
                else:
                    # If it's a file not found or other I/O error, crash loudly
                    raise e

            # 6. Apply Residual Filters
            # (If fallback triggered, this now contains ALL filters)
            if residual:
                self.logger.debug(
                    f"Applying residual filters: {residual.keys()}",
                    extra=self.helper.logger_extra,
                )
                mask_fn = handler.build_mask_fn(residual)
                df = df[mask_fn(df)]

            # 7. Safety Limit (Eager Truncation)
            if limit is not None:
                # Eagerly compute the head to capping memory usage
                self.logger.info(
                    f"Applying safety limit: {limit} rows",
                    extra=self.helper.logger_extra,
                )
                pdf = df.head(limit, npartitions=-1, compute=True)
                # Re-partition back to Dask (lazy) to satisfy return signature
                # Use npartitions=1 as the result is presumably small
                df = dd.from_pandas(pdf, npartitions=1)
                self.total_records = len(pdf)
                return self.total_records, df

            # 7. Check Empty
            if not self.helper._has_any_rows(df):
                self.total_records = 0
                return 0, self._empty_like(df)

            self.total_records = -1
            return self.total_records, df

        except Exception as e:
            self.logger.error(
                f"Parquet load failed: {e}",
                exc_info=True,
                extra=self.helper.logger_extra,
            )
            raise

    @staticmethod
    def _empty_like(ddf: dd.DataFrame) -> dd.DataFrame:
        """Creates an empty Dask DataFrame preserving the original schema."""
        return dd.from_pandas(ddf._meta, npartitions=1)


class HttpBackend(BaseBackend):
    def load(self, **options):
        raise RuntimeError("HttpBackend is async-only. Use `await aload(...)`.")

    async def aload(self, **options) -> tuple[int, dd.DataFrame | pd.DataFrame]:
        backend_config = self.helper.backend_http
        if not backend_config:
            raise ValueError("HTTP backend is not configured.")

        try:
            # 1. Update Parameters (Centralized via ParamsConfig)
            # This handles field mapping, filter distribution, etc.
            local_params = self.helper._backend_params.with_updates(**options)

            # 2. Merge Config
            # Creates a new HttpBackendParams instance merged with runtime options
            # We use local_params.filters which are already mapped/distributed
            config_copy = self._merge_http_config(backend_config, local_params.filters)

            # 2. Fetch
            # Assuming fetch_data is on the config object
            # If fetch_data expects **kwargs for params, adjust accordingly.
            if hasattr(config_copy, "fetch_data"):
                result = await config_copy.fetch_data()
            else:
                # Fallback if fetch logic is external (unlikely in your setup)
                result = await self.helper.backend_http.fetch_data(
                    config_override=config_copy
                )

        except Exception as e:
            self.logger.error(f"HTTP fetch failed: {e}", exc_info=True)
            raise

        # 3. Normalize
        if isinstance(result, (list, dict)):
            pdf = pd.DataFrame(result)
        elif isinstance(result, pd.DataFrame):
            pdf = result
        else:
            return 0, dd.from_pandas(pd.DataFrame(), npartitions=1)

        self.total_records = len(pdf)
        if self.total_records == 0:
            return 0, dd.from_pandas(pd.DataFrame(), npartitions=1)

        npartitions = max(1, min(32, self.total_records // 50_000))
        ddf = dd.from_pandas(pdf, npartitions=npartitions)
        return self.total_records, ddf

    def _merge_http_config(self, base_config: Any, options: Dict[str, Any]) -> Any:
        if hasattr(base_config, "model_copy"):
            new_config = base_config.model_copy(deep=True)
        elif hasattr(base_config, "copy"):
            new_config = base_config.copy(deep=True)
        else:
            new_config = copy.deepcopy(base_config)

        config_fields = set()
        if hasattr(new_config, "model_fields"):
            config_fields = set(new_config.model_fields.keys())

        params_to_merge = {}
        for k, v in options.items():
            if k in config_fields:
                setattr(new_config, k, v)
            else:
                params_to_merge[k] = v

        if params_to_merge and hasattr(new_config, "params"):
            current = getattr(new_config, "params") or {}
            new_params = current.copy()
            new_params.update(params_to_merge)
            setattr(new_config, "params", new_params)

        return new_config
