"""
Core trade date handling module for kepler-pulse.

This module provides the main trade date manipulation functions including
calendar management, date adjustment, and trade day calculations.
"""

from functools import lru_cache

import numpy as np
import pandas as pd
from typing import Union, Optional

from ..config import DEFAULT_BEGIN_DATE, DEFAULT_END_DATE, MAX_STEP_FORWARD, config
from ..utils import dt_handle


class CalendarError(Exception):
    """Custom exception for calendar-related errors."""
    pass


class CalendarBase:
    """
    Base calendar class for managing trade dates.

    This class maintains the original simple interface while adding
    better type hints and error handling for improved maintainability.
    """

    def __init__(self, dates: Optional[pd.Series] = None) -> None:
        """
        Initialize calendar with optional trade dates.

        Args:
            dates: Initial trade dates series. If None, creates empty series.

        Examples:
            >>> calendar = CalendarBase()
            >>> calendar = CalendarBase(pd.Series(['20240101', '20240102']))
        """
        # Initialize with empty Series if no dates provided
        if dates is None:
            self._dates = pd.Series(dtype='str')
        else:
            self._dates = self._format_dates(dates)

    @property
    def dates(self) -> pd.Series:
        """
        Get the trade dates series.

        Returns:
            pandas Series of trade dates in YYYYMMDD format
        """
        return self._dates

    def inject(self, dates: Union[pd.Series, list, tuple]) -> None:
        """
        Inject trade dates into the calendar.

        Args:
            dates: Trade dates to inject (pandas Series, list, or tuple)

        Examples:
            >>> calendar = CalendarBase()
            >>> calendar.inject(['20240101', '20240102', '20240103'])
            >>> calendar.inject(pd.Series(['20240101', '20240102']))
        """
        if isinstance(dates, pd.Series):
            tmp_dates = dates
        else:
            tmp_dates = pd.Series(dates)

        # Convert to datetime and format as YYYYMMDD
        tmp_dates = pd.to_datetime(tmp_dates).dt.strftime('%Y%m%d')
        # Remove duplicates and sort
        tmp_dates = tmp_dates.drop_duplicates().sort_values().reset_index(drop=True)
        self._dates = tmp_dates

    def _format_dates(self, dates: pd.Series) -> pd.Series:
        """
        Helper method to format dates consistently.

        Args:
            dates: Dates to format

        Returns:
            Formatted dates series
        """
        dates = pd.to_datetime(dates).dt.strftime('%Y%m%d')
        return dates.drop_duplicates().sort_values().reset_index(drop=True)

    def is_empty(self) -> bool:
        """Check if calendar is empty."""
        return len(self._dates) == 0


# Global calendar instance for backward compatibility
CALENDAR = CalendarBase()


@lru_cache(maxsize=128)
def get_trade_dts(
    begin_dt: Union[str, int] = DEFAULT_BEGIN_DATE,
    end_dt: Union[str, int] = DEFAULT_END_DATE
) -> pd.Series:
    """
    获取指定时间段的所有交易日.

    Args:
        begin_dt: 开始日期，YYYYMMDD格式或datetime (default: DEFAULT_BEGIN_DATE)
        end_dt: 结束日期，YYYYMMDD格式或datetime (default: DEFAULT_END_DATE)

    Returns:
        pandas Series of trade dates between begin_dt and end_dt

    Examples:
        >>> dates = get_trade_dts('20240101', '20240131')
        >>> print(dates.tolist())
        ['20240102', '20240103', '20240104', ...]

    Raises:
        CalendarError: If calendar is not initialized
    """
    # Check if calendar has dates
    if CALENDAR.is_empty():
        raise CalendarError("Calendar is empty. Please inject trade dates first.")

    # Use cached dates and avoid unnecessary copying
    dates = CALENDAR.dates
    begin_dt, end_dt = dt_handle(begin_dt), dt_handle(end_dt)

    # Validate date range
    if begin_dt > end_dt:
        raise CalendarError(f"Invalid date range: {begin_dt} > {end_dt}")

    # Use boolean indexing for efficient filtering
    mask = (dates >= begin_dt) & (dates <= end_dt)
    result = dates[mask].reset_index(drop=True)
    return result


@lru_cache(maxsize=1024)
def adjust_trade_dt(date: Union[str, int], adjust: str = 'last') -> str:
    """
    调整非交易日到最近的交易日.

    Args:
        date: 要调整的日期，YYYYMMDD格式或datetime
        adjust: 调整方式 - 'last'表示向前一个交易日，'next'表示向后一个交易日 (default: 'last')

    Returns:
        调整后的交易日，YYYYMMDD格式

    Examples:
        >>> adjust_trade_dt('20240106', 'next')  # 周日调整到下周一
        '20240108'
        >>> adjust_trade_dt('20240106', 'last')  # 周日调整到上周五
        '20240105'

    Raises:
        CalendarError: If calendar is not initialized
        ValueError: If adjust parameter is not 'last' or 'next'
    """
    if CALENDAR.is_empty():
        raise CalendarError("Calendar is empty. Please inject trade dates first.")

    dates = CALENDAR.dates
    date = dt_handle(date)

    if adjust == 'last':
        # Use searchsorted for efficient finding
        idx = dates.searchsorted(date, side='right') - 1
        result_idx = max(0, idx)
        if result_idx >= len(dates):
            raise CalendarError(f"No trade date found before {date}")
        return dates.iloc[result_idx]
    elif adjust == 'next':
        # Use searchsorted for efficient finding
        idx = dates.searchsorted(date, side='left')
        result_idx = min(idx, len(dates) - 1)
        if result_idx < 0:
            raise CalendarError(f"No trade date found after {date}")
        return dates.iloc[result_idx]
    else:
        raise ValueError(f"adjust:{adjust} must be 'last' or 'next'!")


