"""
基础API模块
提供通用的数据库查询功能和抽象类
"""

from abc import ABC, abstractmethod
from collections.abc import Iterable
from functools import lru_cache
from typing import Any, Dict, List, Optional, Set, Union

import pandas as pd

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


class BaseAPI(ABC):
    """API基础抽象类

    提供通用的数据库查询功能和接口规范
    """

    def __init__(self, database=None):
        """初始化API实例

        Args:
            database: 数据库连接实例，默认使用全局WIND_DB
        """
        self.db = database or WIND_DB

    @abstractmethod
    def get_table_name(self) -> str:
        """获取表名，子类必须实现

        Returns:
            数据库表名
        """
        pass

    def get_table(self):
        """获取数据库表对象

        Returns:
            数据库表对象
        """
        table_name = self.get_table_name()
        return getattr(self.db, table_name.upper())

    def create_query(self, *columns):
        """创建基础查询

        Args:
            *columns: 要查询的列名，可以是字符串列名或SQLAlchemy列对象

        Returns:
            查询对象
        """
        table = self.get_table()
        if columns:
            # 智能处理不同类型的列参数
            processed_columns = []
            for col in columns:
                if isinstance(col, str):
                    # 字符串列名：使用getattr获取表属性
                    processed_columns.append(getattr(table, col))
                else:
                    # SQLAlchemy对象（如label对象）：直接使用
                    processed_columns.append(col)
            return self.db.query(*processed_columns)
        return self.db.query(table)

    def filter_by_codes(self, query, code_column: str, codes: Union[str, List[str], Set[str]]):
        """根据代码过滤查询结果

        Args:
            query: 查询对象
            code_column: 代码列名
            codes: 股票/指数代码

        Returns:
            过滤后的查询对象
        """
        table = self.get_table()
        code_attr = getattr(table, code_column)

        if isinstance(codes, str):
            return query.filter(code_attr == codes)
        elif isinstance(codes, Iterable):
            return query.filter(code_attr.in_(codes))
        else:
            raise ValueError(f"codes must be str or Iterable, got {type(codes)}")

    def filter_by_date_range(self, query, date_column: str, begin_dt: str, end_dt: str):
        """根据日期范围过滤查询结果

        Args:
            query: 查询对象
            date_column: 日期列名
            begin_dt: 开始日期
            end_dt: 结束日期

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

    def filter_current(self, query, current_column: str = "CUR_SIGN"):
        """只查询当前有效的记录

        Args:
            query: 查询对象
            current_column: 当前有效标识列名

        Returns:
            过滤后的查询对象
        """
        table = self.get_table()
        current_attr = getattr(table, current_column)
        return query.filter(current_attr == "1")

    def to_dataframe(self, query, columns: Optional[List[str]] = None) -> pd.DataFrame:
        """将查询结果转换为DataFrame

        Args:
            query: 查询对象
            columns: 列名列表，用于重命名

        Returns:
            查询结果的DataFrame
        """
        df = query.to_df()
        if columns and len(columns) == len(df.columns):
            df.columns = columns
        return df

    def get_unique_codes(self, code_column: str = "S_CON_WINDCODE",
                        current_only: bool = True) -> Set[str]:
        """获取唯一的代码集合

        Args:
            code_column: 代码列名
            current_only: 是否只获取当前有效的记录

        Returns:
            代码集合
        """
        query = self.create_query(code_column)
        if current_only:
            query = self.filter_current(query)

        df = self.to_dataframe(query, ["sid"])
        return set(df.sid)


class UniverseAPI(BaseAPI):
    """股票池API基础类

    提供股票池查询的通用功能
    """

    def get_universe_set(self, index_code: str, cur_sign: bool = True) -> Set[str]:
        """获取指数成分股集合

        Args:
            index_code: 指数代码
            cur_sign: 是否只获取当前有效的成分股

        Returns:
            股票代码集合
        """
        table = self.get_table()
        query = self.create_query(table.S_CON_WINDCODE).filter(
            table.S_INFO_WINDCODE == index_code
        )

        if cur_sign:
            query = self.filter_current(query)

        df = self.to_dataframe(query, ["sid"])
        return set(df.sid)

    def get_universe_dataframe(self, index_code: str) -> pd.DataFrame:
        """获取指数成分股的完整信息（包含历史变更）

        Args:
            index_code: 指数代码

        Returns:
            包含进入和退出日期的DataFrame
        """
        table = self.get_table()

        # 优雅地使用改进后的create_query方法
        query = self.create_query(
            table.S_CON_WINDCODE.label("sid"),
            table.S_CON_INDATE.label("entry_dt"),
            table.S_CON_OUTDATE.label("out_dt")
        ).filter(table.S_INFO_WINDCODE == index_code)

        df = query.to_df()
        return df


class IndexDataAPI(BaseAPI):
    """指数数据API基础类

    提供指数OHLCV数据查询的通用功能
    """

    def get_ohlcv_data(self, index_list: Union[str, List[str], frozenset],
                      begin_dt: str = "20010101", end_dt: str = "20990101",
                      index_type: str = "normal") -> pd.DataFrame:
        """获取指数OHLCV数据

        Args:
            index_list: 指数代码列表
            begin_dt: 开始日期
            end_dt: 结束日期
            index_type: 指数类型 (normal, zx, wind, sw)

        Returns:
            包含OHLCV数据的DataFrame
        """
        table = self.get_table()

        # 直接构建查询，不使用create_query方法
        query = self.db.query(
            table.S_INFO_WINDCODE,
            table.TRADE_DT,
            table.S_DQ_OPEN,
            table.S_DQ_HIGH,
            table.S_DQ_LOW,
            table.S_DQ_CLOSE,
            table.S_DQ_VOLUME,
            table.S_DQ_AMOUNT,
        ).filter(
            table.TRADE_DT >= begin_dt,
            table.TRADE_DT <= end_dt
        ).order_by(table.TRADE_DT, table.S_INFO_WINDCODE)

        if isinstance(index_list, str):
            query = query.filter(table.S_INFO_WINDCODE == index_list)
        elif isinstance(index_list, (list, set, frozenset)):
            query = query.filter(table.S_INFO_WINDCODE.in_(list(index_list)))

        # 直接构建DataFrame，不使用to_dataframe方法
        df = query.to_df()
        df.columns = [
            "sid", "trade_dt", "s_open", "s_high", "s_low",
            "s_close", "s_volume", "s_amount"
        ]
        return df

    @lru_cache(maxsize=128)
    def get_cached_ohlcv_data(self, *args, **kwargs) -> pd.DataFrame:
        """获取缓存的指数OHLCV数据

        使用LRU缓存提高重复查询的性能

        Args:
            *args, **kwargs: 传递给get_ohlcv_data的参数

        Returns:
            包含OHLCV数据的DataFrame
        """
        return self.get_ohlcv_data(*args, **kwargs)


class DynamicUniverseBase:
    """动态股票池基类

    提供历史任意时间点的股票池查询功能
    """

    def __init__(self, database=None):
        """初始化动态股票池

        Args:
            database: 数据库连接实例
        """
        self.db = database or WIND_DB

    @abstractmethod
    def get_all_members_data(self) -> pd.DataFrame:
        """获取所有成员的历史数据，子类必须实现

        Returns:
            包含sid, entry_dt, out_dt的DataFrame
        """
        pass

    def preview(self, trade_dt: str) -> Set[str]:
        """获取指定日期的股票池

        Args:
            trade_dt: 交易日期

        Returns:
            股票代码集合
        """
        df = self.get_all_members_data()
        trade_dt = dt_handle(trade_dt)

        # 过滤出在指定日期有效的股票
        valid_stocks = df.loc[
            (df["entry_dt"] <= trade_dt) &
            ((df["out_dt"] >= trade_dt) | (df["out_dt"].isnull()))
        ]

        return set(valid_stocks.sid)

    def preview_with_data(self, trade_dt: str) -> pd.DataFrame:
        """获取指定日期的股票池及其详细信息

        Args:
            trade_dt: 交易日期

        Returns:
            包含股票池信息的DataFrame
        """
        df = self.get_all_members_data()
        trade_dt = dt_handle(trade_dt)

        # 过滤出在指定日期有效的股票
        valid_stocks = df.loc[
            (df["entry_dt"] <= trade_dt) &
            ((df["out_dt"] >= trade_dt) | (df["out_dt"].isnull()))
        ].copy()

        # 添加查询日期信息
        valid_stocks['query_dt'] = trade_dt

        return valid_stocks


# 工具函数
def validate_date_format(date_str: str) -> str:
    """验证并标准化日期格式

    Args:
        date_str: 日期字符串

    Returns:
        标准化的日期字符串

    Raises:
        ValueError: 日期格式无效
    """
    try:
        return dt_handle(date_str)
    except Exception as e:
        raise ValueError(f"Invalid date format: {date_str}, error: {e}")


def validate_index_codes(index_list: Union[str, List[str], Set[str]]) -> Set[str]:
    """验证并标准化指数代码列表

    Args:
        index_list: 指数代码或代码列表

    Returns:
        标准化的指数代码集合
    """
    if isinstance(index_list, str):
        return {index_list}
    elif isinstance(index_list, Iterable):
        return set(index_list)
    else:
        raise ValueError("index_list must be str or Iterable")


__all__ = [
    'BaseAPI',
    'UniverseAPI',
    'IndexDataAPI',
    'DynamicUniverseBase',
    'validate_date_format',
    'validate_index_codes'
]