from __future__ import annotations

import datetime as dt
import inspect
from typing import (
    TYPE_CHECKING,
    Any,
)

import dask.dataframe as dd
import pandas as pd

from sibi_flux.logger import Logger

if TYPE_CHECKING:
    from sibi_flux.dataset.protocols import DatasetReaderType


class HybridDataLoader:
    """
    Hybrid loader that merges historical (Parquet) and live (API/DB) data.
    """

    def __init__(
        self,
        start_date: str,
        end_date: str,
        historical_reader: DatasetReaderType,
        live_reader: DatasetReaderType,
        date_field: str,
        reference_date: dt.date | None = None,
        **kwargs: Any,
    ) -> None:
        self.start_date = self._validate_date_format(start_date)
        self.end_date = self._validate_date_format(end_date)
        self.historical_reader = historical_reader
        self.live_reader = live_reader
        self.date_field = date_field

        self.logger: Logger = kwargs.get("logger") or Logger.default_logger(
            logger_name=__name__
        )
        self.debug = kwargs.get("debug", False)

        self._validate_date_range()

        # Reference date for time-based logic (defaults to today)
        ref_date = reference_date or dt.date.today()
        self._today_str = ref_date.strftime("%Y-%m-%d")
        self._yesterday_str = (ref_date - dt.timedelta(days=1)).strftime("%Y-%m-%d")

        self._read_live_flag = self.end_date == self._today_str
        self._is_single_today = self.start_date == self.end_date == self._today_str

    @staticmethod
    def _validate_date_format(date_str: str) -> str:
        try:
            dt.datetime.strptime(date_str, "%Y-%m-%d")
            return date_str
        except ValueError:
            raise ValueError(f"Date '{date_str}' is not in valid YYYY-MM-DD format")

    def _validate_date_range(self) -> None:
        start = dt.datetime.strptime(self.start_date, "%Y-%m-%d").date()
        end = dt.datetime.strptime(self.end_date, "%Y-%m-%d").date()
        if end < start:
            raise ValueError(
                f"End date ({self.end_date}) cannot be before start date ({self.start_date})"
            )

    @staticmethod
    def _normalize_datetimes(df: dd.DataFrame, cols: list[str]) -> dd.DataFrame:
        """Normalize datetime columns to UTC safely and efficiently."""
        if not cols:
            return df

        valid_cols = [c for c in cols if c in df.columns]
        if not valid_cols:
            return df

        # Vectorized batch transformation within map_partitions
        def _to_ts_vectorized(pdf: pd.DataFrame) -> pd.DataFrame:
            pdf = pdf.copy() # Avoid SettingWithCopyWarning
            for c in valid_cols:
                pdf[c] = pd.to_datetime(pdf[c], errors="coerce", utc=True)
            return pdf

        return df.map_partitions(_to_ts_vectorized)

    @staticmethod
    def _create_empty_dataframe(meta: pd.DataFrame | None = None) -> dd.DataFrame:
        if meta is not None and not meta.empty:
            return dd.from_pandas(meta.iloc[0:0], npartitions=1)
        return dd.from_pandas(pd.DataFrame(), npartitions=1)

    async def _load_today_data(self, **kwargs: Any) -> dd.DataFrame | None:
        """Load today's live data."""
        self.logger.debug("Loading today's live data...")
        date_filter = {f"{self.date_field}__date": self._today_str}
        filters = {**kwargs, **date_filter}

        try:
            reader_obj = self.live_reader(logger=self.logger, debug=self.debug)
            today_df = await reader_obj.aload(**filters)
            if today_df is not None:
                today_df = self._normalize_datetimes(today_df, [self.date_field])
            return today_df
        except Exception as e:
            self.logger.error(f"Failed to load today's data: {e}")
            return None

    async def _load_historical_data(
        self, start_date: str, end_date: str, **kwargs: Any
    ) -> dd.DataFrame:
        """Load historical data."""
        self.logger.debug(f"Loading historical data from {start_date} to {end_date}...")
        try:
            reader_params = {"logger": self.logger, "debug": self.debug}
            
            sig = inspect.signature(self.historical_reader.__init__)
            if "parquet_start_date" in sig.parameters:
                reader_params["parquet_start_date"] = start_date
                reader_params["parquet_end_date"] = end_date

            reader_obj = self.historical_reader(**reader_params)
            df = await reader_obj.aload(**kwargs)
            df = self._normalize_datetimes(df, [self.date_field])
            return df
        except Exception as e:
            self.logger.error(f"Failed to load historical data: {e}")
            return self._create_empty_dataframe()

    async def aload(self, **kwargs: Any) -> dd.DataFrame:
        """Load and concatenate data from historical and live sources."""
        self.logger.debug(
            f"[HybridLoader] start={self.start_date}, end={self.end_date}, "
            f"read_live={self._read_live_flag}"
        )

        # Case 1: only today
        if self._is_single_today:
            today_df = await self._load_today_data(**kwargs)
            return today_df if today_df is not None else self._create_empty_dataframe()

        # Case 2: purely historical
        if not self._read_live_flag:
            return await self._load_historical_data(
                self.start_date, self.end_date, **kwargs
            )

        # Case 3: mixed historical + live
        hist_df = await self._load_historical_data(
            self.start_date, self._yesterday_str, **kwargs
        )
        live_df = await self._load_today_data(**kwargs)

        if hist_df is None:
            hist_df = self._create_empty_dataframe()
        if live_df is None:
            live_df = self._create_empty_dataframe()

        try:
            # Partition stability check
            n_hist = hist_df.npartitions
            n_live = live_df.npartitions

            # Only interleave if divisions are definitely known, otherwise simple concat is faster
            known_divisions = (
                hist_df.known_divisions and live_df.known_divisions
            )
            
            # If divisions are NOT known, interleave_partitions=True causes a full shuffle
            # We strictly avoid interleaving unless divisions are known and counts are comparable.
            do_interleave = known_divisions and (
                0.2 < (n_hist / n_live) < 5
            )

            return dd.concat(
                 [hist_df, live_df],
                 axis=0,
                 ignore_unknown_divisions=True,
                 interleave_partitions=do_interleave,
            )
        except Exception as e:
            self.logger.warning(
                f"Concat failed: {e}. Falling back to basic concat."
            )
            return dd.concat(
                [hist_df, live_df], axis=0, ignore_unknown_divisions=True
            )

    def __repr__(self) -> str:
        return (
            f"HybridDataLoader(start='{self.start_date}', end='{self.end_date}', "
            f"read_live={self._read_live_flag})"
        )
