import datetime
from typing import Union, Tuple, Callable, Dict, Optional

import fsspec
import numpy as np
import pandas as pd
import dask.dataframe as dd
from .log_utils import Logger


class DateUtils:
    """
    Utility class for date-related operations.

    The DateUtils class provides a variety of operations to manipulate and retrieve
    information about dates, such as calculating week ranges, determining start or
    end dates for specific periods (quarters, months, years), and dynamically
    registering custom time period functions. It also supports parsing specific
    periods for date range computations and ensuring the input date is correctly
    converted to the desired format.

    :ivar logger: Logger instance used for logging messages. Defaults to the logger
                  for the current class if not provided.
    :type logger: Logger

    :ivar _PERIOD_FUNCTIONS: Stores dynamically registered period functions that
                             return start and end dates.
    :type _PERIOD_FUNCTIONS: Dict[str, Callable[[], Tuple[datetime.date, datetime.date]]]
    """
    _PERIOD_FUNCTIONS: Dict[str, Callable[[], Tuple[datetime.date, datetime.date]]] = {}

    def __init__(self, logger=None, debug=False):
        self.logger = logger or Logger.default_logger(logger_name=self.__class__.__name__)
        self.debug = debug

    @classmethod
    def _ensure_date(cls, value: Union[str, datetime.date, datetime.datetime, pd.Timestamp]) -> datetime.date:
        """
        Ensure the input is converted to a datetime.date object.
        """
        if isinstance(value, datetime.date) and not isinstance(value, datetime.datetime):
            return value
        elif isinstance(value, datetime.datetime):
            return value.date()
        elif isinstance(value, pd.Timestamp):
            return value.to_pydatetime().date()
        elif isinstance(value, str):
            for fmt in ('%Y-%m-%d %H:%M:%S', '%Y-%m-%d'):
                try:
                    return datetime.datetime.strptime(value, fmt).date()
                except ValueError:
                    continue
        raise ValueError(f"Unsupported date format: {value}")

    # Public alias to access _ensure_date from other classes
    ensure_date = _ensure_date

    @classmethod
    def calc_week_range(cls, reference_date: Union[str, datetime.date, datetime.datetime, pd.Timestamp]) -> Tuple[
        datetime.date, datetime.date]:
        """
        Calculate the start and end of the week for a given reference date.
        """
        reference_date = cls._ensure_date(reference_date)
        start = reference_date - datetime.timedelta(days=reference_date.weekday())
        end = start + datetime.timedelta(days=6)
        return start, end

    @staticmethod
    def get_year_timerange(year: int) -> Tuple[datetime.date, datetime.date]:
        """
        Get the start and end dates for a given year.
        """
        return datetime.date(year, 1, 1), datetime.date(year, 12, 31)

    @classmethod
    def get_first_day_of_the_quarter(cls, reference_date: Union[
        str, datetime.date, datetime.datetime, pd.Timestamp]) -> datetime.date:
        """
        Get the first day of the quarter for a given date.
        """
        reference_date = cls._ensure_date(reference_date)
        quarter = (reference_date.month - 1) // 3 + 1
        return datetime.date(reference_date.year, 3 * quarter - 2, 1)

    @classmethod
    def get_last_day_of_the_quarter(cls, reference_date: Union[
        str, datetime.date, datetime.datetime, pd.Timestamp]) -> datetime.date:
        """
        Get the last day of the quarter for a given date.
        """
        reference_date = cls._ensure_date(reference_date)
        quarter = (reference_date.month - 1) // 3 + 1
        first_day_of_next_quarter = datetime.date(reference_date.year, 3 * quarter + 1, 1)
        return first_day_of_next_quarter - datetime.timedelta(days=1)

    @classmethod
    def get_month_range(cls, n: int = 0) -> Tuple[datetime.date, datetime.date]:
        """
        Get the date range for the current month or the month `n` months in the past or future.
        """
        today = datetime.date.today()
        target_month = (today.month - 1 + n) % 12 + 1
        target_year = today.year + (today.month - 1 + n) // 12
        start = datetime.date(target_year, target_month, 1)
        if n == 0:
            return start, today
        next_month = (target_month % 12) + 1
        next_year = target_year + (target_month == 12)
        end = datetime.date(next_year, next_month, 1) - datetime.timedelta(days=1)
        return start, end

    @classmethod
    def register_period(cls, name: str, func: Callable[[], Tuple[datetime.date, datetime.date]]):
        """
        Dynamically register a new period function.
        """
        cls._PERIOD_FUNCTIONS[name] = func

    @classmethod
    def parse_period(cls, **kwargs) -> Tuple[datetime.date, datetime.date]:
        """
        Parse the period keyword to determine the start and end date for date range operations.
        """
        period = kwargs.setdefault('period', 'today')
        period_functions = cls._get_default_periods()
        period_functions.update(cls._PERIOD_FUNCTIONS)
        if period not in period_functions:
            raise ValueError(f"Unknown period '{period}'. Available periods: {list(period_functions.keys())}")
        return period_functions[period]()

    @classmethod
    def _get_default_periods(cls) -> Dict[str, Callable[[], Tuple[datetime.date, datetime.date]]]:
        """
        Get default period functions.
        """
        today = datetime.date.today
        return {
            'today': lambda: (today(), today()),
            'yesterday': lambda: (today() - datetime.timedelta(days=1), today() - datetime.timedelta(days=1)),
            'current_week': lambda: cls.calc_week_range(today()),
            'last_week': lambda: cls.calc_week_range(today() - datetime.timedelta(days=7)),
            'current_month': lambda: cls.get_month_range(n=0),
            'last_month': lambda: cls.get_month_range(n=-1),
            'current_year': lambda: cls.get_year_timerange(today().year),
            'current_quarter': lambda: (
                cls.get_first_day_of_the_quarter(today()), cls.get_last_day_of_the_quarter(today())),
            'ytd': lambda: (datetime.date(today().year, 1, 1), today()),
        }


