# coding:utf-8

"""
每日筹码分布数据采集工具
功能：从Tushare获取每日筹码分布数据，存储到gp_chip_distribution表
API文档: https://tushare.pro/document/2?doc_id=294
注意：该接口需要至少10000积分权限，单次最大返回2000条数据
每天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 ChipDistributionCollector:
    """每日筹码分布数据采集器（按年分表）"""

    def __init__(self, db_url, tushare_token, base_table_name='gp_chip_distribution'):
        """
        初始化采集器
        
        Args:
            db_url: 数据库连接URL
            tushare_token: Tushare API Token
            base_table_name: 基础表名，默认为gp_chip_distribution（实际表名会加上年度后缀，如gp_chip_distribution_2025）
        """
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600
        )
        self.base_table_name = base_table_name
        
        # 初始化Tushare
        ts.set_token(tushare_token)
        self.pro = ts.pro_api()
        
        print("=" * 60)
        print("每日筹码分布数据采集工具（按年分表）")
        print(f"基础表名: {self.base_table_name}")
        print("=" * 60)

    def get_table_name_by_date(self, date_str):
        """
        根据日期获取对应的分表名称
        
        Args:
            date_str: 日期字符串，格式YYYY-MM-DD
            
        Returns:
            分表名称，如gp_chip_distribution_2025
        """
        date_obj = datetime.strptime(date_str, '%Y-%m-%d')
        year = date_obj.strftime('%Y')
        return f"{self.base_table_name}_{year}"
    
    def create_yearly_table(self, table_name):
        """
        创建年度分表（如果不存在）
        
        Args:
            table_name: 表名，如gp_chip_distribution_2025
        """
        try:
            create_table_sql = f"""
            CREATE TABLE IF NOT EXISTS `{table_name}` (
              `id` bigint NOT NULL AUTO_INCREMENT COMMENT '主键',
              `symbol` varchar(20) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL COMMENT '股票代码（数据库格式，如SH600000）',
              `ts_code` varchar(20) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin NOT NULL COMMENT 'Tushare股票代码（如600000.SH）',
              `trade_date` date NOT NULL COMMENT '交易日期',
              `price` decimal(10, 2) NOT NULL COMMENT '价格',
              `percent` decimal(10, 4) NULL DEFAULT NULL COMMENT '占比（%）',
              `create_time` datetime NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
              PRIMARY KEY (`id`) USING BTREE,
              UNIQUE INDEX `idx_symbol_date_price`(`symbol` ASC, `trade_date` ASC, `price` ASC) USING BTREE,
              INDEX `idx_symbol`(`symbol` ASC) USING BTREE,
              INDEX `idx_trade_date`(`trade_date` ASC) USING BTREE,
              INDEX `idx_ts_code`(`ts_code` ASC) USING BTREE
            ) ENGINE = InnoDB 
            CHARACTER SET = utf8mb4 
            COLLATE = utf8mb4_bin 
            COMMENT = '每日筹码分布数据表（Tushare）' 
            ROW_FORMAT = Dynamic;
            """
            
            with self.engine.begin() as conn:
                conn.execute(text(create_table_sql))
            
            print(f"✓ 表 {table_name} 已确认存在")
            
        except Exception as e:
            print(f"创建表 {table_name} 失败: {e}")
            raise
    
    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_copy"
            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_distribution(self, ts_code=None, trade_date=None, start_date=None, end_date=None):
        """
        从Tushare获取每日筹码分布数据
        API文档: https://tushare.pro/document/2?doc_id=294
        注意：单次最大返回2000条数据
        
        Args:
            ts_code: Tushare格式的股票代码，为None时获取全市场数据
            trade_date: 交易日期，格式YYYYMMDD（单日查询）
            start_date: 开始日期，格式YYYYMMDD（范围查询）
            end_date: 结束日期，格式YYYYMMDD（范围查询）
            
        Returns:
            筹码分布DataFrame
        """
        try:
            df = self.pro.cyq_chips(
                ts_code=ts_code,
                trade_date=trade_date,
                start_date=start_date,
                end_date=end_date,
                fields='ts_code,trade_date,price,percent'
            )
            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['price'] = df['price']
        result_df['percent'] = df['percent']
        
        return result_df

    def run_full_collection(self, start_date=None, end_date=None, batch_size=100):
        """
        执行全量覆盖采集（按年分表，按日期倒序采集）
        从最新日期（end_date）开始往前采集到最早日期（start_date）
        自动跳过周六和周日（非交易日）
        
        Args:
            start_date: 开始日期（最早日期），格式YYYY-MM-DD，默认为8年前
            end_date: 结束日期（最新日期），格式YYYY-MM-DD，默认为今天
            batch_size: 批量入库的股票数量，默认100
        """
        print("=" * 60)
        print("开始执行全量覆盖采集（按年分表，按日期倒序采集）")
        print("=" * 60)
        
        try:
            # 确定日期范围
            if end_date is None:
                end_date_obj = datetime.now()
            else:
                end_date_obj = datetime.strptime(end_date, '%Y-%m-%d')
            
            if start_date is None:
                # 默认采集近8年历史数据
                start_date_obj = end_date_obj - timedelta(days=8*365)
            else:
                start_date_obj = datetime.strptime(start_date, '%Y-%m-%d')
            
            print(f"日期范围: {start_date_obj.strftime('%Y-%m-%d')} 至 {end_date_obj.strftime('%Y-%m-%d')}")
            print(f"批量入库设置: 每{batch_size}只股票入库一次")
            print("=" * 60)
            
            # 获取所有股票代码
            stock_codes = self.fetch_all_stock_codes()
            if not stock_codes:
                print("未获取到股票代码，终止采集")
                return
            
            # 按日期倒序逐日采集（从最新日期往前采集，跳过周六周日）
            current_date = end_date_obj
            total_days = (end_date_obj - start_date_obj).days + 1
            processed_days = 0
            skipped_weekends = 0
            
            while current_date >= start_date_obj:
                # 跳过周六（weekday() == 5）和周日（weekday() == 6）
                weekday = current_date.weekday()
                if weekday == 5:  # 周六
                    skipped_weekends += 1
                    current_date -= timedelta(days=1)
                    continue
                elif weekday == 6:  # 周日
                    skipped_weekends += 1
                    current_date -= timedelta(days=1)
                    continue
                
                date_str = current_date.strftime('%Y-%m-%d')
                
                # 调用每日采集方法，自动按年分表
                print(f"\n[{processed_days + 1}/{total_days - skipped_weekends}] 采集日期: {date_str} (倒序采集)")
                
                try:
                    # 使用每日采集方法，自动按年分表
                    self.run_daily_collection(date=date_str, batch_size=batch_size)
                    processed_days += 1
                except Exception as e:
                    print(f"采集 {date_str} 失败: {e}")
                
                # 移到前一天（倒序）
                current_date -= timedelta(days=1)
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("全量覆盖采集完成")
            print("=" * 60)
            print(f"日期范围: {start_date_obj.strftime('%Y-%m-%d')} 至 {end_date_obj.strftime('%Y-%m-%d')}")
            print(f"总天数: {total_days}天")
            print(f"跳过周末: {skipped_weekends}天")
            print(f"实际采集: {processed_days}天")
            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, batch_size=100):
        """
        执行每日增量采集（按年分表）
        获取指定日期的全市场筹码分布数据
        注意：单次返回限制10000条
        
        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')
            
            # 获取当年分表名称
            table_name = self.get_table_name_by_date(date_str)
            
            print(f"采集日期: {date_str} ({trade_date})")
            print(f"目标表名: {table_name}")
            print(f"批量入库设置: 每{batch_size}只股票入库一次")
            
            # 创建当年分表（如果不存在）
            self.create_yearly_table(table_name)
            
            # 删除指定日期的旧数据
            with self.engine.begin() as conn:
                delete_query = text(f"DELETE FROM {table_name} WHERE DATE(trade_date) = :date_str")
                result = conn.execute(delete_query, {"date_str": date_str})
                print(f"已删除{table_name}表中{date_str}的旧数据，共{result.rowcount}条")
            
            # 获取所有股票代码
            stock_codes = self.fetch_all_stock_codes()
            if not stock_codes:
                print("未获取到股票代码，终止采集")
                return
            
            # 统计信息
            success_count = 0
            failed_count = 0
            total_records = 0
            
            # 批量数据缓存
            batch_data_list = []
            
            # 遍历所有股票，获取当日筹码分布数据
            print(f"正在获取{date_str}全市场筹码分布数据...")
            for idx, symbol in enumerate(tqdm(stock_codes, desc=f"采集{date_str}筹码分布"), 1):
                try:
                    # 转换股票代码格式
                    ts_code = self.convert_db_code_to_tushare(symbol)
                    
                    # 获取当日筹码分布数据
                    df = self.fetch_chip_distribution(ts_code=ts_code, trade_date=trade_date)
                    
                    if df.empty:
                        failed_count += 1
                        continue
                    
                    # 转换数据格式
                    result_df = self.transform_data(df)
                    
                    if result_df.empty:
                        failed_count += 1
                        continue
                    
                    # 添加到批量缓存
                    batch_data_list.append(result_df)
                    success_count += 1
                    total_records += len(result_df)
                    
                    # 每达到batch_size或最后一只股票时，批量入库
                    if len(batch_data_list) >= batch_size or idx == len(stock_codes):
                        if batch_data_list:
                            # 合并所有DataFrame
                            combined_df = pd.concat(batch_data_list, ignore_index=True)
                            # 批量保存到数据库（使用分表）
                            combined_df.to_sql(table_name, self.engine, if_exists='append', index=False)
                            # 清空缓存
                            batch_data_list = []
                    
                except Exception as e:
                    failed_count += 1
                    continue
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("每日增量采集完成")
            print("=" * 60)
            print(f"采集日期: {date_str}")
            print(f"总股票数: {len(stock_codes)}")
            print(f"成功采集: {success_count}")
            print(f"失败: {failed_count}")
            print(f"总记录数: {total_records}")
            print(f"数据表: {table_name}")
            print("=" * 60)
            
        except Exception as e:
            print(f"每日增量采集失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.engine.dispose()


def collect_chip_distribution(db_url, tushare_token, mode='daily', date=None, start_date=None, end_date=None, batch_size=100):
    """
    筹码分布数据采集入口函数（按年分表）
    
    Args:
        db_url: 数据库连接URL
        tushare_token: Tushare API Token
        mode: 采集模式，'daily'=增量采集当日数据，'full'=全量覆盖采集（倒序）
        date: 指定日期（仅mode='daily'时有效），格式YYYY-MM-DD
        start_date: 最早日期（仅mode='full'时有效），格式YYYY-MM-DD，默认为8年前
        end_date: 最新日期（仅mode='full'时有效），格式YYYY-MM-DD，默认为今天
        batch_size: 批量入库的股票数量，默认100
    
    注意：
        - full模式采用倒序采集：从最新日期（end_date）往前采集到最早日期（start_date）
        - 这样可以优先采集最新数据，如果中途中断，最新数据已经采集完成
        - 自动跳过周六和周日（非交易日），只采集工作日数据
    """
    collector = ChipDistributionCollector(db_url, tushare_token)
    
    if mode == 'daily':
        collector.run_daily_collection(date, batch_size)
    elif mode == 'full':
        collector.run_full_collection(start_date, end_date, batch_size)
    else:
        print(f"未知的采集模式: {mode}")


if __name__ == "__main__":
    # 数据库配置
    db_url = 'mysql+pymysql://fac_pattern:Chlry$%.8pattern@192.168.16.150:3306/factordb_mysql'
    
    # Tushare Token（从config.py读取）
    tushare_token = TUSHARE_TOKEN
    
    # 使用方式示例：
    
    # 1. 每日增量采集（推荐每日使用）
    # 自动按年分表，如2025-11-03的数据会存到gp_chip_distribution_2025表
    # collect_chip_distribution(db_url, tushare_token, mode='daily')
    
    # 2. 全量覆盖采集（首次使用或需要重建数据时）
    # 采集近8年数据，自动按年创建分表（如gp_chip_distribution_2025, gp_chip_distribution_2024等）
    # collect_chip_distribution(db_url, tushare_token, mode='full')
    
    # 3. 采集指定日期的数据
    # collect_chip_distribution(db_url, tushare_token, mode='daily', date='2025-11-25')
    
    # 4. 采集指定日期范围的数据
    collect_chip_distribution(db_url, tushare_token, mode='full',
                             start_date='2025-11-28', end_date='2025-12-09')
    
    # 5. 调整批量入库大小（默认100只股票一批）
    # collect_chip_distribution(db_url, tushare_token, mode='daily', batch_size=200)

