from functools import lru_cache

import attr
import pandas as pd

from sangreal_wind.utils.base_api import UniverseAPI, DynamicUniverseBase
from sangreal_wind.utils.commons import INDEX_DICT
from sangreal_wind.utils.datetime_handle import dt_handle
from sangreal_wind.utils.config_manager import WIND_DB

indx_error = f"请输入正确的指数简称，如{list(INDEX_DICT.keys())}，或指数wind代码！"


class AStockUniverseAPI(UniverseAPI):
    """A股股票池API"""

    def get_table_name(self) -> str:
        return "AINDEXMEMBERSWIND"

    def get_a_stocks(self, cur_sign: bool = True) -> set:
        """获取全A股成份股

        Args:
            cur_sign: 是否需要最新的股票池

        Returns:
            股票代码集合
        """
        return self.get_universe_set("881001.WI", cur_sign)

    def get_universe_dataframe(self, index_code: str = "881001.WI") -> pd.DataFrame:
        """获取全A股成分股的完整信息

        Args:
            index_code: 指数代码

        Returns:
            包含进入和退出日期的DataFrame
        """
        table = self.get_table()
        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.F_INFO_WINDCODE == index_code)
        df = query.to_df()
        return df


class NormalIndexUniverseAPI(UniverseAPI):
    """普通指数股票池API"""

    def get_table_name(self) -> str:
        return "AINDEXMEMBERS"

    def get_index_stocks(self, index_code: str, cur_sign: bool = True) -> set:
        """获取普通指数成分股

        Args:
            index_code: 指数代码
            cur_sign: 是否需要最新的股票池

        Returns:
            股票代码集合
        """
        return self.get_universe_set(index_code, cur_sign)


class MSCIUniverseAPI(UniverseAPI):
    """MSCI指数股票池API"""

    def get_table_name(self) -> str:
        return "ASHAREMSCIMEMBERS"

    def get_msci_stocks(self, cur_sign: bool = True) -> set:
        """获取MSCI成分股

        Args:
            cur_sign: 是否需要最新的股票池

        Returns:
            股票代码集合
        """
        table = self.get_table()

        # 直接构建查询，避免create_query的问题
        query = self.db.query(table.S_INFO_WINDCODE)

        if cur_sign:
            query = query.filter(table.CUR_SIGN == "1")

        df = query.to_df()
        df.columns = ["sid"]
        return set(df.sid)

    def get_universe_dataframe(self, index_code: str = "") -> pd.DataFrame:
        """获取MSCI成分股的完整信息

        Args:
            index_code: 指数代码（MSCI不需要）

        Returns:
            包含进入和退出日期的DataFrame
        """
        table = self.get_table()
        query = self.create_query(
            table.S_INFO_WINDCODE.label("sid"),
            table.ENTRY_DT.label("entry_dt"),
            table.REMOVE_DT.label("out_dt")
        )
        df = query.to_df()
        return df


# 创建API实例
_a_stock_api = AStockUniverseAPI()
_normal_index_api = NormalIndexUniverseAPI()
_msci_api = MSCIUniverseAPI()


def universe_A(cur_sign=True):
    """[返回最新全A成份股]

    Keyword Arguments:
        cur_sign {bool} -- [是否需要最新的股票池] (default: {True})

    Returns:
        [set] -- [set of stk code]
    """
    return _a_stock_api.get_a_stocks(cur_sign)


def universe_normal(indx, cur_sign=True):
    """[返回指数的最新份股]

    Arguments:
        indx {[str]} -- [wind code of index]
        cur_sign {bool} -- [是否需要最新的股票池] (default: {True})
    Raises:
        ValueError -- [description]
        ValueError -- [description]

    Returns:
        [set] -- [set of stk code]
    """
    try:
        indx_code = INDEX_DICT[indx]
    except KeyError:
        if "." not in indx:
            raise ValueError(indx_error)
        indx_code = indx

    stocks = _normal_index_api.get_index_stocks(indx_code, cur_sign)
    if not stocks:
        raise ValueError(indx_error)
    return stocks


def universe_msci(cur_sign=True):
    """[返回MSCI最新成分股]
    Arguments:
        cur_sign {bool} -- [是否需要最新的股票池] (default: {True})

    Returns:
        [set] -- [set of stk code]
    """
    return _msci_api.get_msci_stocks(cur_sign)


def Universe(indx, cur_sign=True):
    """[返回指数的最新成分股]

    Arguments:
        indx {[str]} -- [wind code of index or abbrev]
        cur_sign {bool} -- [是否需要最新的股票池] (default: {True})

    Returns:
        [set] -- [set of stk code]
    """
    if indx == "MSCI":
        return universe_msci(cur_sign=cur_sign)
    elif indx == "A":
        return universe_A(cur_sign=cur_sign)
    else:
        return universe_normal(indx, cur_sign=cur_sign)


