"""
Trading Calendar Module.
"""

import datetime as dt
from typing import Union, Literal

import numpy as np
import pandas as pd
from pandas.tseries.offsets import WeekOfMonth

from ..config import (
    DEFAULT_BEGIN_DATE, DEFAULT_END_DATE, MAX_STEP_FORWARD,
    DAILY_SEARCH_STEPS, WEEKLY_OFFSET, MONTHLY_OFFSET,
    QUARTERLY_OFFSET, YEARLY_OFFSET, YEARS_FORWARD,
    CONTRACT_SWITCH_DAYS, FUTURES_START_DATE, TREASURY_START_DATE
)
from ..utils import validate_date


class CalendarError(Exception):
    """Calendar exception."""
    pass


class FuturesAccessor:
    """
    期货访问器.

    支持:
        - cal.futures.stock('20240115')      # 股指交割日
        - cal.futures.stock(all=True)        # 所有股指交割日
        - cal.futures.treasury()             # 国债交割日
        - cal.futures.contracts('20240115')  # 主力合约
    """

    def __init__(self, calendar: 'Calendar'):
        self._cal = calendar

    def stock(self, date: str = None, all: bool = False) -> Union[pd.Series, list]:
        """股指期货交割日."""
        if all:
            return self._cal._stock_delivery_dates_all()
        return self._cal._stock_delivery_dates(date)

    def treasury(self) -> list:
        """国债期货交割日."""
        return self._cal._treasury_delivery_dates()

    def contracts(self, date: str) -> pd.Series:
        """主力合约（股指期货）."""
        return self._cal._main_contracts(date)


class BoundaryAccessor:
    """
    边界访问器（first/last）.

    支持:
        - cal.monthly.first('20240101', '20241231')  # 获取月度首日序列
        - cal.monthly.first.next('20240115')         # 下一个月度首日
        - cal.monthly.first.prev('20240115')         # 上一个月度首日
    """

    def __init__(self, freq_accessor: 'FrequencyAccessor', keep: str):
        self._fa = freq_accessor
        self._keep = keep

    def __call__(self, begin=DEFAULT_BEGIN_DATE, end=DEFAULT_END_DATE) -> pd.Series:
        """获取边界日期序列."""
        return self._fa._cal._get_freq_dates(self._fa._freq, begin, end, self._keep)

    def next(self, date: Union[str, int], step: int = 1) -> str:
        """下一个边界日期."""
        return self._fa._cal._next_freq_date(self._fa._freq, date, step, self._keep)

    def prev(self, date: Union[str, int], step: int = 1) -> str:
        """上一个边界日期."""
        return self._fa._cal._prev_freq_date(self._fa._freq, date, step, self._keep)


class FrequencyAccessor:
    """
    频率访问器.

    支持:
        - cal.monthly.first('20240101', '20241231')   # 获取月度首日序列
        - cal.monthly.first.next('20240115')          # 下一个月度首日
        - cal.monthly.last('20240101', '20241231')    # 获取月末序列
        - cal.monthly.last.next('20240115')           # 下一个月末
        - cal.monthly.both('20240101', '20241231')    # 月初+月末
    """

    def __init__(self, calendar: 'Calendar', freq: str):
        self._cal = calendar
        self._freq = freq
        self._first = BoundaryAccessor(self, 'first')
        self._last = BoundaryAccessor(self, 'last')

    @property
    def first(self) -> BoundaryAccessor:
        """周期首日."""
        return self._first

    @property
    def last(self) -> BoundaryAccessor:
        """周期末日."""
        return self._last

    def both(self, begin=DEFAULT_BEGIN_DATE, end=DEFAULT_END_DATE) -> pd.Series:
        """每个周期的第一个和最后一个交易日."""
        first_dates = self._cal._get_freq_dates(self._freq, begin, end, 'first')
        last_dates = self._cal._get_freq_dates(self._freq, begin, end, 'last')
        combined = pd.concat([first_dates, last_dates])
        return combined.drop_duplicates().sort_values().reset_index(drop=True)