class FileAgeChecker:
    def __init__(self, debug=False, logger=None):
        self.logger = logger or Logger.default_logger(logger_name=self.__class__.__name__)
        self.logger.set_level(Logger.DEBUG if debug else Logger.INFO)
    def is_file_older_than(
            self,
            file_path: str,
            max_age_minutes: int,
            fs: Optional[fsspec.AbstractFileSystem] = None,
            ignore_missing: bool = False,
            verbose: bool = False,
    ) -> bool:
        """
        Check if a file or directory is older than the specified max_age_minutes.

        :param file_path: Path to the file or directory.
        :param max_age_minutes: Maximum allowed age in minutes.
        :param fs: Filesystem object. Defaults to local filesystem.
        :param ignore_missing: Treat missing paths as not old if True.
        :param verbose: Enable detailed logging.
        :return: True if older than max_age_minutes, False otherwise.
        """
        fs = fs or fsspec.filesystem("file")
        self.logger.debug(f"Checking age for {file_path}...")

        try:
            if not fs.exists(file_path):
                self.logger.debug(f"Path not found: {file_path}.")
                return not ignore_missing

            if fs.isdir(file_path):
                self.logger.debug(f"Found directory: {file_path}")
                age = self._get_directory_age_minutes(file_path, fs, verbose)
            elif fs.isfile(file_path):
                age = self._get_file_age_minutes(file_path, fs, verbose)
            else:
                self.logger.warning(f"Path {file_path} is neither file nor directory.")
                return True

            return age > max_age_minutes

        except Exception as e:
            self.logger.warning(f"Error checking {file_path}: {str(e)}")
            return True

    def get_file_or_dir_age_minutes(
            self,
            file_path: str,
            fs: Optional[fsspec.AbstractFileSystem] = None,
    ) -> float:
        """
        Get age of file/directory in minutes. Returns infinity for errors/missing paths.

        :param file_path: Path to check.
        :param fs: Filesystem object. Defaults to local filesystem.
        :return: Age in minutes or infinity if unavailable.
        """
        fs = fs or fsspec.filesystem("file")
        try:
            if not fs.exists(file_path):
                self.logger.debug(f"Path not found: {file_path}")
                return float("inf")

            if fs.isdir(file_path):
                return self._get_directory_age_minutes(file_path, fs, verbose=False)
            if fs.isfile(file_path):
                return self._get_file_age_minutes(file_path, fs, verbose=False)

            self.logger.warning(f"Invalid path type: {file_path}")
            return float("inf")

        except Exception as e:
            self.logger.warning(f"Error getting age for {file_path}: {str(e)}")
            return float("inf")

    def _get_directory_age_minutes(
            self,
            dir_path: str,
            fs: fsspec.AbstractFileSystem,
            verbose: bool,
    ) -> float:
        """Calculate age of oldest file in directory."""
        try:
            all_files = fs.ls(dir_path)
        except Exception as e:
            self.logger.warning(f"Error listing {dir_path}: {str(e)}")
            return float("inf")

        if not all_files:
            self.logger.debug(f"Empty directory: {dir_path}")
            return float("inf")

        modification_times = []
        for file in all_files:
            try:
                info = fs.info(file)
                mod_time = self._get_modification_time(info, file)
                modification_times.append(mod_time)
            except Exception as e:
                self.logger.warning(f"Skipping {file}: {str(e)}")

        if not modification_times:
            self.logger.warning(f"No valid files in {dir_path}")
            return float("inf")

        oldest = min(modification_times)
        age = (datetime.datetime.now(datetime.timezone.utc) - oldest).total_seconds() / 60
        self.logger.debug(f"Oldest in {dir_path}: {age:.2f} minutes")

        return age

    def _get_file_age_minutes(
            self,
            file_path: str,
            fs: fsspec.AbstractFileSystem,
            verbose: bool,
    ) -> float:
        """Calculate file age in minutes."""
        try:
            info = fs.info(file_path)
            mod_time = self._get_modification_time(info, file_path)
            age = (datetime.datetime.now(datetime.timezone.utc) - mod_time).total_seconds() / 60

            if verbose:
                self.logger.debug(f"{file_path} info: {info}")
                self.logger.debug(f"File age: {age:.2f} minutes")

            return age

        except Exception as e:
            self.logger.warning(f"Error processing {file_path}: {str(e)}")
            return float("inf")

    def _get_modification_time(self, info: Dict, file_path: str) -> datetime.datetime:
        """Extract modification time from filesystem info with timezone awareness."""
        try:
            if "LastModified" in info:  # S3-like
                lm = info["LastModified"]
                return lm if isinstance(lm, datetime.datetime) else datetime.datetime.fromisoformat(
                    lm[:-1]).astimezone()

            if "mtime" in info:  # Local filesystem
                return datetime.datetime.fromtimestamp(info["mtime"], tz=datetime.timezone.utc)

            if "modified" in info:  # FTP/SSH
                return datetime.datetime.strptime(
                    info["modified"], "%Y-%m-%d %H:%M:%S"
                ).replace(tzinfo=datetime.timezone.utc)

            raise KeyError("No valid modification time key found")

        except (KeyError, ValueError) as e:
            self.logger.warning(f"Invalid mod time for {file_path}: {str(e)}")
            raise ValueError(f"Unsupported modification time format for {file_path}") from e


