"""
收益率数据API基础模块

提供统一的收益率数据查询和处理基础类
"""

from abc import ABC, abstractmethod
from functools import lru_cache
from typing import Optional, Set, Union, List, Dict, Any
import pandas as pd
from sqlalchemy import func

from sangreal_wind.utils.base_api import BaseAPI
from sangreal_wind.utils.config_manager import WIND_DB
from sangreal_wind.utils.datetime_handle import dt_handle


class ReturnDataAPI(BaseAPI):
    """收益率数据API基础类"""

    # 市量表名映射
    UNIVERSE_TABLES = {
        "A": "AShareEODPrices",
        "HK": "HKShareEODPrices"
    }

    # 月度收益率表名
    MONTHLY_TABLE = "ASHAREMONTHLYYIELD"

    def __init__(self, database=None):
        """初始化收益率API

        Args:
            database: 数据库连接实例，默认使用全局WIND_DB
        """
        super().__init__(database or WIND_DB)

    def get_table_by_universe(self, universe: str):
        """根据universe获取表名

        Args:
            universe: 市场类型 ('A' 或 'HK')

        Returns:
            数据库表对象

        Raises:
            ValueError: 不支持的universe类型
        """
        table_name = self.UNIVERSE_TABLES.get(universe)
        if table_name is None:
            raise ValueError(f"universe only support {list(self.UNIVERSE_TABLES.keys())}!")
        return getattr(self.db, table_name.upper())

    def create_return_query(self, table, columns):
        """创建收益率查询

        Args:
            table: 数据库表对象
            columns: 要查询的列

        Returns:
            查询对象
        """
        return self.db.query(*columns)

    def filter_by_stocks(self, query, table, sid: Optional[Union[str, List[str], Set[str]]]):
        """统一的股票代码过滤逻辑

        Args:
            query: 查询对象
            table: 数据库表对象
            sid: 股票代码

        Returns:
            过滤后的查询对象
        """
        if sid is not None:
            if isinstance(sid, str):
                return query.filter(table.S_INFO_WINDCODE == sid)
            else:
                return query.filter(table.S_INFO_WINDCODE.in_(sid))
        return query

    def create_date_filter(self, query, table, begin_dt: str, end_dt: str):
        """创建日期范围过滤器

        Args:
            query: 查询对象
            table: 数据库表对象
            begin_dt: 开始日期
            end_dt: 结束日期

        Returns:
            过滤后的查询对象
        """
        return query.filter(
            table.TRADE_DT >= begin_dt, table.TRADE_DT <= end_dt
        )

    def execute_query(self, query, order_by_col=None):
        """执行查询并返回DataFrame

        Args:
            query: 查询对象
            order_by_col: 排序列名

        Returns:
            查询结果的DataFrame
        """
        if order_by_col:
            query = query.order_by(order_by_col)
        return query.to_df()

    def filter_and_execute(self, query, table, sid=None, begin_dt="20010101",
                           end_dt="20990101", order_by_col=None):
        """过滤并执行查询

        Args:
            query: 查询对象
            table: 数据库表对象
            sid: 股票代码
            begin_dt: 开始日期
            end_dt: 结束日期
            order_by_col: 排序列名

        Returns:
            查询结果的DataFrame
        """
        query = self.filter_by_stocks(query, table, sid)
        query = self.create_date_filter(query, table, begin_dt, end_dt)
        return self.execute_query(query, order_by_col)

    @abstractmethod
    def get_return_data(self, *args, **kwargs) -> pd.DataFrame:
        """获取收益率数据，子类必须实现

        Returns:
            收益率DataFrame
        """
        pass

    def process_return_dataframe(self, df, column_mapping: Dict[str, str] = None):
        """处理收益率DataFrame

        Args:
            df: 原始DataFrame
            column_mapping: 列名映射

        Returns:
            处理后的DataFrame
        """
        if column_mapping:
            df = df.rename(columns=column_mapping)
        return df


class DailyReturnAPI(ReturnDataAPI):
    """日收益率API"""

    def get_table_name(self) -> str:
        """获取日收益率表名"""
        return "daily_return"

    @lru_cache(maxsize=128)
    def get_return_data(self, sid: Optional[Union[str, List[str], Set[str]]] = None,
                        begin_dt: str = "20030101", end_dt: str = "20990101",
                        universe: str = "A") -> pd.DataFrame:
        """获取日收益率数据

        Args:
            sid: 股票代码
            begin_dt: 开始日期
            end_dt: 结束日期
            universe: 市场类型

        Returns:
            日收益率DataFrame
        """
        table = self.get_table_by_universe(universe)
        columns = [
            table.S_INFO_WINDCODE,
            table.TRADE_DT,
            table.S_DQ_ADJCLOSE
        ]

        query = self.create_return_query(table, columns)
        return self.filter_and_execute(query, table, sid, begin_dt, end_dt, table.TRADE_DT)


class MonthlyReturnAPI(ReturnDataAPI):
    """月收益率API"""

    def get_table_name(self) -> str:
        """获取月收益率表名"""
        return "monthly_return"

    @lru_cache(maxsize=128)
    def get_return_data(self, sid: Optional[Union[str, List[str], Set[str]]] = None,
                        begin_dt: str = "20030101", end_dt: str = "20990101",
                        trade_dt: Optional[str] = None) -> pd.DataFrame:
        """获取月收益率数据

        Args:
            sid: 股票代码
            begin_dt: 开始日期
            end_dt: 结束日期
            trade_dt: 特定交易日期

        Returns:
            月收益率DataFrame
        """
        table = getattr(self.db, self.MONTHLY_TABLE.upper())
        columns = [
            table.S_INFO_WINDCODE,
            table.TRADE_DT,
            table.S_MQ_PCTCHANGE
        ]

        query = self.create_return_query(table, columns)
        query = self.filter_by_stocks(query, table, sid)

        if trade_dt is not None:
            trade_dt = dt_handle(trade_dt)
            df = query.filter(table.TRADE_DT == trade_dt).order_by(table.TRADE_DT).to_df()
            column_mapping = {"S_MQ_PCTCHANGE": "close"}
        else:
            df = query.filter(
                table.TRADE_DT >= begin_dt,
                table.TRADE_DT <= end_dt
            ).order_by(table.TRADE_DT).to_df()
            column_mapping = {"S_MQ_PCTCHANGE": "close"}

        df = self.process_return_dataframe(df, column_mapping)
        df.close = df.close / 100.0

        if trade_dt is not None:
            return df
        else:
            # 对于日期范围查询，需要透视数据
            df = df.pivot(values="close", index="trade_dt", columns="sid")
            df.dropna(how="all", inplace=True)
            return df.T


# 创建全局API实例
_daily_return_api = DailyReturnAPI()
_monthly_return_api = MonthlyReturnAPI()

__all__ = [
    'ReturnDataAPI',
    'DailyReturnAPI',
    'MonthlyReturnAPI',
    '_daily_return_api',
    '_monthly_return_api'
]