# coding:utf-8
"""
交易日历采集工具
功能：从 Tushare 获取交易日历数据并落库
API 文档: https://tushare.pro/document/2?doc_id=26
说明：每次调用时全量覆盖，不需要每日更新
"""

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 TradeCalCollector:
    """交易日历采集器"""

    def __init__(self, db_url: str, tushare_token: str, table_name: str = "trade_cal"):
        """
        Args:
            db_url: 数据库连接 URL
            tushare_token: Tushare Token
            table_name: 目标表名，默认 trade_cal
        """
        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_trade_cal(
        self,
        exchange: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
    ) -> pd.DataFrame:
        """
        获取交易日历数据
        API: trade_cal
        文档: https://tushare.pro/document/2?doc_id=26
        
        Args:
            exchange: 交易所代码 SSE-上交所 SZSE-深交所 BSE-北交所，None表示全部
            start_date: 开始日期，格式YYYYMMDD，None表示不限制
            end_date: 结束日期，格式YYYYMMDD，None表示不限制
        
        Returns:
            交易日历 DataFrame
        """
        try:
            params = {}
            if exchange:
                params['exchange'] = exchange
            if start_date:
                params['start_date'] = start_date
            if end_date:
                params['end_date'] = end_date
            
            print("正在获取交易日历数据...")
            if params:
                print(f"筛选条件: {params}")
            
            df = self.pro.trade_cal(**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["exchange"] = df["exchange"]
        result["cal_date"] = pd.to_datetime(df["cal_date"], format="%Y%m%d", errors='coerce')
        result["is_open"] = df.get("is_open", None)
        result["pretrade_date"] = pd.to_datetime(df["pretrade_date"], format="%Y%m%d", errors='coerce') if "pretrade_date" in df.columns else 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,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
    ) -> None:
        """
        执行全量覆盖采集：
        - 清空目标表
        - 获取交易日历数据
        - 全量写入数据库
        
        Args:
            exchange: 交易所代码，None表示全部
            start_date: 开始日期，格式YYYYMMDD，None表示不限制
            end_date: 结束日期，格式YYYYMMDD，None表示不限制
        """
        print("=" * 60)
        print("开始执行全量覆盖采集（交易日历）")
        print("=" * 60)

        try:
            # 清空表
            with self.engine.begin() as conn:
                # 如果指定了筛选条件，只删除符合条件的记录；否则清空全部
                if exchange or start_date or end_date:
                    delete_conditions = []
                    params = {}
                    if exchange:
                        delete_conditions.append("exchange = :exchange")
                        params["exchange"] = exchange
                    if start_date:
                        delete_conditions.append("cal_date >= :start_date")
                        params["start_date"] = datetime.strptime(start_date, "%Y%m%d").date()
                    if end_date:
                        delete_conditions.append("cal_date <= :end_date")
                        params["end_date"] = datetime.strptime(end_date, "%Y%m%d").date()
                    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} 已清空")
            
            # 获取交易日历数据
            cal_df = self.fetch_trade_cal(exchange=exchange, start_date=start_date, end_date=end_date)
            if cal_df.empty:
                print("未获取到交易日历数据，采集终止")
                return

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

            # 保存数据
            self.save_dataframe(result_df)
            print(f"成功写入 {len(result_df)} 条记录")
            
            # 按交易所统计
            if "exchange" in result_df.columns:
                print("\n按交易所统计:")
                for exchange_name in result_df["exchange"].unique():
                    exchange_data = result_df[result_df["exchange"] == exchange_name]
                    total_days = len(exchange_data)
                    trade_days = int(exchange_data["is_open"].sum()) if "is_open" in exchange_data.columns and exchange_data["is_open"].notna().any() else 0
                    min_date = exchange_data["cal_date"].min()
                    max_date = exchange_data["cal_date"].max()
                    print(f"  {exchange_name}: 总天数={total_days}, 交易日={trade_days}, 日期范围={min_date} ~ {max_date}")
            
            print("\n" + "=" * 60)
            print("全量覆盖采集完成")
            print("=" * 60)
        except Exception as exc:
            print(f"全量采集失败: {exc}")
            import traceback
            traceback.print_exc()
        finally:
            self.engine.dispose()


def collect_trade_cal(
    db_url: str,
    tushare_token: str,
    exchange: Optional[str] = None,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
):
    """
    采集入口 - 全量覆盖采集
    
    Args:
        db_url: 数据库连接URL
        tushare_token: Tushare Token
        exchange: 交易所代码 SSE-上交所 SZSE-深交所 BSE-北交所，None表示全部
        start_date: 开始日期，格式YYYYMMDD，None表示不限制
        end_date: 结束日期，格式YYYYMMDD，None表示不限制
    """
    collector = TradeCalCollector(db_url, tushare_token)
    collector.run_full_collection(exchange=exchange, start_date=start_date, end_date=end_date)


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_trade_cal(DB_URL, TOKEN)
    
    # 如果只想获取特定交易所，可以这样调用：
    # collect_trade_cal(DB_URL, TOKEN, exchange='SSE')