@lru_cache()
def get_all_normal_index(index):
    return _normal_index_api.get_universe_dataframe(index)


@lru_cache()
def get_all_msci():
    return _msci_api.get_universe_dataframe("")


@lru_cache()
def get_all_stk():
    return _a_stock_api.get_universe_dataframe("881001.WI")


class HKUniverseAPI(UniverseAPI):
    """港股股票池API"""

    def get_table_name(self) -> str:
        return "HKSTOCKINDEXMEMBERS"

    def get_universe_dataframe(self, index_code: str) -> pd.DataFrame:
        """获取港股成分股的完整信息

        Args:
            index_code: 指数代码

        Returns:
            包含进入和退出日期的DataFrame
        """
        table = self.get_table()
        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 BondUniverseAPI(UniverseAPI):
    """可转债股票池API"""

    def get_table_name(self) -> str:
        return "CBINDEXMEMBERS"

    def get_universe_dataframe(self, index_code: str = "931078.CSI") -> pd.DataFrame:
        """获取可转债成分股的完整信息

        Args:
            index_code: 指数代码

        Returns:
            包含进入和退出日期的DataFrame
        """
        table = self.get_table()
        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


# 创建港股和可转债API实例
_hk_api = HKUniverseAPI()
_bond_api = BondUniverseAPI()


@lru_cache()
def get_all_hk(index):
    return _hk_api.get_universe_dataframe(index)


@lru_cache()
def get_all_bond():
    return _bond_api.get_universe_dataframe("931078.CSI")


class DynamicUniverseNormal(NormalIndexUniverseAPI, DynamicUniverseBase):
    """普通指数动态股票池"""

    def get_all_members_data(self) -> pd.DataFrame:
        """获取普通指数的所有历史成员数据"""
        return self.get_universe_dataframe(self.index)


class DynamicUniverseA(AStockUniverseAPI, DynamicUniverseBase):
    """A股动态股票池"""

    def get_all_members_data(self) -> pd.DataFrame:
        """获取A股的所有历史成员数据"""
        return self.get_universe_dataframe("881001.WI")


class DynamicUniverseMSCI(MSCIUniverseAPI, DynamicUniverseBase):
    """MSCI动态股票池"""

    def get_all_members_data(self) -> pd.DataFrame:
        """获取MSCI的所有历史成员数据"""
        return self.get_universe_dataframe("")


class DynamicUniverseHK(HKUniverseAPI, DynamicUniverseBase):
    """港股动态股票池"""

    def get_all_members_data(self) -> pd.DataFrame:
        """获取港股的所有历史成员数据"""
        return self.get_universe_dataframe(self.index)


class DynamicUniverseBond(BondUniverseAPI, DynamicUniverseBase):
    """可转债动态股票池"""

    def get_all_members_data(self) -> pd.DataFrame:
        """获取可转债的所有历史成员数据"""
        return self.get_universe_dataframe("931078.CSI")


@attr.s
class DynamicUniverse:
    """[get stock_list of universe on trade_dt]

    Args:
        indx: 指数代码或简称
        members: 自定义成员数据（可选）

    Raises:
        ValueError -- 无效的指数代码

    Returns:
        set -- 指定日期的股票代码集合
    """

    indx = attr.ib()
    index = attr.ib(init=False)
    members = attr.ib(default=None)
    _dynamic_api = attr.ib(init=False)

    @indx.validator
    def check(self, attribute, value):
        if value not in INDEX_DICT.keys():
            if "." not in value:
                raise ValueError(indx_error)

    def __attrs_post_init__(self):
        try:
            self.index = INDEX_DICT[self.indx]
        except KeyError:
            self.index = self.indx

        # 根据指数类型选择相应的动态API
        if self.indx == "MSCI":
            self._dynamic_api = DynamicUniverseMSCI()
        elif self.indx == "A":
            self._dynamic_api = DynamicUniverseA()
        elif self.indx == "CBOND":
            self._dynamic_api = DynamicUniverseBond()
        elif self.index.endswith("HI"):
            self._dynamic_api = DynamicUniverseHK()
        elif self.index != "":
            self._dynamic_api = DynamicUniverseNormal()
            self._dynamic_api.index = self.index
        else:
            raise ValueError(f"Invalid index: {self.indx}")

    def preview(self, trade_dt):
        """获取指定日期的股票池

        Args:
            trade_dt: 交易日期

        Returns:
            set: 股票代码集合
        """
        if isinstance(self.members, pd.DataFrame):
            df = self.members.copy()
        else:
            df = self._dynamic_api.get_all_members_data()

        trade_dt = dt_handle(trade_dt)
        df = df.loc[
            (df["entry_dt"] <= trade_dt)
            & ((df["out_dt"] >= trade_dt) | (df["out_dt"].isnull()))
        ]
        return set(df.sid)


if __name__ == "__main__":
    f_list = DynamicUniverse("HS300").preview("20180105")
    print(len(f_list))
