# coding:utf-8

"""
股票集合竞价数据采集工具
功能：从Tushare获取开盘和收盘集合竞价数据，合并存储到gp_auction_data表
API文档: 
  开盘集合竞价 https://tushare.pro/document/2?doc_id=353
  收盘集合竞价 https://tushare.pro/document/2?doc_id=354
注意：
  1. 该接口需要至少2000积分权限
  2. 单次最大返回10000条数据
  3. 采用按日期遍历策略，每天获取全市场数据
  4. 每天15:30之后数据更新完成
采集策略：
  - 按日期遍历，每个日期获取全市场的开盘和收盘集合竞价数据
  - 将同一股票同一日期的开盘和收盘数据合并为一条记录
"""

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 AuctionDataCollector:
    """股票集合竞价数据采集器"""

    def __init__(self, db_url, tushare_token, table_name='gp_auction_data'):
        """
        初始化采集器
        
        Args:
            db_url: 数据库连接URL
            tushare_token: Tushare API Token
            table_name: 数据表名，默认为gp_auction_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 fetch_open_auction(self, trade_date):
        """
        获取开盘集合竞价数据
        
        Args:
            trade_date: 交易日期，格式YYYYMMDD
            
        Returns:
            开盘集合竞价DataFrame
        """
        try:
            df = self.pro.stk_auction_o(trade_date=trade_date)
            return df
        except Exception as e:
            print(f"获取开盘集合竞价数据失败: {e}")
            return pd.DataFrame()

    def fetch_close_auction(self, trade_date):
        """
        获取收盘集合竞价数据
        
        Args:
            trade_date: 交易日期，格式YYYYMMDD
            
        Returns:
            收盘集合竞价DataFrame
        """
        try:
            df = self.pro.stk_auction_c(trade_date=trade_date)
            time.sleep(0.2)  # API限频控制
            return df
        except Exception as e:
            print(f"获取收盘集合竞价数据失败: {e}")
            return pd.DataFrame()

    def merge_and_transform_data(self, open_df, close_df):
        """
        合并开盘和收盘集合竞价数据，并转换为数据库格式
        
        Args:
            open_df: 开盘集合竞价DataFrame
            close_df: 收盘集合竞价DataFrame
            
        Returns:
            合并后的DataFrame
        """
        if open_df.empty and close_df.empty:
            return pd.DataFrame()
        
        # 处理开盘集合竞价数据
        if not open_df.empty:
            open_df = open_df.rename(columns={
                'close': 'open_close',
                'open': 'open_open',
                'high': 'open_high',
                'low': 'open_low',
                'vol': 'open_vol',
                'amount': 'open_amount',
                'vwap': 'open_vwap'
            })
            open_df['trade_date'] = pd.to_datetime(open_df['trade_date'], format='%Y%m%d')
        
        # 处理收盘集合竞价数据
        if not close_df.empty:
            close_df = close_df.rename(columns={
                'close': 'close_close',
                'open': 'close_open',
                'high': 'close_high',
                'low': 'close_low',
                'vol': 'close_vol',
                'amount': 'close_amount',
                'vwap': 'close_vwap'
            })
            close_df['trade_date'] = pd.to_datetime(close_df['trade_date'], format='%Y%m%d')
        
        # 合并数据
        if not open_df.empty and not close_df.empty:
            # 两个都有数据，按ts_code和trade_date合并
            result_df = pd.merge(
                open_df[['ts_code', 'trade_date', 'open_close', 'open_open', 'open_high', 
                        'open_low', 'open_vol', 'open_amount', 'open_vwap']],
                close_df[['ts_code', 'trade_date', 'close_close', 'close_open', 'close_high',
                         'close_low', 'close_vol', 'close_amount', 'close_vwap']],
                on=['ts_code', 'trade_date'],
                how='outer'
            )
        elif not open_df.empty:
            # 只有开盘数据
            result_df = open_df[['ts_code', 'trade_date', 'open_close', 'open_open', 'open_high',
                                'open_low', 'open_vol', 'open_amount', 'open_vwap']].copy()
        else:
            # 只有收盘数据
            result_df = close_df[['ts_code', 'trade_date', 'close_close', 'close_open', 'close_high',
                                 'close_low', 'close_vol', 'close_amount', 'close_vwap']].copy()
        
        # 添加symbol字段
        result_df.insert(0, 'symbol', result_df['ts_code'].apply(self.convert_tushare_code_to_db))
        
        return result_df

    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}条")
            
            # 获取开盘集合竞价数据
            print(f"正在获取{date_str}开盘集合竞价数据...")
            open_df = self.fetch_open_auction(trade_date)
            if not open_df.empty:
                print(f"成功获取开盘集合竞价数据，共{len(open_df)}只股票")
            else:
                print(f"{date_str}没有开盘集合竞价数据")
            
            # 获取收盘集合竞价数据
            print(f"正在获取{date_str}收盘集合竞价数据...")
            close_df = self.fetch_close_auction(trade_date)
            if not close_df.empty:
                print(f"成功获取收盘集合竞价数据，共{len(close_df)}只股票")
            else:
                print(f"{date_str}没有收盘集合竞价数据")
            
            # 合并数据
            if open_df.empty and close_df.empty:
                print(f"{date_str}没有集合竞价数据（可能是非交易日）")
                return
            
            print(f"正在合并数据...")
            result_df = self.merge_and_transform_data(open_df, close_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(open_df)}只股票" if not open_df.empty else "开盘集合竞价: 无数据")
            print(f"收盘集合竞价: {len(close_df)}只股票" if not close_df.empty else "收盘集合竞价: 无数据")
            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 run_full_collection(self, start_date=None, end_date=None):
        """
        执行全量覆盖采集（按日期逐日采集）
        
        Args:
            start_date: 开始日期，格式YYYY-MM-DD，默认为30天前
            end_date: 结束日期，格式YYYY-MM-DD，默认为今天
        """
        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:
                # 默认采集近4000天历史数据
                start_date_obj = end_date_obj - timedelta(days=4000)
            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("=" * 60)
            
            # 按日期逐日采集
            current_date = start_date_obj
            total_days = (end_date_obj - start_date_obj).days + 1
            processed_days = 0
            success_days = 0
            total_records = 0
            
            while current_date <= end_date_obj:
                date_str = current_date.strftime('%Y-%m-%d')
                
                print(f"\n[{processed_days + 1}/{total_days}] 采集日期: {date_str}")
                
                try:
                    # 使用每日采集方法
                    trade_date = current_date.strftime('%Y%m%d')
                    
                    # 获取开盘集合竞价数据
                    open_df = self.fetch_open_auction(trade_date)
                    
                    # 获取收盘集合竞价数据
                    close_df = self.fetch_close_auction(trade_date)
                    
                    # 合并并保存数据
                    if not open_df.empty or not close_df.empty:
                        result_df = self.merge_and_transform_data(open_df, close_df)
                        
                        if not result_df.empty:
                            # 删除该日期旧数据
                            # with self.engine.begin() as conn:
                            #     delete_query = text(f"DELETE FROM {self.table_name} WHERE DATE(trade_date) = :date_str")
                            #     conn.execute(delete_query, {"date_str": date_str})
                            
                            # 保存新数据
                            result_df.to_sql(self.table_name, self.engine, if_exists='append', index=False)
                            total_records += len(result_df)
                            success_days += 1
                            print(f"  成功: {len(result_df)}条记录")
                    else:
                        print(f"  无数据（可能是非交易日）")

                except Exception as e:
                    print(f"  采集失败: {e}")

                # 移到下一天
                current_date += timedelta(days=1)
                processed_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"成功天数: {success_days}")
            print(f"总记录数: {total_records}条")
            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_auction_data(db_url, tushare_token, mode='daily', date=None, start_date=None, end_date=None):
    """
    股票集合竞价数据采集入口函数
    
    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，默认为30天前
        end_date: 结束日期（仅mode='full'时有效），格式YYYY-MM-DD，默认为今天
    """
    collector = AuctionDataCollector(db_url, tushare_token)
    
    if mode == 'daily':
        collector.run_daily_collection(date)
    elif mode == 'full':
        collector.run_full_collection(start_date, end_date)
    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. 每日增量采集（推荐每日15:30后使用）
    # collect_auction_data(db_url, tushare_token, mode='daily')
    
    # 2. 全量覆盖采集（首次使用或需要重建数据时）
    collect_auction_data(db_url, tushare_token, mode='full')
    
    # 3. 采集指定日期的数据
    # collect_auction_data(db_url, tushare_token, mode='daily', date='2025-11-03')
    
    # 4. 采集指定日期范围的数据
    # collect_auction_data(db_url, tushare_token, mode='full',
    #                      start_date='2024-01-01', end_date='2024-12-31')

