# coding:utf-8
"""
同花顺概念和行业指数采集工具
功能：从 Tushare 获取同花顺概念和行业指数数据并落库
API 文档: https://tushare.pro/document/2?doc_id=259
说明：每次调用时全量覆盖，不需要每日更新
"""

import os
import sys
from datetime import datetime
from typing import Optional

import pandas as pd
import tushare as ts
from sqlalchemy import create_engine, text

# 添加项目根目录到路径，确保能够读取配置
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(PROJECT_ROOT)

from src.scripts.config import TUSHARE_TOKEN


class THSIndexCollector:
    """同花顺概念和行业指数采集器"""

    def __init__(self, db_url: str, tushare_token: str, table_name: str = "ths_index"):
        """
        Args:
            db_url: 数据库连接 URL
            tushare_token: Tushare Token
            table_name: 目标表名，默认 ths_index
        """
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600,
        )
        self.table_name = table_name

        ts.set_token(tushare_token)
        self.pro = ts.pro_api()

        print("=" * 60)
        print("同花顺概念和行业指数采集工具")
        print(f"目标数据表: {self.table_name}")
        print("=" * 60)

    def fetch_index_list(
        self,
        exchange: Optional[str] = None,
        index_type: Optional[str] = None,
    ) -> pd.DataFrame:
        """
        获取同花顺概念和行业指数列表
        API: ths_index
        文档: https://tushare.pro/document/2?doc_id=259
        
        Args:
            exchange: 市场类型 A-a股 HK-港股 US-美股，None表示全部
            index_type: 指数类型 N-概念指数 I-行业指数 R-地域指数 S-同花顺特色指数 
                       ST-同花顺风格指数 TH-同花顺主题指数 BB-同花顺宽基指数，None表示全部
        
        Returns:
            指数列表 DataFrame
        """
        try:
            params = {}
            if exchange:
                params['exchange'] = exchange
            if index_type:
                params['type'] = index_type
            
            print(f"正在获取同花顺指数列表...")
            if params:
                print(f"筛选条件: {params}")
            
            df = self.pro.ths_index(**params)
            if df.empty:
                print("未获取到指数数据")
                return pd.DataFrame()
            print(f"成功获取 {len(df)} 个指数")
            return df
        except Exception as exc:
            print(f"获取指数列表失败: {exc}")
            return pd.DataFrame()

    def transform_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        将 Tushare 返回的数据转换为数据库入库格式
        
        Args:
            df: 指数信息 DataFrame
            
        Returns:
            转换后的 DataFrame
        """
        if df.empty:
            return pd.DataFrame()

        result = pd.DataFrame()
        result["ts_code"] = df["ts_code"]
        result["name"] = df["name"]
        result["count"] = df.get("count", None)
        result["exchange"] = df.get("exchange", None)
        result["list_date"] = pd.to_datetime(df["list_date"], format="%Y%m%d", errors='coerce') if "list_date" in df.columns else None
        result["type"] = df.get("type", None)
        result["created_at"] = datetime.now()
        result["updated_at"] = datetime.now()
        return result

    def save_dataframe(self, df: pd.DataFrame) -> None:
        """
        将数据写入数据库
        """
        if df.empty:
            return
        df.to_sql(self.table_name, self.engine, if_exists="append", index=False)

    def run_full_collection(
        self,
        exchange: Optional[str] = None,
        index_type: Optional[str] = None,
    ) -> None:
        """
        执行全量覆盖采集：
        - 清空目标表
        - 获取所有指数数据
        - 全量写入数据库
        
        Args:
            exchange: 市场类型，None表示全部
            index_type: 指数类型，None表示全部
        """
        print("=" * 60)
        print("开始执行全量覆盖采集（同花顺概念和行业指数）")
        print("=" * 60)

        try:
            # 清空表
            with self.engine.begin() as conn:
                # 如果指定了筛选条件，只删除符合条件的记录；否则清空全部
                if exchange or index_type:
                    delete_conditions = []
                    params = {}
                    if exchange:
                        delete_conditions.append("exchange = :exchange")
                        params["exchange"] = exchange
                    if index_type:
                        delete_conditions.append("type = :type")
                        params["type"] = index_type
                    delete_sql = text(f"DELETE FROM {self.table_name} WHERE {' AND '.join(delete_conditions)}")
                    conn.execute(delete_sql, params)
                    print(f"{self.table_name} 已删除符合条件的旧数据")
                else:
                    conn.execute(text(f"TRUNCATE TABLE {self.table_name}"))
                    print(f"{self.table_name} 已清空")
            
            # 获取指数列表
            index_df = self.fetch_index_list(exchange=exchange, index_type=index_type)
            if index_df.empty:
                print("未获取到指数数据，采集终止")
                return

            # 转换数据格式
            result_df = self.transform_data(index_df)
            if result_df.empty:
                print("数据转换失败")
                return

            # 保存数据
            self.save_dataframe(result_df)
            print(f"成功写入 {len(result_df)} 条记录")
            
            # 按类型统计
            if "type" in result_df.columns:
                type_stats = result_df.groupby("type").size()
                print("\n按类型统计:")
                for idx, count in type_stats.items():
                    print(f"  {idx}: {count} 个")
            
            print("\n" + "=" * 60)
            print("全量覆盖采集完成")
            print("=" * 60)
        except Exception as exc:
            print(f"全量采集失败: {exc}")
            import traceback
            traceback.print_exc()
        finally:
            self.engine.dispose()


def collect_ths_index(
    db_url: str,
    tushare_token: str,
    exchange: Optional[str] = None,
    index_type: Optional[str] = None,
):
    """
    采集入口 - 全量覆盖采集
    
    Args:
        db_url: 数据库连接URL
        tushare_token: Tushare Token
        exchange: 市场类型 A-a股 HK-港股 US-美股，None表示全部
        index_type: 指数类型 N-概念指数 I-行业指数 R-地域指数 S-同花顺特色指数 
                   ST-同花顺风格指数 TH-同花顺主题指数 BB-同花顺宽基指数，None表示全部
    """
    collector = THSIndexCollector(db_url, tushare_token)
    collector.run_full_collection(exchange=exchange, index_type=index_type)


if __name__ == "__main__":
    # DB_URL = "mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj"
    DB_URL = "mysql+pymysql://fac_pattern:Chlry$%.8_app@192.168.16.153:3307/my_quant_db"
    TOKEN = TUSHARE_TOKEN

    # 执行全量覆盖采集（获取所有类型的指数）
    collect_ths_index(DB_URL, TOKEN)
    
    # 如果只想获取特定类型，可以这样调用：
    # collect_ths_index(DB_URL, TOKEN, exchange='A', index_type='N')  # 只获取A股概念指数

