# coding:utf-8

"""
每日筹码及胜率数据采集工具
功能：从Tushare获取每日筹码及胜率数据，存储到gp_chip_data表
API文档: https://tushare.pro/document/2?doc_id=293
注意：该接口需要至少10000积分权限
每天17~18点左右更新
"""

import tushare as ts
import pandas as pd
from sqlalchemy import create_engine, text
from datetime import datetime, timedelta
from tqdm import tqdm
import time
import sys
import os

# 添加项目根目录到路径
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 ChipDataCollector:
    """每日筹码及胜率数据采集器"""

    def __init__(self, db_url, tushare_token, table_name='gp_chip_data'):
        """
        初始化采集器
        
        Args:
            db_url: 数据库连接URL
            tushare_token: Tushare API Token
            table_name: 数据表名，默认为gp_chip_data
        """
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600
        )
        self.table_name = table_name
        
        # 初始化Tushare
        ts.set_token(tushare_token)
        self.pro = ts.pro_api()
        
        print("=" * 60)
        print("每日筹码及胜率数据采集工具")
        print(f"目标数据表: {self.table_name}")
        print("=" * 60)

    def convert_tushare_code_to_db(self, ts_code):
        """
        将Tushare格式的股票代码转换为数据库格式
        Tushare格式：600000.SH、000001.SZ
        数据库格式：SH600000、SZ000001
        
        Args:
            ts_code: Tushare格式的股票代码
            
        Returns:
            数据库格式的股票代码
        """
        if '.SH' in ts_code:
            return f"SH{ts_code.replace('.SH', '')}"
        elif '.SZ' in ts_code:
            return f"SZ{ts_code.replace('.SZ', '')}"
        elif '.BJ' in ts_code:
            return f"BJ{ts_code.replace('.BJ', '')}"
        else:
            return ts_code

    def convert_db_code_to_tushare(self, symbol):
        """
        将数据库格式的股票代码转换为Tushare格式
        数据库格式：SH600000、SZ000001
        Tushare格式：600000.SH、000001.SZ
        
        Args:
            symbol: 数据库格式的股票代码
            
        Returns:
            Tushare格式的股票代码
        """
        if symbol.startswith('SH'):
            return f"{symbol[2:]}.SH"
        elif symbol.startswith('SZ'):
            return f"{symbol[2:]}.SZ"
        elif symbol.startswith('BJ'):
            return f"{symbol[2:]}.BJ"
        else:
            return symbol

    def fetch_all_stock_codes(self):
        """
        从数据库获取所有A股股票代码
        
        Returns:
            股票代码列表
        """
        try:
            query = "SELECT gp_code FROM gp_code_all"
            df = pd.read_sql(query, self.engine)
            codes = df['gp_code'].tolist()
            print(f"获取到A股股票代码: {len(codes)}个")
            return codes
        except Exception as e:
            print(f"获取股票代码失败: {e}")
            return []

    def fetch_chip_data(self, ts_code=None, trade_date=None, start_date=None, end_date=None):
        """
        从Tushare获取每日筹码及胜率数据
        API文档: https://tushare.pro/document/2?doc_id=293
        
        Args:
            ts_code: Tushare格式的股票代码，为None时获取全市场数据
            trade_date: 交易日期，格式YYYYMMDD（单日查询）
            start_date: 开始日期，格式YYYYMMDD（范围查询）
            end_date: 结束日期，格式YYYYMMDD（范围查询）
            
        Returns:
            筹码数据DataFrame
        """
        try:
            df = self.pro.cyq_perf(
                ts_code=ts_code,
                trade_date=trade_date,
                start_date=start_date,
                end_date=end_date,
                fields='ts_code,trade_date,his_low,his_high,cost_5pct,cost_15pct,cost_50pct,cost_85pct,cost_95pct,weight_avg,winner_rate'
            )
            return df
        except Exception as e:
            print(f"获取筹码数据失败: {e}")
            return pd.DataFrame()

    def transform_data(self, df):
        """
        转换Tushare数据为数据库格式
        
        Args:
            df: Tushare返回的DataFrame
            
        Returns:
            转换后的DataFrame
        """
        if df.empty:
            return pd.DataFrame()
        
        result_df = pd.DataFrame()
        
        # 字段映射
        result_df['symbol'] = df['ts_code'].apply(self.convert_tushare_code_to_db)
        result_df['ts_code'] = df['ts_code']
        result_df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
        result_df['his_low'] = df['his_low']
        result_df['his_high'] = df['his_high']
        result_df['cost_5pct'] = df['cost_5pct']
        result_df['cost_15pct'] = df['cost_15pct']
        result_df['cost_50pct'] = df['cost_50pct']
        result_df['cost_85pct'] = df['cost_85pct']
        result_df['cost_95pct'] = df['cost_95pct']
        result_df['weight_avg'] = df['weight_avg']
        result_df['winner_rate'] = df['winner_rate']
        
        return result_df

    def run_full_collection(self):
        """
        执行全量覆盖采集
        先清空表，再重新采集所有股票的数据
        """
        print("=" * 60)
        print("开始执行全量覆盖采集")
        print("=" * 60)
        
        try:
            # 获取所有股票代码
            stock_codes = self.fetch_all_stock_codes()
            if not stock_codes:
                print("未获取到股票代码，终止采集")
                return
            
            # 清空表
            print(f"正在清空{self.table_name}表...")
            with self.engine.begin() as conn:
                truncate_query = text(f"TRUNCATE TABLE {self.table_name}")
                conn.execute(truncate_query)
            print("表已清空")
            
            # 统计信息
            success_count = 0
            failed_count = 0
            total_records = 0
            
            # 遍历所有股票
            for symbol in tqdm(stock_codes, desc="全量采集筹码数据"):
                try:
                    # 转换股票代码格式
                    ts_code = self.convert_db_code_to_tushare(symbol)
                    
                    # 获取该股票的所有历史筹码数据
                    df = self.fetch_chip_data(ts_code=ts_code)
                    
                    if df.empty:
                        failed_count += 1
                        continue
                    
                    # 转换数据格式
                    result_df = self.transform_data(df)
                    
                    if result_df.empty:
                        failed_count += 1
                        continue
                    
                    # 保存到数据库
                    result_df.to_sql(self.table_name, self.engine, if_exists='append', index=False)
                    success_count += 1
                    total_records += len(result_df)
                    
                except Exception as e:
                    print(f"\n处理{symbol}时发生错误: {e}")
                    failed_count += 1
                    continue
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("全量采集完成")
            print("=" * 60)
            print(f"总股票数: {len(stock_codes)}")
            print(f"成功采集: {success_count}")
            print(f"失败: {failed_count}")
            print(f"总记录数: {total_records}")
            print("=" * 60)
            
        except Exception as e:
            print(f"全量采集失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.engine.dispose()

    def run_daily_collection(self, date=None):
        """
        执行每日增量采集
        获取指定日期的全市场筹码数据
        
        Args:
            date: 指定日期，格式YYYY-MM-DD，默认为今天
        """
        print("=" * 60)
        print("开始执行每日增量采集")
        print("=" * 60)
        
        try:
            # 确定采集日期
            if date is None:
                date_str = datetime.now().strftime('%Y-%m-%d')
            else:
                date_str = date
            
            # 转换为YYYYMMDD格式
            date_obj = datetime.strptime(date_str, '%Y-%m-%d')
            trade_date = date_obj.strftime('%Y%m%d')
            
            print(f"采集日期: {date_str} ({trade_date})")
            
            # 删除指定日期的旧数据
            with self.engine.begin() as conn:
                delete_query = text(f"DELETE FROM {self.table_name} WHERE DATE(trade_date) = :date_str")
                result = conn.execute(delete_query, {"date_str": date_str})
                print(f"已删除{self.table_name}表中{date_str}的旧数据，共{result.rowcount}条")
            
            # 获取全市场当日筹码数据（不传ts_code，获取所有股票）
            print(f"正在获取{date_str}全市场筹码数据...")
            df = self.fetch_chip_data(trade_date=trade_date)
            
            if df.empty:
                print(f"{date_str}没有筹码数据（可能是非交易日或接口无数据）")
                return
            
            print(f"成功获取筹码数据，共{len(df)}只股票")
            
            # 转换数据格式
            result_df = self.transform_data(df)
            
            if result_df.empty:
                print("数据转换失败")
                return
            
            # 保存到数据库
            print(f"正在保存数据到{self.table_name}表...")
            result_df.to_sql(self.table_name, self.engine, if_exists='append', index=False)
            print(f"成功保存{len(result_df)}条记录")
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("每日增量采集完成")
            print("=" * 60)
            print(f"采集日期: {date_str}")
            print(f"成功采集: {len(result_df)}只股票")
            print(f"数据表: {self.table_name}")
            print("=" * 60)
            
        except Exception as e:
            print(f"每日增量采集失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.engine.dispose()


def collect_chip_data(db_url, tushare_token, mode='daily', date=None):
    """
    筹码数据采集入口函数
    
    Args:
        db_url: 数据库连接URL
        tushare_token: Tushare API Token
        mode: 采集模式，'daily'=增量采集当日数据，'full'=全量覆盖采集
        date: 指定日期（仅mode='daily'时有效），格式YYYY-MM-DD
    """
    collector = ChipDataCollector(db_url, tushare_token)
    
    if mode == 'daily':
        collector.run_daily_collection(date)
    elif mode == 'full':
        collector.run_full_collection()
    else:
        print(f"未知的采集模式: {mode}")


if __name__ == "__main__":
    # 数据库配置
    db_url = 'mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj'
    
    # Tushare Token（从config.py读取）
    tushare_token = TUSHARE_TOKEN
    
    # 使用方式示例：
    
    # 1. 每日增量采集（推荐每日使用）
    # collect_chip_data(db_url, tushare_token, mode='daily')
    
    # 2. 全量覆盖采集（首次使用或需要重建数据时）
    # collect_chip_data(db_url, tushare_token, mode='full')
    
    # 3. 采集指定日期的数据
    collect_chip_data(db_url, tushare_token, mode='daily', date='2025-11-07')