@lru_cache(maxsize=1024)
def step_trade_dt(date: Union[str, int], step: int = 1) -> str:
    """
    按指定步长跳转交易日.

    Args:
        date: 起始日期，YYYYMMDD格式或datetime
        step: 步长，正数向前，负数向后 (default: 1)

    Returns:
        结果交易日，YYYYMMDD格式

    Examples:
        >>> step_trade_dt('20240101', 3)  # 向前3个交易日
        '20240104'
        >>> step_trade_dt('20240108', -2)  # 向后2个交易日
        '20240103'

    Raises:
        CalendarError: If calendar is not initialized or step limit exceeded
        ValueError: If step is not an integer
    """
    if CALENDAR.is_empty():
        raise CalendarError("Calendar is empty. Please inject trade dates first.")

    if not isinstance(step, (int, np.integer)):
        raise ValueError(f"Step must be an integer, got {type(step)}")

    if abs(step) > MAX_STEP_FORWARD:
        raise CalendarError(f"Step {step} exceeds maximum limit of {MAX_STEP_FORWARD}")

    dates = CALENDAR.dates
    date = dt_handle(date)

    if step >= 0:
        # Find first date >= input, then step forward
        idx = dates.searchsorted(date, side='left')
        target_idx = idx + step
        result_idx = min(target_idx, len(dates) - 1)
        if result_idx >= len(dates):
            raise CalendarError(f"Cannot step {step} days forward from {date}: no more trade dates")
        return dates.iloc[result_idx]
    else:
        # Find first date <= input, then step backward
        idx = dates.searchsorted(date, side='right') - 1
        target_idx = idx + step  # step is negative
        result_idx = max(target_idx, 0)
        if result_idx < 0:
            raise CalendarError(f"Cannot step {step} days backward from {date}: no more trade dates")
        return dates.iloc[result_idx]


@lru_cache(maxsize=1024)
def delta_trade_dt(begin_dt: Union[str, int], end_dt: Union[str, int]) -> int:
    """
    计算两个日期之间的交易日数量（包含开始和结束日期）.

    Args:
        begin_dt: 开始日期，YYYYMMDD格式或datetime
        end_dt: 结束日期，YYYYMMDD格式或datetime

    Returns:
        交易日数量

    Examples:
        >>> delta_trade_dt('20240101', '20240105')  # 假设都是交易日
        5
        >>> delta_trade_dt('20240101', '20240131')
        22  # 1月份的交易日数量

    Raises:
        CalendarError: If calendar is not initialized
        ValueError: If date range is invalid
    """
    if CALENDAR.is_empty():
        raise CalendarError("Calendar is empty. Please inject trade dates first.")

    dates = CALENDAR.dates
    begin_dt, end_dt = dt_handle(begin_dt), dt_handle(end_dt)

    # Validate date range
    if begin_dt > end_dt:
        raise ValueError(f"Begin date {begin_dt} must be before end date {end_dt}")

    # Use searchsorted for efficient range calculation
    start_idx = dates.searchsorted(begin_dt, side='left')
    end_idx = dates.searchsorted(end_dt, side='right')
    return end_idx - start_idx


if __name__ == "__main__":
    # Example usage
    try:
        # Initialize with some sample dates
        sample_dates = ['20240102', '20240103', '20240104', '20240105',
                       '20240108', '20240109', '20240110', '20240111', '20240112']
        CALENDAR.inject(sample_dates)

        # Test functions
        print("Trade dates from 20240101 to 20240115:")
        trade_dates = get_trade_dts('20240101', '20240115')
        print(trade_dates.tolist())

        print(f"\nAdjust weekend date 20240106:")
        print(f"Next trade date: {adjust_trade_dt('20240106', 'next')}")
        print(f"Last trade date: {adjust_trade_dt('20240106', 'last')}")

        print(f"\nStep operations:")
        print(f"3 days forward from 20240102: {step_trade_dt('20240102', 3)}")
        print(f"2 days backward from 20240112: {step_trade_dt('20240112', -2)}")

        print(f"\nDelta calculation:")
        print(f"Trade days between 20240102 and 20240112: {delta_trade_dt('20240102', '20240112')}")

    except CalendarError as e:
        print(f"Calendar error: {e}")
    except Exception as e:
        print(f"Unexpected error: {e}")