class Calendar:
    """
    交易日历类.

    Examples:
        >>> cal = Calendar(trade_dates)
        >>> cal.range('20240101', '20240131')
        >>> cal.monthly.first('20240101', '20241231')
        >>> cal.monthly.first.next('20240115')
    """

    def __init__(self, trade_dates: Union[pd.Series, list, tuple, None] = None) -> None:
        """
        初始化日历.

        Args:
            trade_dates: 交易日数据，支持 list/tuple/pd.Series
        """
        self._cache: dict = {}
        self._freq_accessors: dict = {}
        if trade_dates is None:
            self._dates = pd.Series(dtype='str')
        else:
            self._dates = self._format_dates(trade_dates)

    @property
    def dates(self) -> pd.Series:
        """交易日序列."""
        return self._dates

    def inject(self, trade_dates: Union[pd.Series, list, tuple]) -> 'Calendar':
        """注入交易日数据."""
        if isinstance(trade_dates, pd.Series):
            tmp_dates = trade_dates
        else:
            tmp_dates = pd.Series(trade_dates)

        tmp_dates = pd.to_datetime(tmp_dates).dt.strftime('%Y%m%d')
        tmp_dates = tmp_dates.drop_duplicates().sort_values().reset_index(drop=True)
        self._dates = tmp_dates
        self._cache.clear()
        return self

    def is_empty(self) -> bool:
        return len(self._dates) == 0

    def _format_dates(self, trade_dates) -> pd.Series:
        dates = pd.Series(pd.to_datetime(trade_dates))
        dates = dates.dt.strftime('%Y%m%d')
        return dates.drop_duplicates().sort_values().reset_index(drop=True)

    def _cached(self, key: tuple, compute_fn):
        if key not in self._cache:
            self._cache[key] = compute_fn()
        return self._cache[key]

    def _get_freq_accessor(self, freq: str) -> FrequencyAccessor:
        if freq not in self._freq_accessors:
            self._freq_accessors[freq] = FrequencyAccessor(self, freq)
        return self._freq_accessors[freq]

    @property
    def futures(self) -> FuturesAccessor:
        """期货访问器."""
        if not hasattr(self, '_futures_accessor'):
            self._futures_accessor = FuturesAccessor(self)
        return self._futures_accessor

    # ============ 基础操作 ============

    def range(
        self,
        begin: Union[str, int] = DEFAULT_BEGIN_DATE,
        end: Union[str, int] = DEFAULT_END_DATE
    ) -> pd.Series:
        """获取时间范围内的交易日."""
        if self.is_empty():
            raise CalendarError("Calendar is empty.")

        begin, end = validate_date(begin), validate_date(end)
        key = ('range', begin, end)
        return self._cached(key, lambda: self._compute_range(begin, end))

    def _compute_range(self, begin: str, end: str) -> pd.Series:
        if begin > end:
            raise CalendarError(f"Invalid range: {begin} > {end}")
        mask = (self._dates >= begin) & (self._dates <= end)
        return self._dates[mask].reset_index(drop=True)

    def adjust(self, date: Union[str, int], direction: Literal['prev', 'next'] = 'prev') -> str:
        """
        调整到最近的交易日.

        Args:
            date: 日期
            direction: 'prev' 向前找, 'next' 向后找
        """
        if self.is_empty():
            raise CalendarError("Calendar is empty.")

        date = validate_date(date)
        key = ('adjust', date, direction)
        return self._cached(key, lambda: self._compute_adjust(date, direction))

    def _compute_adjust(self, date: str, direction: str) -> str:
        if direction == 'prev':
            idx = self._dates.searchsorted(date, side='right') - 1
            result_idx = max(0, idx)
            if result_idx >= len(self._dates):
                raise CalendarError(f"No trade date before {date}")
            return self._dates.iloc[result_idx]
        elif direction == 'next':
            idx = self._dates.searchsorted(date, side='left')
            result_idx = min(idx, len(self._dates) - 1)
            if result_idx < 0:
                raise CalendarError(f"No trade date after {date}")
            return self._dates.iloc[result_idx]
        else:
            raise ValueError("direction must be 'prev' or 'next'")

    def step(self, date: Union[str, int], n: int = 1) -> str:
        """
        按步长跳转.

        Args:
            date: 起始日期
            n: 步长（正数向前，负数向后）
        """
        if self.is_empty():
            raise CalendarError("Calendar is empty.")

        if not isinstance(n, (int, np.integer)):
            raise ValueError("n must be an integer")

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

        date = validate_date(date)
        key = ('step', date, n)
        return self._cached(key, lambda: self._compute_step(date, n))

    def _compute_step(self, date: str, n: int) -> str:
        if n >= 0:
            idx = self._dates.searchsorted(date, side='left')
            target_idx = min(idx + n, len(self._dates) - 1)
            return self._dates.iloc[target_idx]
        else:
            idx = self._dates.searchsorted(date, side='right') - 1
            target_idx = max(idx + n, 0)
            return self._dates.iloc[target_idx]

    def delta(self, begin: Union[str, int], end: Union[str, int]) -> int:
        """计算两日之间的交易日数量."""
        if self.is_empty():
            raise CalendarError("Calendar is empty.")

        begin, end = validate_date(begin), validate_date(end)
        key = ('delta', begin, end)
        return self._cached(key, lambda: self._compute_delta(begin, end))

    def _compute_delta(self, begin: str, end: str) -> int:
        if begin > end:
            raise ValueError("begin must be before end")
        start_idx = self._dates.searchsorted(begin, side='left')
        end_idx = self._dates.searchsorted(end, side='right')
        return end_idx - start_idx

    # ============ 频率访问器 ============

    @property
    def daily(self) -> FrequencyAccessor:
        """日频."""
        return self._get_freq_accessor('daily')

    @property
    def weekly(self) -> FrequencyAccessor:
        """周频."""
        return self._get_freq_accessor('weekly')

    @property
    def biweekly(self) -> FrequencyAccessor:
        """双周频."""
        return self._get_freq_accessor('biweekly')

    @property
    def monthly(self) -> FrequencyAccessor:
        """月频."""
        return self._get_freq_accessor('monthly')

    @property
    def quarterly(self) -> FrequencyAccessor:
        """季频."""
        return self._get_freq_accessor('quarterly')

    @property
    def yearly(self) -> FrequencyAccessor:
        """年频."""
        return self._get_freq_accessor('yearly')

    @property
    def halfyearly(self) -> FrequencyAccessor:
        """半年频."""
        return self._get_freq_accessor('halfyearly')

    @property
    def reportly(self) -> FrequencyAccessor:
        """财报频."""
        return self._get_freq_accessor('reportly')

    # ============ 频率内部实现 ============

    FREQ_CONFIG = {
        'daily': {'offset': 0},
        'weekly': {'offset': WEEKLY_OFFSET, 'period': lambda x: f"{x.isocalendar()[0]}{x.isocalendar()[1]}"},
        'biweekly': {'offset': WEEKLY_OFFSET, 'period': lambda x: f"{x.isocalendar()[0]}{x.isocalendar()[1]}"},
        'monthly': {'offset': MONTHLY_OFFSET, 'period': lambda x: f"{x.year}{x.month}"},
        'quarterly': {'offset': QUARTERLY_OFFSET, 'period': lambda x: f"{x.year}{x.quarter}"},
        'yearly': {'offset': YEARLY_OFFSET, 'period': lambda x: f"{x.year}"},
        'halfyearly': {'offset': MONTHLY_OFFSET},
        'reportly': {'offset': QUARTERLY_OFFSET},
    }

    def _get_freq_dates(self, freq: str, begin, end, keep) -> pd.Series:
        """获取频率日期."""
        if freq == 'daily':
            return self.range(begin, end).copy()

        begin, end = validate_date(begin), validate_date(end)
        config = self.FREQ_CONFIG[freq]
        offset = config['offset']

        if freq == 'halfyearly':
            return self._get_halfyearly(begin, end, keep)
        elif freq == 'reportly':
            return self._get_reportly(begin, end, keep)
        elif freq == 'biweekly':
            dates = self._get_frequency(begin, end, config['period'], offset, keep)
            return dates[::2].reset_index(drop=True)
        else:
            return self._get_frequency(begin, end, config['period'], offset, keep)

    def _get_frequency(self, begin, end, period_func, offset, keep):
        begin, end = validate_date(begin), validate_date(end)
        df = self.range(
            self.step(begin, -offset),
            self.step(end, offset)
        ).to_frame('trade_dt')
        df['_dt'] = pd.to_datetime(df['trade_dt'])
        df['period'] = df['_dt'].map(period_func)
        result = df.drop_duplicates('period', keep=keep)['trade_dt']
        result = result[(result >= begin) & (result <= end)].reset_index(drop=True)
        return result

    def _get_halfyearly(self, begin, end, keep):
        def to_period(x):
            date_int = int(x)
            if date_int <= int(x[:4] + '0630'):
                return x[:4] + '01'
            return x[:4] + '07'

        df = self.range(
            self.step(begin, -MONTHLY_OFFSET),
            self.step(end, MONTHLY_OFFSET)
        ).to_frame('trade_dt')
        df['period'] = df['trade_dt'].map(to_period)
        result = df.drop_duplicates('period', keep=keep)['trade_dt']
        result = result[(result >= begin) & (result <= end)].sort_values().reset_index(drop=True)
        return result

    def _get_reportly(self, begin, end, keep):
        def to_period(x):
            date_int = int(x)
            if date_int <= int(x[:4] + '0430'):
                return str(int(x[:4]) - 1) + '11'
            elif date_int <= int(x[:4] + '0831'):
                return x[:4] + '05'
            elif date_int <= int(x[:4] + '1031'):
                return x[:4] + '09'
            return x[:4] + '11'

        df = self.range(
            self.step(begin, -QUARTERLY_OFFSET),
            self.step(end, QUARTERLY_OFFSET)
        ).to_frame('trade_dt')
        df['period'] = df['trade_dt'].map(to_period)
        result = df.drop_duplicates('period', keep=keep)['trade_dt']

        if keep == 'first':
            result = result.map(lambda x: self.adjust(x[:6] + '01', 'next'))
        else:
            result = result.map(lambda x: self.step(x[:6] + '01', -1))

        result = result[(result >= begin) & (result <= end)].sort_values().reset_index(drop=True)
        return result

    def _next_freq_date(self, freq: str, date: Union[str, int], step: int, keep: str) -> str:
        """下一个频率日期."""
        if not isinstance(step, (int, np.integer)) or step < 1:
            raise CalendarError("step must be positive integer")

        date = validate_date(date)
        end = self.step(date, DAILY_SEARCH_STEPS)
        df = self._get_freq_dates(freq, date, end, keep).tolist()

        try:
            if df and df[0] == date and len(df) > step:
                return df[step]
            if df and len(df) >= step:
                return df[step - 1]
            if df:
                return df[-1]
            raise CalendarError(f"No {freq} dates after {date}")
        except IndexError:
            if df:
                return df[-1]
            raise CalendarError(f"Cannot get next {freq} from {date}")

    def _prev_freq_date(self, freq: str, date: Union[str, int], step: int, keep: str) -> str:
        """上一个频率日期."""
        if not isinstance(step, (int, np.integer)) or step < 1:
            raise CalendarError("step must be positive integer")

        date = validate_date(date)
        begin = self.step(date, -DAILY_SEARCH_STEPS)
        df = self._get_freq_dates(freq, begin, date, keep).tolist()

        try:
            if df and df[-1] == date and len(df) > step:
                return df[-step - 1]
            if df and len(df) >= step:
                return df[-step]
            if df:
                return df[0]
            raise CalendarError(f"No {freq} dates before {date}")
        except IndexError:
            if df:
                return df[0]
            raise CalendarError(f"Cannot get prev {freq} from {date}")

    # ============ 期货内部实现 ============

    def _stock_delivery_dates_all(self) -> list:
        today = dt.datetime.today().strftime('%Y%m%d')
        dates = self.monthly.first(FUTURES_START_DATE, f'{int(today[:4]) + YEARS_FORWARD}0101')
        dates = pd.to_datetime(dates).dt.date + WeekOfMonth(1, week=2, weekday=4)
        dates = dates[dates < pd.to_datetime(self._dates.iloc[-1])].copy()
        dates = dates.map(lambda x: self.adjust(x, 'next'))
        dates = sorted([d for d in dates if d <= self.monthly.first.next(today, 2)])
        return dates

    def _stock_delivery_dates(self, date: str) -> pd.Series:
        dates = self.monthly.first(FUTURES_START_DATE, self.monthly.first.next(date))
        dates = pd.to_datetime(dates).dt.date + WeekOfMonth(1, week=2, weekday=4)
        dates = dates.map(lambda x: self.adjust(x, 'next'))
        if dates.iloc[-2] >= date:
            dates = dates.iloc[:-1].copy()
        return dates

    def _treasury_delivery_dates(self) -> list:
        today = dt.datetime.today().strftime('%Y%m%d')
        dates = self.quarterly.last(TREASURY_START_DATE, f'{int(today[:4]) + YEARS_FORWARD}0101')
        dates = dates.map(lambda x: self.step(self.monthly.first.prev(x), -2))
        dates = dates[dates < self._dates.iloc[-1]].copy()
        dates = sorted([d for d in dates if d <= self.quarterly.first.next(today, 2)])
        return dates

    def _main_contracts(self, date: str) -> pd.Series:
        delivery = self._stock_delivery_dates(date).iloc[-1]
        c0 = delivery[2:6]
        c1 = self.monthly.first.next(delivery)[2:6]

        if date == delivery:
            return pd.Series([c1])
        elif date > self.step(delivery, -CONTRACT_SWITCH_DAYS):
            return pd.Series([c0, c1])
        else:
            return pd.Series([c0])