# --- Vectorized Helper Functions ---
# These replace the slow, row-by-row .apply() logic. They operate
# on entire DataFrame partitions for maximum efficiency.

def _vectorized_busday_count(
        partition: pd.DataFrame,
        begin_col: str,
        end_col: str,
        holidays: list
) -> pd.Series:
    """Vectorized function to count business days on a DataFrame partition."""
    if partition.empty:
        return pd.Series([], dtype=float)

    # Convert entire columns to datetime at once, coercing errors to NaT
    start_dates = pd.to_datetime(partition[begin_col], errors='coerce').dt.date
    end_dates = pd.to_datetime(partition[end_col], errors='coerce').dt.date

    # Create a result series filled with NaN to handle rows with invalid dates
    result = pd.Series(np.nan, index=partition.index, dtype=float)

    # Create a boolean mask for valid, non-NaT date pairs
    valid_mask = pd.notna(start_dates) & pd.notna(end_dates)

    # Perform the vectorized calculation only on the valid subset of dates
    result.loc[valid_mask] = np.busday_count(
        start_dates[valid_mask],
        end_dates[valid_mask],
        holidays=holidays
    )
    return result


def _vectorized_sla_end_date(
        partition: pd.DataFrame,
        start_col: str,
        n_days_col: str,
        holidays: list
) -> pd.Series:
    """Vectorized function to calculate the SLA end date on a DataFrame partition."""
    if partition.empty:
        return pd.Series([], dtype='datetime64[ns]')

    start_dates = pd.to_datetime(partition[start_col], errors='coerce').dt.date
    sla_days = partition[n_days_col]

    # Create a result series filled with NaT for rows with invalid start dates
    result = pd.Series(pd.NaT, index=partition.index, dtype='datetime64[ns]')

    # Create a boolean mask for valid start dates and sla_days
    valid_mask = pd.notna(start_dates) & pd.notna(sla_days)

    # Perform the vectorized calculation only on the valid subset
    result.loc[valid_mask] = np.busday_offset(
        start_dates[valid_mask],
        sla_days[valid_mask].astype(int),  # Ensure days are integers
        roll='forward',
        holidays=holidays
    )
    return result


