# coding:utf-8

"""
股票日线数据采集器V3版本 - 基于Tushare数据源
作者: AI Assistant
创建时间: 2025-11-03
功能:
1. 从Tushare采集A股日线行情数据
2. 从Tushare采集每日指标数据
3. 整合两个数据源并存储到gp_day_data表
4. 支持历史数据采集(1800天)
5. 支持增量更新(当日数据)
6. 自动处理除权除息问题
每天15~16点左右更新
"""

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 gc
import sys
import os

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

# 导入配置
from src.scripts.config import TUSHARE_TOKEN


class StockDailyDataCollectorV3:
    """股票日线数据采集器V3版本 - 基于Tushare"""

    def __init__(self, db_url, tushare_token, table_name='gp_day_data', basic_table_name='gp_day_basic'):
        """
        初始化采集器
        
        Args:
            db_url: 数据库连接URL
            tushare_token: Tushare API Token
            table_name: 日线数据表名，默认为gp_day_data（不影响原有业务）
            basic_table_name: 每日指标数据表名，默认为gp_day_basic
        """
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600
        )
        
        # 目标数据表名
        self.table_name = table_name
        self.basic_table_name = basic_table_name
        
        # 初始化Tushare
        ts.set_token(tushare_token)
        self.pro = ts.pro_api()
        
        print("=" * 60)
        print(f"股票日线数据采集器V3 - Tushare版本")
        print(f"日线数据表: {self.table_name}")
        print(f"指标数据表: {self.basic_table_name}")
        print("=" * 60)

    def convert_stock_code_to_tushare(self, symbol):
        """
        将数据库格式的股票代码转换为Tushare格式
        数据库格式：SH600000、SZ000001、BJ430047 等
        Tushare格式：600000.SH、000001.SZ、430047.BJ 等
        
        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:
            # 如果已经是Tushare格式，直接返回
            return symbol

    def convert_tushare_code_to_db(self, ts_code):
        """
        将Tushare格式的股票代码转换为数据库格式
        Tushare格式：600000.SH、000001.SZ、430047.BJ 等
        数据库格式：SH600000、SZ000001、BJ430047 等
        
        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_all_stock_codes(self):
        """
        从数据库获取所有需要采集的A股股票代码
        仅从gp_code_all表获取，指数(gp_code_zs)和港股(gp_code_hk)由其他脚本处理
        
        Returns:
            股票代码列表
        """
        try:
            # 只从gp_code_all获取A股股票代码
            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)}个")
            print(f"注意: 指数和港股由其他采集脚本处理")
            
            return codes
            
        except Exception as e:
            print(f"获取股票代码失败: {e}")
            return []

    def fetch_stock_daily_data(self, ts_code=None, start_date=None, end_date=None, trade_date=None):
        """
        从Tushare获取日线行情数据
        API文档: https://tushare.pro/document/2?doc_id=27
        
        Args:
            ts_code: Tushare格式的股票代码，为None时获取全市场数据
            start_date: 开始日期，格式YYYYMMDD（用于历史数据查询）
            end_date: 结束日期，格式YYYYMMDD（用于历史数据查询）
            trade_date: 交易日期，格式YYYYMMDD（用于单日查询，优先级高于start_date/end_date）
            
        Returns:
            日线行情DataFrame
        """
        try:
            if trade_date:
                # 单日查询：获取指定日期的全市场或单只股票数据
                df = self.pro.daily(ts_code=ts_code, trade_date=trade_date)
            else:
                # 范围查询：用于历史数据采集
                df = self.pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
            
            time.sleep(0.1)  # API限频控制
            return df
        except Exception as e:
            print(f"获取日线行情失败: {e}")
            return pd.DataFrame()

    def fetch_stock_daily_basic(self, ts_code=None, start_date=None, end_date=None, trade_date=None):
        """
        从Tushare获取每日指标数据
        API文档: https://tushare.pro/document/2?doc_id=32
        
        Args:
            ts_code: Tushare格式的股票代码，为None时获取全市场数据
            start_date: 开始日期，格式YYYYMMDD（用于历史数据查询）
            end_date: 结束日期，格式YYYYMMDD（用于历史数据查询）
            trade_date: 交易日期，格式YYYYMMDD（用于单日查询，优先级高于start_date/end_date）
            
        Returns:
            每日指标DataFrame
        """
        try:
            if trade_date:
                # 单日查询：获取指定日期的全市场或单只股票数据
                df = self.pro.daily_basic(
                    ts_code=ts_code,
                    trade_date=trade_date,
                    fields='ts_code,trade_date,turnover_rate,turnover_rate_f,volume_ratio,pe,pe_ttm,pb,ps,ps_ttm,dv_ratio,dv_ttm,total_share,float_share,free_share,total_mv,circ_mv'
                )
            else:
                # 范围查询：用于历史数据采集
                df = self.pro.daily_basic(
                    ts_code=ts_code,
                    start_date=start_date,
                    end_date=end_date,
                    fields='ts_code,trade_date,turnover_rate,turnover_rate_f,volume_ratio,pe,pe_ttm,pb,ps,ps_ttm,dv_ratio,dv_ttm,total_share,float_share,free_share,total_mv,circ_mv'
                )
            
            time.sleep(0.1)  # API限频控制
            return df
        except Exception as e:
            print(f"获取每日指标失败: {e}")
            return pd.DataFrame()

    def merge_and_transform_data(self, df_daily, df_basic, symbol=None):
        """
        合并日线行情和每日指标数据，并转换为数据库格式
        
        Args:
            df_daily: 日线行情DataFrame
            df_basic: 每日指标DataFrame
            symbol: 数据库格式的股票代码（单只股票时使用），为None时从ts_code转换
            
        Returns:
            转换后的DataFrame
        """
        if df_daily.empty:
            return pd.DataFrame()
        
        # 合并两个数据源
        if not df_basic.empty:
            df = pd.merge(
                df_daily,
                df_basic[['ts_code', 'trade_date', 'turnover_rate', 'pe_ttm', 'pb', 'ps']],
                on=['ts_code', 'trade_date'],
                how='left'
            )
        else:
            df = df_daily.copy()
            # 如果没有指标数据，添加空列
            df['turnover_rate'] = None
            df['pe_ttm'] = None
            df['pb'] = None
            df['ps'] = None
        
        # 创建符合gp_day_data表结构的DataFrame
        result_df = pd.DataFrame()
        
        # 字段映射和转换
        result_df['timestamp'] = pd.to_datetime(df['trade_date'].values, format='%Y%m%d')
        result_df['volume'] = df['vol'].values * 100  # Tushare的vol单位是手(100股)，需要转换为股
        result_df['open'] = df['open'].values
        result_df['high'] = df['high'].values
        result_df['low'] = df['low'].values
        result_df['close'] = df['close'].values
        result_df['chg'] = df['change'].values
        result_df['percent'] = df['pct_chg'].values
        result_df['turnoverrate'] = df['turnover_rate'].values
        result_df['amount'] = df['amount'].values * 1000  # Tushare的amount单位是千元，需要转换为元
        result_df['pb'] = df['pb'].values
        result_df['pe'] = df['pe_ttm'].values
        result_df['ps'] = df['ps'].values
        result_df['pre_close'] = df['pre_close'].values
        
        # symbol字段处理：如果传入了symbol则使用，否则从ts_code转换
        if symbol is not None:
            result_df['symbol'] = symbol  # 单只股票：标量值会自动广播到所有行
        else:
            # 批量数据：从ts_code转换为数据库格式
            result_df['symbol'] = df['ts_code'].apply(self.convert_tushare_code_to_db)
        
        # 按时间排序
        result_df = result_df.sort_values('timestamp').reset_index(drop=True)
        
        return result_df

    def transform_basic_data(self, df_daily, df_basic, symbol=None):
        """
        转换每日指标数据为数据库格式
        
        Args:
            df_daily: 日线行情DataFrame（用于获取收盘价）
            df_basic: 每日指标DataFrame
            symbol: 数据库格式的股票代码（单只股票时使用），为None时从ts_code转换
            
        Returns:
            转换后的DataFrame
        """
        if df_basic.empty:
            return pd.DataFrame()
        
        # 如果有日线数据，合并获取收盘价
        if not df_daily.empty:
            df = pd.merge(
                df_basic,
                df_daily[['ts_code', 'trade_date', 'close']],
                on=['ts_code', 'trade_date'],
                how='left'
            )
        else:
            df = df_basic.copy()
            df['close'] = None
        
        # 创建符合gp_day_basic表结构的DataFrame
        basic_df = pd.DataFrame()
        
        # 字段映射和转换
        basic_df['timestamp'] = pd.to_datetime(df['trade_date'].values, format='%Y%m%d')
        basic_df['close'] = df['close'].values
        basic_df['turnover_rate'] = df['turnover_rate'].values
        basic_df['turnover_rate_f'] = df.get('turnover_rate_f', pd.Series([None] * len(df))).values
        basic_df['volume_ratio'] = df.get('volume_ratio', pd.Series([None] * len(df))).values
        basic_df['pe'] = df.get('pe', pd.Series([None] * len(df))).values
        basic_df['pe_ttm'] = df['pe_ttm'].values
        basic_df['pb'] = df['pb'].values
        basic_df['ps'] = df['ps'].values
        basic_df['ps_ttm'] = df.get('ps_ttm', pd.Series([None] * len(df))).values
        basic_df['dv_ratio'] = df.get('dv_ratio', pd.Series([None] * len(df))).values
        basic_df['dv_ttm'] = df.get('dv_ttm', pd.Series([None] * len(df))).values
        basic_df['total_share'] = df.get('total_share', pd.Series([None] * len(df))).values
        basic_df['float_share'] = df.get('float_share', pd.Series([None] * len(df))).values
        basic_df['free_share'] = df.get('free_share', pd.Series([None] * len(df))).values
        basic_df['total_mv'] = df.get('total_mv', pd.Series([None] * len(df))).values
        basic_df['circ_mv'] = df.get('circ_mv', pd.Series([None] * len(df))).values
        
        # symbol字段处理
        if symbol is not None:
            basic_df['symbol'] = symbol
        else:
            basic_df['symbol'] = df['ts_code'].apply(self.convert_tushare_code_to_db)
        
        # 按时间排序
        basic_df = basic_df.sort_values('timestamp').reset_index(drop=True)
        
        return basic_df

    def fetch_historical_data(self, symbol, days=1800):
        """
        获取单只股票近N天的历史数据
        
        Args:
            symbol: 数据库格式的股票代码
            days: 获取的天数，默认1800天
            
        Returns:
            是否成功
        """
        try:
            # 转换股票代码格式
            ts_code = self.convert_stock_code_to_tushare(symbol)
            
            # 计算日期范围
            end_date = datetime.now().strftime('%Y%m%d')
            start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')
            
            # 获取日线行情数据
            df_daily = self.fetch_stock_daily_data(ts_code, start_date, end_date)
            if df_daily.empty:
                print(f"未获取到{symbol}的日线行情数据")
                return False
            
            # 获取每日指标数据
            df_basic = self.fetch_stock_daily_basic(ts_code, start_date, end_date)
            
            # 合并并转换数据
            df_result = self.merge_and_transform_data(df_daily, df_basic, symbol)
            
            if df_result.empty:
                print(f"数据转换失败: {symbol}")
                return False
            
            # 保存到日线数据表（使用配置的表名）
            df_result.to_sql(self.table_name, self.engine, if_exists='append', index=False)
            print(f"成功保存{symbol}的历史数据到{self.table_name}表，共{len(df_result)}条记录")
            
            # 转换并保存每日指标数据
            basic_result = self.transform_basic_data(df_daily, df_basic, symbol)
            if not basic_result.empty:
                basic_result.to_sql(self.basic_table_name, self.engine, if_exists='append', index=False)
                print(f"成功保存{symbol}的指标数据到{self.basic_table_name}表，共{len(basic_result)}条记录")
            
            return True
            
        except Exception as e:
            print(f"获取{symbol}历史数据失败: {e}")
            return False

    def fetch_today_data(self, symbol, target_date=None):
        """
        获取单只股票指定日期的数据
        
        Args:
            symbol: 数据库格式的股票代码
            target_date: 目标日期，格式YYYY-MM-DD，默认为今天
            
        Returns:
            转换后的DataFrame，失败返回空DataFrame
        """
        try:
            # 转换股票代码格式
            ts_code = self.convert_stock_code_to_tushare(symbol)
            
            # 确定目标日期
            if target_date:
                # 将YYYY-MM-DD格式转换为YYYYMMDD
                date_obj = datetime.strptime(target_date, '%Y-%m-%d')
                trade_date = date_obj.strftime('%Y%m%d')
            else:
                # 默认为今天
                trade_date = datetime.now().strftime('%Y%m%d')
            
            # 获取日线行情数据（使用同一天作为起止日期）
            df_daily = self.fetch_stock_daily_data(ts_code, trade_date, trade_date)
            if df_daily.empty:
                return pd.DataFrame()
            
            # 获取每日指标数据（使用同一天作为起止日期）
            df_basic = self.fetch_stock_daily_basic(ts_code, trade_date, trade_date)
            
            # 合并并转换数据
            df_result = self.merge_and_transform_data(df_daily, df_basic, symbol)
            
            return df_result
            
        except Exception as e:
            print(f"获取{symbol}在{target_date or '今日'}的数据失败: {e}")
            return pd.DataFrame()

    def check_ex_rights_before_save(self, df, date_str):
        """
        在保存数据库之前检查除权情况，返回除权股票列表和除权日志数据
        
        Args:
            df: 待保存的日线数据DataFrame
            date_str: 当前采集日期，格式YYYY-MM-DD
            
        Returns:
            tuple: (除权股票列表, 除权日志数据列表)
        """
        print("正在检查除权情况...")
        
        ex_rights_stocks = []
        ex_rights_log_data = []
        
        for _, row in tqdm(df.iterrows(), total=len(df), desc="检查除权"):
            symbol = row['symbol']
            current_pre_close = row['pre_close']
            
            # 如果pre_close为空，跳过
            if pd.isna(current_pre_close):
                continue
            
            # 查询数据库中该股票的最近两条收盘价记录
            query = text("""
                SELECT `close`, `timestamp` FROM {table_name}
                WHERE symbol = :symbol 
                ORDER BY `timestamp` DESC 
                LIMIT 2
            """.format(table_name=self.table_name))
            
            try:
                with self.engine.connect() as conn:
                    results = conn.execute(query, {"symbol": symbol}).fetchall()
                    
                if results:
                    # 检查最新记录是否为今天的数据
                    latest_record = results[0]
                    latest_timestamp = latest_record[1]
                    latest_date_str = latest_timestamp.strftime('%Y-%m-%d')
                    
                    if latest_date_str == date_str and len(results) > 1:
                        # 如果最新记录是今天的，且有第二条记录，则用第二条记录比较
                        db_last_close = float(results[1][0])
                    else:
                        # 如果最新记录不是今天的，或者只有一条记录，则用最新记录比较
                        db_last_close = float(results[0][0])
                    
                    # 比较pre_close和数据库中的收盘价
                    if abs(db_last_close - current_pre_close) > 0.001:
                        print(f"发现除权股票: {symbol}, 数据库收盘价: {db_last_close}, 当前昨收价: {current_pre_close}")
                        ex_rights_stocks.append(symbol)
                        
                        # 收集除权日志数据
                        ex_rights_log_data.append({
                            'stock_code': symbol,
                            'change_date': date_str,
                            'before_price': db_last_close,
                            'after_price': current_pre_close,
                            'update_time': datetime.now()
                        })
            except Exception as e:
                # 如果查询失败（比如股票是新上市的），跳过
                continue
        
        if ex_rights_stocks:
            print(f"检测到 {len(ex_rights_stocks)} 只除权股票: {ex_rights_stocks}")
        else:
            print("未发现除权股票")
            
        return ex_rights_stocks, ex_rights_log_data

    def save_ex_rights_log(self, log_data):
        """
        将除权日志保存到数据库
        
        Args:
            log_data: 除权日志数据列表
        """
        if not log_data:
            return
        
        print(f"正在保存 {len(log_data)} 条除权日志到gp_ex_rights_log表...")
        try:
            df = pd.DataFrame(log_data)
            df.to_sql('gp_ex_rights_log', self.engine, if_exists='append', index=False)
            print("除权日志保存成功")
        except Exception as e:
            print(f"保存除权日志失败: {e}")

    def handle_ex_rights_stocks(self, ex_rights_stocks, ex_rights_log_data, days=1800):
        """
        处理除权股票：保存日志、删除历史日线数据并重新获取
        注意：每日指标数据不需要重新获取，因为除权前后这些基础数据是一致的
        
        Args:
            ex_rights_stocks: 除权股票列表
            ex_rights_log_data: 除权日志数据列表
            days: 重新获取的历史天数，默认1800天
        """
        if not ex_rights_stocks:
            return
            
        print("=" * 60)
        print("开始处理除权股票...")
        print("=" * 60)
        
        # 1. 保存除权日志
        if ex_rights_log_data:
            self.save_ex_rights_log(ex_rights_log_data)
        
        # 2. 重新获取历史日线数据（不包括每日指标数据）
        print(f"开始处理 {len(ex_rights_stocks)} 只除权股票，重新获取历史日线数据...")
        
        success_count = 0
        fail_count = 0
        
        for symbol in tqdm(ex_rights_stocks, desc="处理除权股票"):
            try:
                # 转换股票代码格式
                ts_code = self.convert_stock_code_to_tushare(symbol)
                
                # 删除该股票的所有历史日线数据（仅删除日线数据，不删除指标数据）
                delete_query = text(f"DELETE FROM {self.table_name} WHERE symbol = :symbol")
                with self.engine.begin() as conn:
                    result = conn.execute(delete_query, {"symbol": symbol})
                    print(f"已删除 {symbol} 的 {result.rowcount} 条历史日线数据")
                
                # 计算日期范围
                end_date = datetime.now().strftime('%Y%m%d')
                start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')
                
                # 重新获取历史日线数据
                df_daily = self.fetch_stock_daily_data(ts_code, start_date, end_date)
                if df_daily.empty:
                    print(f"未获取到{symbol}的日线行情数据")
                    fail_count += 1
                    continue
                
                # 获取每日指标数据（用于合并，获取turnover_rate等字段）
                df_basic = self.fetch_stock_daily_basic(ts_code, start_date, end_date)
                
                # 合并并转换数据
                df_result = self.merge_and_transform_data(df_daily, df_basic, symbol)
                
                if df_result.empty:
                    print(f"数据转换失败: {symbol}")
                    fail_count += 1
                    continue
                
                # 保存到日线数据表（只保存日线数据，不更新指标数据）
                df_result.to_sql(self.table_name, self.engine, if_exists='append', index=False)
                print(f"成功重新获取 {symbol} 的历史日线数据，共{len(df_result)}条记录")
                success_count += 1
                    
            except Exception as e:
                print(f"处理除权股票 {symbol} 失败: {e}")
                fail_count += 1
        
        print("\n" + "=" * 60)
        print("除权股票处理完成")
        print("=" * 60)
        print(f"成功: {success_count}只")
        print(f"失败: {fail_count}只")
        print("=" * 60)

    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})")
            
            # 删除指定日期的旧数据（日线数据表和指标数据表）
            # 使用timestamp范围查询代替DATE()函数，避免全表扫描，提升性能
            # 添加LENGTH(symbol) = 8限定条件，只删除A股数据（避免误删港股等其他数据）
            delete_start = f"{date_str} 00:00:00"
            delete_end = f"{date_str} 23:59:59"
            
            with self.engine.begin() as conn:
                # 删除日线数据表的旧数据（只删除A股，symbol长度为8，如SH600000）
                delete_query_day = text(f"DELETE FROM {self.table_name} WHERE `timestamp` >= :start_time AND `timestamp` <= :end_time AND LENGTH(symbol) = 8")
                result_day = conn.execute(delete_query_day, {"start_time": delete_start, "end_time": delete_end})
                print(f"已删除{self.table_name}表中{date_str}的A股旧数据，共{result_day.rowcount}条")
                
                # 删除指标数据表的旧数据（只删除A股）
                delete_query_basic = text(f"DELETE FROM {self.basic_table_name} WHERE `timestamp` >= :start_time AND `timestamp` <= :end_time AND LENGTH(symbol) = 8")
                result_basic = conn.execute(delete_query_basic, {"start_time": delete_start, "end_time": delete_end})
                print(f"已删除{self.basic_table_name}表中{date_str}的A股旧数据，共{result_basic.rowcount}条")
            
            # 批量获取全市场日线数据（不传ts_code，获取所有股票）
            print(f"正在获取{date_str}全市场日线行情数据...")
            df_daily = self.fetch_stock_daily_data(trade_date=trade_date)
            
            if df_daily.empty:
                print(f"{date_str}没有交易数据（可能是非交易日）")
                return
            
            print(f"成功获取日线行情数据，共{len(df_daily)}只股票")
            
            # 批量获取全市场每日指标数据
            print(f"正在获取{date_str}全市场每日指标数据...")
            df_basic = self.fetch_stock_daily_basic(trade_date=trade_date)
            print(f"成功获取每日指标数据，共{len(df_basic)}只股票")
            
            # 合并并转换数据（不传symbol参数，自动从ts_code转换）
            print("正在合并和转换数据...")
            result_df = self.merge_and_transform_data(df_daily, df_basic)
            
            if result_df.empty:
                print("数据转换失败")
                return
            
            # 检查除权（保存前）
            ex_rights_stocks, ex_rights_log_data = self.check_ex_rights_before_save(result_df, date_str)
            
            # 保存到日线数据表
            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(f"正在保存指标数据到{self.basic_table_name}表...")
            basic_result = self.transform_basic_data(df_daily, df_basic)
            if not basic_result.empty:
                basic_result.to_sql(self.basic_table_name, self.engine, if_exists='append', index=False)
                print(f"成功保存{len(basic_result)}只股票的指标数据")
            else:
                print("警告：指标数据为空")
            
            # 处理除权股票（保存后）
            self.handle_ex_rights_stocks(ex_rights_stocks, ex_rights_log_data)
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("每日数据采集完成")
            print("=" * 60)
            print(f"采集日期: {date_str}")
            print(f"成功采集: {len(result_df)}只股票")
            print(f"除权处理: {len(ex_rights_stocks)}只股票")
            print(f"日线数据表: {self.table_name}")
            print(f"指标数据表: {self.basic_table_name}")
            print("=" * 60)
            
        except Exception as e:
            print(f"每日数据采集失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            self.engine.dispose()
            gc.collect()

    def run_historical_collection(self, days=1800):
        """
        执行历史数据批量采集
        
        Args:
            days: 获取的天数，默认1800天
        """
        print("=" * 60)
        print(f"开始执行历史数据批量采集 - 近{days}天")
        print("=" * 60)
        
        try:
            # 获取所有股票代码
            stock_codes = self.fetch_all_stock_codes()
            if not stock_codes:
                print("未获取到股票代码，终止采集")
                return
            
            # 统计信息
            success_count = 0
            failed_count = 0
            
            # 遍历所有股票
            for symbol in tqdm(stock_codes, desc=f"采集近{days}天历史数据"):
                try:
                    # 先删除该股票的历史数据（日线数据表和指标数据表）
                    with self.engine.begin() as conn:
                        # 删除日线数据
                        delete_query_day = text(f"DELETE FROM {self.table_name} WHERE symbol = :symbol")
                        conn.execute(delete_query_day, {"symbol": symbol})
                        
                        # 删除指标数据
                        delete_query_basic = text(f"DELETE FROM {self.basic_table_name} WHERE symbol = :symbol")
                        conn.execute(delete_query_basic, {"symbol": symbol})
                    
                    # 获取历史数据
                    if self.fetch_historical_data(symbol, days):
                        success_count += 1
                    else:
                        failed_count += 1
                    
                    # 休息一下，避免API限频
                    # time.sleep(0.2)
                    
                except Exception as e:
                    print(f"处理{symbol}时发生错误: {e}")
                    failed_count += 1
                    continue
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("历史数据采集完成")
            print("=" * 60)
            print(f"采集天数: {days}")
            print(f"总股票数: {len(stock_codes)}")
            print(f"成功采集: {success_count}")
            print(f"失败: {failed_count}")
            print("=" * 60)
            
        except Exception as e:
            print(f"历史数据采集失败: {e}")
        finally:
            self.engine.dispose()
            gc.collect()

    def refetch_single_stock_history(self, symbol, days=1800):
        """
        重新获取单只股票的历史数据（用于除权处理）
        
        Args:
            symbol: 数据库格式的股票代码
            days: 获取的天数，默认1800天
            
        Returns:
            是否成功
        """
        try:
            # 删除该股票的历史数据（日线数据表和指标数据表）
            with self.engine.begin() as conn:
                # 删除日线数据
                delete_query_day = text(f"DELETE FROM {self.table_name} WHERE symbol = :symbol")
                conn.execute(delete_query_day, {"symbol": symbol})
                
                # 删除指标数据
                delete_query_basic = text(f"DELETE FROM {self.basic_table_name} WHERE symbol = :symbol")
                conn.execute(delete_query_basic, {"symbol": symbol})
            
            print(f"已删除{self.table_name}和{self.basic_table_name}表中{symbol}的历史数据")
            
            # 重新获取历史数据
            success = self.fetch_historical_data(symbol, days)
            if success:
                print(f"成功重新获取{symbol}的历史数据")
            else:
                print(f"重新获取{symbol}的历史数据失败")
            
            return success
            
        except Exception as e:
            print(f"重新获取{symbol}历史数据失败: {e}")
            return False


def collect_stock_daily_data_v3(db_url, tushare_token, mode='daily', date=None, days=1800, table_name='gp_day_data', basic_table_name='gp_day_basic'):
    """
    V3版本的股票日线数据采集入口函数
    
    Args:
        db_url: 数据库连接URL
        tushare_token: Tushare API Token
        mode: 采集模式，'daily'=增量采集当日数据，'historical'=批量采集历史数据
        date: 指定日期（仅mode='daily'时有效），格式YYYY-MM-DD
        days: 历史数据天数（仅mode='historical'时有效），默认1800天
        table_name: 日线数据表名，默认为gp_day_data（不影响原有业务）
        basic_table_name: 每日指标数据表名，默认为gp_day_basic
    """
    collector = StockDailyDataCollectorV3(db_url, tushare_token, table_name, basic_table_name)
    
    if mode == 'daily':
        collector.run_daily_collection(date)
    elif mode == 'historical':
        collector.run_historical_collection(days)
    else:
        print(f"未知的采集模式: {mode}")


def get_stock_price_data_from_db(stock_code, db_url='mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj'):
    """
    从gp_day_data和gp_day_basic表获取股票最新价格数据
    替代原StockPriceCollector.get_stock_price_data()方法
    
    Args:
        stock_code: 股票代码，支持多种格式（SH600000、600000.SH、600000等）
        db_url: 数据库连接URL
        
    Returns:
        dict: 包含股票实时价格信息的字典，格式与原接口兼容
        {
            'stock_code': '600000',
            'stock_name': '浦发银行',
            'current_price': 10.5,
            'change_percent': 2.5,
            'change_amount': 0.25,
            'volume': 1234567,
            'amount': 12345678.9,
            'high': 10.8,
            'low': 10.2,
            'open': 10.3,
            'pre_close': 10.25,
            'update_time': '2024-01-20 14:30:00',
            'total_market_value': 100000000.0,  # 总市值（万元）
            'pe': 8.5,
            'pb': 0.6,
            'ps': 1.2
        }
    """
    try:
        # 统一股票代码格式为数据库格式（如SH600000）
        db_code = _normalize_stock_code(stock_code)
        
        # 创建数据库连接
        engine = create_engine(db_url)
        
        # 从gp_day_data获取最新日线数据
        day_query = f"""
            SELECT 
                symbol,
                timestamp,
                open,
                high,
                low,
                close,
                chg,
                percent,
                volume,
                amount,
                pre_close
            FROM gp_day_data
            WHERE symbol = '{db_code}'
            ORDER BY timestamp DESC
            LIMIT 1
        """
        
        day_df = pd.read_sql(day_query, engine)
        
        if day_df.empty:
            engine.dispose()
            return None
        
        # 从gp_day_basic获取最新指标数据
        basic_query = f"""
            SELECT 
                pe_ttm,
                pb,
                ps,
                total_mv,
                circ_mv
            FROM gp_day_basic
            WHERE symbol = '{db_code}'
            ORDER BY timestamp DESC
            LIMIT 1
        """
        
        basic_df = pd.read_sql(basic_query, engine)
        engine.dispose()
        
        # 组装返回数据
        day_row = day_df.iloc[0]
        
        result = {
            'stock_code': stock_code,  # 保持原格式
            'stock_name': None,  # 数据库中没有股票名称，保持为None
            'current_price': float(day_row['close']) if day_row['close'] else None,
            'change_percent': float(day_row['percent']) if day_row['percent'] else None,
            'change_amount': float(day_row['chg']) if day_row['chg'] else None,
            'volume': int(day_row['volume']) if day_row['volume'] else None,
            'amount': float(day_row['amount']) if day_row['amount'] else None,
            'high': float(day_row['high']) if day_row['high'] else None,
            'low': float(day_row['low']) if day_row['low'] else None,
            'open': float(day_row['open']) if day_row['open'] else None,
            'pre_close': float(day_row['pre_close']) if day_row['pre_close'] else None,
            'update_time': day_row['timestamp'].strftime('%Y-%m-%d %H:%M:%S') if day_row['timestamp'] else None,
        }
        
        # 添加估值指标（如果有）
        if not basic_df.empty:
            basic_row = basic_df.iloc[0]
            result['total_market_value'] = float(basic_row['total_mv']) * 10000 if basic_row['total_mv'] else None  # 转换为元
            result['pe'] = float(basic_row['pe_ttm']) if basic_row['pe_ttm'] else None
            result['pb'] = float(basic_row['pb']) if basic_row['pb'] else None
            result['ps'] = float(basic_row['ps']) if basic_row['ps'] else None
        else:
            result['total_market_value'] = None
            result['pe'] = None
            result['pb'] = None
            result['ps'] = None
        
        return result
        
    except Exception as e:
        print(f"从数据库获取股票价格数据失败: {e}")
        import traceback
        traceback.print_exc()
        return None


def _normalize_stock_code(stock_code):
    """
    统一股票代码格式为数据库格式
    支持的输入格式：
    - SH600000、SZ000001 （数据库格式）
    - 600000.SH、000001.SZ （Tushare格式）
    - 600000、000001 （纯数字）
    
    Args:
        stock_code: 股票代码
        
    Returns:
        数据库格式的股票代码（如SH600000）
    """
    stock_code = stock_code.upper().strip()
    
    # 已经是数据库格式
    if stock_code.startswith(('SH', 'SZ', 'BJ')):
        return stock_code
    
    # Tushare格式转换
    if '.' in stock_code:
        code, market = stock_code.split('.')
        if market == 'SH':
            return f'SH{code}'
        elif market == 'SZ':
            return f'SZ{code}'
        elif market == 'BJ':
            return f'BJ{code}'
    
    # 纯数字，根据规则判断
    if stock_code.isdigit():
        if stock_code.startswith('6'):
            return f'SH{stock_code}'
        elif stock_code.startswith(('0', '3')):
            return f'SZ{stock_code}'
        elif stock_code.startswith('8') or stock_code.startswith('4'):
            return f'BJ{stock_code}'
    
    # 无法识别，返回原值
    return stock_code


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_stock_daily_data_v3(db_url, tushare_token, mode='daily')
    
    # 2. 批量采集历史数据（首次使用或需要重建数据时）
    # collect_stock_daily_data_v3(db_url, tushare_token, mode='historical', days=4000)
    
    # 3. 采集指定日期的数据
    collect_stock_daily_data_v3(db_url, tushare_token, mode='daily', date='2025-11-07')
    
    # 4. 单只股票历史数据重新获取（用于除权处理）
    # collector = StockDailyDataCollectorV3(db_url, tushare_token)
    # collector.refetch_single_stock_history('SH600000', 1800)

