"""
行业分类数据API基础模块

提供统一的行业数据查询和处理基础类
"""

import re
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


class IndustryAPI(BaseAPI):
    """行业分类API基础类"""

    # 类级别的正则表达式，避免重复编译
    CHINESE_PATTERN = re.compile(r"[^\u4e00-\u9fa5]")

    def __init__(self, database=None):
        """初始化行业API

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

    def create_industry_query(self, classification_table, code_table, level: int = 1):
        """创建行业分类查询的基础查询

        Args:
            classification_table: 行业分类表
            code_table: 行业代码表
            level: 行业级别

        Returns:
            查询对象
        """
        query = self.db.query(
            classification_table.S_INFO_WINDCODE,
            classification_table.ENTRY_DT,
            classification_table.REMOVE_DT,
            code_table.INDUSTRIESNAME
        ).filter(code_table.LEVELNUM == level)

        return query

    def filter_by_industry_code(self, query, classification_table, code_table, level: int = 1):
        """根据行业代码过滤查询

        Args:
            query: 查询对象
            classification_table: 行业分类表
            code_table: 行业代码表
            level: 行业级别

        Returns:
            过滤后的查询对象
        """
        try:
            return query.filter(
                func.substring(classification_table.CITICS_IND_CODE, 1, 2 + 2 * level)
                == func.substring(code_table.INDUSTRIESCODE, 1, 2 + 2 * level)
            )
        except:
            # 处理不同数据库系统的兼容性
            return query.filter(
                func.substr(classification_table.CITICS_IND_CODE, 1, 2 + 2 * level)
                == func.substr(code_table.INDUSTRIESCODE, 1, 2 + 2 * level)
            )

    def process_industry_dataframe(self, df, adjust: bool = False) -> pd.DataFrame:
        """处理行业DataFrame，包括中文处理和调整

        Args:
            df: 原始DataFrame
            adjust: 是否进行行业调整

        Returns:
            处理后的DataFrame
        """
        df.columns = ["sid", "entry_dt", "out_dt", "ind"]

        # 去除行业中的罗马数字，使用类级别的正则表达式
        df.ind = df.ind.str.replace(self.CHINESE_PATTERN, "", regex=True)

        # 行业调整
        if adjust:
            df.ind = df.ind.map(self._get_industry_mapping())

        return df

    def filter_by_trade_date(self, df, trade_dt: str) -> pd.DataFrame:
        """根据交易日期过滤数据

        Args:
            df: DataFrame
            trade_dt: 交易日期

        Returns:
            过滤后的DataFrame
        """
        return df.loc[
            (df["entry_dt"] <= trade_dt)
            & ((df["out_dt"] >= trade_dt) | (df["out_dt"].isnull()))
        ].copy()

    def filter_by_stocks(self, df, sid: Optional[Union[str, List[str], Set[str]]]) -> pd.DataFrame:
        """根据股票代码过滤数据

        Args:
            df: DataFrame
            sid: 股票代码

        Returns:
            过滤后的DataFrame
        """
        if sid is not None:
            sid_set = {sid} if isinstance(sid, str) else set(sid)
            return df[df["sid"].isin(sid_set)].copy()
        return df

    def _get_industry_mapping(self) -> Dict[str, str]:
        """获取行业映射关系

        Returns:
            行业映射字典
        """
        return {
            "综合金融": "非银行金融",
            "多领域控股": "多元金融",
            "资产管理": "多元金融",
            "新兴金融服务": "多元金融",
            "全国性股份制银行": "股份制与城商行",
            "区域性银行": "股份制与城商行",
        }

    @abstractmethod
    def get_all_industries_data(self, *args, **kwargs) -> pd.DataFrame:
        """获取所有行业数据，子类必须实现

        Returns:
            包含行业信息的DataFrame
        """
        pass


class ZXIndustryAPI(IndustryAPI):
    """中信行业分类API"""

    def get_table_name(self) -> str:
        """获取中信行业分类表名"""
        return "ASHAREINDUSTRIESCLASSCITICS"

    def get_code_table(self):
        """获取行业代码表"""
        return self.db.ASHAREINDUSTRIESCODE

    @lru_cache(maxsize=128)
    def get_all_industries_data(self, level: int = 1, adjust: bool = False) -> pd.DataFrame:
        """获取所有中信行业数据

        Args:
            level: 行业级别
            adjust: 是否进行行业调整

        Returns:
            包含行业信息的DataFrame
        """
        classification_table = self.get_table()
        code_table = self.get_code_table()

        query = self.create_industry_query(classification_table, code_table, level + 1)
        query = self.filter_by_industry_code(query, classification_table, code_table, level)

        df = query.to_df()
        return self.process_industry_dataframe(df, adjust)

    def get_industry_at_date(self, trade_dt: str, sid: Optional[Union[str, List[str], Set[str]]] = None,
                           level: int = 1, adjust: bool = False) -> pd.DataFrame:
        """获取指定日期的行业分类

        Args:
            trade_dt: 交易日期
            sid: 股票代码
            level: 行业级别
            adjust: 是否进行行业调整

        Returns:
            行业分类DataFrame
        """
        df = self.get_all_industries_data(level=level, adjust=adjust)
        df = self.filter_by_stocks(df, sid)
        df = self.filter_by_trade_date(df, trade_dt)
        return df.set_index("sid")[["ind"]].copy()


class WindHKIndustryAPI(IndustryAPI):
    """Wind港股行业分类API"""

    def get_table_name(self) -> str:
        """获取Wind港股行业分类表名"""
        return "HKSTOCKWINDINDUSTRIESMEMBERS"

    def get_code_table(self):
        """获取行业代码表"""
        return self.db.HKSTOCKINDUSTRIESCODE

    @lru_cache(maxsize=128)
    def get_all_industries_data(self, level: int = 1) -> pd.DataFrame:
        """获取所有Wind港股行业数据

        Args:
            level: 行业级别

        Returns:
            包含行业信息的DataFrame
        """
        classification_table = self.get_table()
        code_table = self.get_code_table()

        query = self.db.query(
            classification_table.S_INFO_WINDCODE,
            classification_table.ENTRY_DT,
            classification_table.REMOVE_DT,
            code_table.INDUSTRIESNAME
        ).filter(
            code_table.LEVELNUM == level,
            code_table.industriesclass == "Wind行业分类"
        )

        try:
            query = query.filter(
                func.substring(classification_table.WIND_IND_CODE, 1, 2 + 2 * level)
                == func.substring(code_table.INDUSTRIESCODE, 1, 2 + 2 * level)
            )
        except:
            query = query.filter(
                func.substr(classification_table.WIND_IND_CODE, 1, 2 + 2 * level)
                == func.substr(code_table.INDUSTRIESCODE, 1, 2 + 2 * level)
            )

        df = query.to_df()
        df.columns = ["sid", "entry_dt", "out_dt", "ind"]
        df.ind = df.ind.str.replace(self.CHINESE_PATTERN, "", regex=True)
        return df


class HKCITICSIndustryAPI(IndustryAPI):
    """港股中信行业分类API"""

    def get_table_name(self) -> str:
        """获取港股中信行业分类表名"""
        return "HKSHARESCITICSINDUSTRIESCLASS"

    @lru_cache(maxsize=128)
    def get_all_industries_data(self) -> pd.DataFrame:
        """获取所有港股中信行业数据

        Returns:
            包含行业信息的DataFrame
        """
        table = self.get_table()
        df = (
            self.db.query(
                table.S_INFO_WINDCODE.label("sid"),
                table.CHANGE_DT.label("change_dt"),
                table.citics_ind_name.label("ind"),
            )
            .filter(table.change_type == 1)
            .to_df()
        )

        df.ind = df.ind.str.slice(0, -7)
        return df

    def get_industry_at_date(self, trade_dt: str, sid: Optional[Union[str, List[str], Set[str]]] = None) -> pd.DataFrame:
        """获取指定日期的港股中信行业分类

        Args:
            trade_dt: 交易日期
            sid: 股票代码

        Returns:
            行业分类DataFrame
        """
        df = self.get_all_industries_data()
        df = self.filter_by_stocks(df, sid)
        df = df[df["change_dt"] <= trade_dt].copy()
        df.sort_values(["sid", "change_dt"], inplace=True)
        df.drop_duplicates("sid", keep="last", inplace=True)
        return df.set_index("sid")[["ind"]].copy()


# 创建全局API实例
_zx_industry_api = ZXIndustryAPI()
_wind_hk_industry_api = WindHKIndustryAPI()
_hk_citics_industry_api = HKCITICSIndustryAPI()

__all__ = [
    'IndustryAPI',
    'ZXIndustryAPI',
    'WindHKIndustryAPI',
    'HKCITICSIndustryAPI',
    '_zx_industry_api',
    '_wind_hk_industry_api',
    '_hk_citics_industry_api'
]