# coding:utf-8

"""
股票技术因子数据采集工具
功能：从Tushare获取股票技术因子数据，存储到gp_stock_factor表
API文档: https://tushare.pro/document/2?doc_id=296
注意：该接口需要至少2000积分权限，单次最大返回10000条数据
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 StockFactorCollector:
    """股票技术因子数据采集器"""

    def __init__(self, db_url, tushare_token, table_name='gp_stock_factor'):
        """
        初始化采集器
        
        Args:
            db_url: 数据库连接URL
            tushare_token: Tushare API Token
            table_name: 数据表名，默认为gp_stock_factor
        """
        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_stock_factor(self, ts_code=None, trade_date=None, start_date=None, end_date=None):
        """
        从Tushare获取股票技术因子数据
        API文档: https://tushare.pro/document/2?doc_id=296
        单次最大返回10000条数据
        
        Args:
            ts_code: Tushare格式的股票代码，为None时获取全市场数据
            trade_date: 交易日期，格式YYYYMMDD（单日查询）
            start_date: 开始日期，格式YYYYMMDD（范围查询）
            end_date: 结束日期，格式YYYYMMDD（范围查询）
            
        Returns:
            技术因子DataFrame
        """
        try:
            df = self.pro.stk_factor(
                ts_code=ts_code,
                trade_date=trade_date,
                start_date=start_date,
                end_date=end_date
            )
            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['close'] = df.get('close', None)
        result_df['open'] = df.get('open', None)
        result_df['high'] = df.get('high', None)
        result_df['low'] = df.get('low', None)
        result_df['pre_close'] = df.get('pre_close', None)
        result_df['change'] = df.get('change', None)
        result_df['pct_change'] = df.get('pct_change', None)
        result_df['vol'] = df.get('vol', None)
        result_df['amount'] = df.get('amount', None)
        
        # 复权因子和复权价格
        result_df['adj_factor'] = df.get('adj_factor', None)
        result_df['open_hfq'] = df.get('open_hfq', None)
        result_df['open_qfq'] = df.get('open_qfq', None)
        result_df['close_hfq'] = df.get('close_hfq', None)
        result_df['close_qfq'] = df.get('close_qfq', None)
        result_df['high_hfq'] = df.get('high_hfq', None)
        result_df['high_qfq'] = df.get('high_qfq', None)
        result_df['low_hfq'] = df.get('low_hfq', None)
        result_df['low_qfq'] = df.get('low_qfq', None)
        result_df['pre_close_hfq'] = df.get('pre_close_hfq', None)
        result_df['pre_close_qfq'] = df.get('pre_close_qfq', None)
        
        # 技术指标
        result_df['macd_dif'] = df.get('macd_dif', None)
        result_df['macd_dea'] = df.get('macd_dea', None)
        result_df['macd'] = df.get('macd', None)
        result_df['kdj_k'] = df.get('kdj_k', None)
        result_df['kdj_d'] = df.get('kdj_d', None)
        result_df['kdj_j'] = df.get('kdj_j', None)
        result_df['rsi_6'] = df.get('rsi_6', None)
        result_df['rsi_12'] = df.get('rsi_12', None)
        result_df['rsi_24'] = df.get('rsi_24', None)
        result_df['boll_upper'] = df.get('boll_upper', None)
        result_df['boll_mid'] = df.get('boll_mid', None)
        result_df['boll_lower'] = df.get('boll_lower', None)
        result_df['cci'] = df.get('cci', None)
        
        return result_df

    def run_full_collection(self, start_date=None, end_date=None):
        """
        执行全量覆盖采集（按日期逐日采集）
        
        Args:
            start_date: 开始日期，格式YYYY-MM-DD，默认为100天前
            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
            
            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:
                    # 使用每日采集方法
                    self.run_daily_collection(date=date_str)
                except Exception as e:
                    print(f"采集 {date_str} 失败: {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"已处理: {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):
        """
        执行每日增量采集
        获取指定日期的全市场技术因子数据
        单次可获取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')
            
            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_stock_factor(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_stock_factor(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，默认为100天前
        end_date: 结束日期（仅mode='full'时有效），格式YYYY-MM-DD，默认为今天
    """
    collector = StockFactorCollector(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. 每日增量采集（推荐每日使用）
    # collect_stock_factor(db_url, tushare_token, mode='daily')
    
    # 2. 全量覆盖采集（首次使用或需要重建数据时）
    # collect_stock_factor(db_url, tushare_token, mode='full')
    
    # 3. 采集指定日期的数据
    collect_stock_factor(db_url, tushare_token, mode='daily', date='2025-11-06')
    
    # 4. 采集指定日期范围的数据
    # collect_stock_factor(db_url, tushare_token, mode='full', 
    #                      start_date='2024-01-01', end_date='2024-12-31')