# --- Refactored BusinessDays Class ---

class BusinessDays:
    """
    Business days calculations with a custom holiday list.
    Supports scalar and efficient, vectorized Dask DataFrame operations.
    """

    def __init__(self, holiday_list: dict[str, list[str]], logger) -> None:
        self.logger = logger
        self.HOLIDAY_LIST = holiday_list

        # Flatten and store as tuple for determinism
        bd_holidays = [day for year in self.HOLIDAY_LIST for day in self.HOLIDAY_LIST[year]]
        self.holidays = tuple(bd_holidays)

    def get_business_days_count(
            self,
            begin_date: str | datetime.date | pd.Timestamp,
            end_date: str | datetime.date | pd.Timestamp,
    ) -> int:
        """Scalar method to count business days between two dates."""
        begin = pd.to_datetime(begin_date)
        end = pd.to_datetime(end_date)
        return int(np.busday_count(begin.date(), end.date(), holidays=list(self.holidays)))

    def calc_business_days_from_df(
            self,
            df: dd.DataFrame,
            begin_date_col: str,
            end_date_col: str,
            result_col: str = "business_days",
    ) -> dd.DataFrame:
        """Calculates business days between two columns in a Dask DataFrame."""
        missing = {begin_date_col, end_date_col} - set(df.columns)
        if missing:
            self.logger.error(f"Missing columns: {missing}")
            raise ValueError("Required columns are missing from DataFrame")

        return df.assign(
            **{result_col: df.map_partitions(
                _vectorized_busday_count,
                begin_col=begin_date_col,
                end_col=end_date_col,
                holidays=list(self.holidays),
                meta=(result_col, 'f8')  # f8 is float64
            )}
        )

    def add_business_days(
            self,
            start_date: str | datetime.date | pd.Timestamp,
            n_days: int,
    ) -> np.datetime64:
        """Scalar method to add N business days to a start date."""
        start = pd.to_datetime(start_date)
        return np.busday_offset(
            start.date(),
            n_days,
            roll='forward',
            holidays=list(self.holidays),
        )

    def calc_sla_end_date(
            self,
            df: dd.DataFrame,
            start_date_col: str,
            n_days_col: str,
            result_col: str = "sla_end_date",
    ) -> dd.DataFrame:
        """Calculates an SLA end date column for a Dask DataFrame."""
        missing = {start_date_col, n_days_col} - set(df.columns)
        if missing:
            self.logger.error(f"Missing columns: {missing}")
            raise ValueError("Required columns are missing from DataFrame")

        return df.assign(
            **{result_col: df.map_partitions(
                _vectorized_sla_end_date,
                start_col=start_date_col,
                n_days_col=n_days_col,
                holidays=list(self.holidays),
                meta=(result_col, 'datetime64[ns]')
            )}
        )
# Class enhancements
# DateUtils.register_period('next_week', lambda: (datetime.date.today() + datetime.timedelta(days=7),
#                                                 datetime.date.today() + datetime.timedelta(days=13)))
# start, end = DateUtils.parse_period(period='next_week')
# print(f"Next Week: {start} to {end}")
