# coding:utf-8

"""
美股日线数据采集工具
功能：从Tushare获取美股日线行情数据，存储到us_day_data表
API文档: https://tushare.pro/document/2?doc_id=254
注意：
  1. 该接口需要2000积分权限
  2. 单次最大返回5000条数据
  3. 美股交易时间：北京时间21:30-次日4:00（夏令时）或22:30-次日5:00（冬令时）
  4. 可按日期一次性获取全市场数据
数据存储说明：
  - 存储表：us_day_data（美股专用表）
  - symbol格式：与Tushare一致（如AAPL、TSLA等）
  - PB/PE：从API获取市净率和市盈率数据
  - PS：API中无此数据，保持为NULL
  - turnoverrate：美股接口无换手率数据，保持为NULL
采集策略：
  - 每日增量：按日期获取全市场数据（推荐）
  - 历史数据：按股票逐个采集历史数据
  - 每天美东时间收盘后（北京时间早上5-6点）更新
"""

import tushare as ts
import pandas as pd
from sqlalchemy import create_engine, text
from datetime import datetime, timedelta
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 USDailyDataCollector:
    """美股日线数据采集器"""

    def __init__(self, db_url, tushare_token, table_name='us_day_data'):
        """
        初始化采集器
        
        Args:
            db_url: 数据库连接URL
            tushare_token: Tushare API Token
            table_name: 数据表名，默认为us_day_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 fetch_us_stock_list(self):
        """
        从数据库获取美股代码列表
        
        Returns:
            美股代码列表
        """
        try:
            query = """
                SELECT DISTINCT symbol 
                FROM us_code 
                WHERE symbol IS NOT NULL 
                ORDER BY symbol
            """
            df = pd.read_sql(query, self.engine)
            stock_codes = df['symbol'].tolist()
            print(f"从数据库获取到{len(stock_codes)}只美股代码")
            return stock_codes
        except Exception as e:
            print(f"从数据库获取美股代码失败: {e}")
            # 如果数据库查询失败，返回空列表
            return []

    def fetch_us_daily_data(self, trade_date=None, ts_code=None, start_date=None, end_date=None):
        """
        获取美股日线数据
        
        Args:
            trade_date: 交易日期，格式YYYYMMDD（单日查询）
            ts_code: 股票代码，Tushare格式（单股查询）
            start_date: 开始日期，格式YYYYMMDD（范围查询）
            end_date: 结束日期，格式YYYYMMDD（范围查询）
            
        Returns:
            美股日线数据DataFrame
        """
        try:
            # 构建API参数
            params = {}
            if trade_date:
                params['trade_date'] = trade_date
            if ts_code:
                params['ts_code'] = ts_code
            if start_date:
                params['start_date'] = start_date
            if end_date:
                params['end_date'] = end_date
            
            df = self.pro.us_daily(**params,fields=[
                                                    "ts_code",
                                                    "trade_date",
                                                    "close",
                                                    "open",
                                                    "high",
                                                    "low",
                                                    "pre_close",
                                                    "pct_change",
                                                    "vol",
                                                    "amount",
                                                    "vwap",
                                                    "pe",
                                                    "pb",
                                                    "change",
                                                    "adj_factor",
                                                    "turnover_ratio",
                                                    "total_mv"
                                                ])
            return df
        except Exception as e:
            print(f"获取美股日线数据失败: {e}")
            return pd.DataFrame()

    def transform_data(self, df):
        """
        转换数据格式为数据库格式（与us_day_data表结构一致）
        
        Args:
            df: 原始数据DataFrame
            
        Returns:
            转换后的DataFrame，存储到us_day_data表
            
        字段说明：
            - symbol: 股票代码（如AAPL、TSLA）
            - timestamp: 交易日期
            - volume: 成交量（股）
            - open/high/low/close: 价格（转为字符串）
            - chg: 涨跌额（转为字符串）
            - percent: 涨跌幅（%）
            - turnoverrate: 换手率（美股接口无此数据，设为NULL）
            - amount: 成交额（美元）
            - pb/pe/ps: 估值指标（从接口获取）
            - pre_close: 昨收价
        """
        if df.empty:
            return pd.DataFrame()
        
        # 创建符合us_day_data表结构的DataFrame
        result_df = pd.DataFrame()
        
        # 字段映射和转换
        result_df['symbol'] = df['ts_code']  # 美股代码保持原格式（如AAPL）
        result_df['timestamp'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
        result_df['volume'] = df['vol']  # 成交量（股）
        result_df['open'] = df['open'].astype(str) if 'open' in df.columns else None
        result_df['high'] = df['high'].astype(str) if 'high' in df.columns else None
        result_df['low'] = df['low'].astype(str) if 'low' in df.columns else None
        result_df['close'] = df['close'].astype(str)
        
        # 涨跌额：优先使用change字段，转为字符串
        if 'change' in df.columns:
            result_df['chg'] = df['change'].fillna('').astype(str)
        else:
            result_df['chg'] = None
        
        result_df['percent'] = df.get('pct_change', None) if 'pct_change' in df.columns else None
        
        result_df['turnoverrate'] = None  # 美股接口没有换手率数据
        result_df['amount'] = df.get('amount', None)  # 成交额（如果接口有）
        
        # 估值指标：从API获取
        result_df['pb'] = df.get('pb', None)  # 市净率
        result_df['pe'] = df.get('pe', None)  # 市盈率
        result_df['ps'] = None  # PS字段API中可能没有
        
        result_df['pre_close'] = df.get('pre_close', None)
        
        # 按时间排序
        result_df = result_df.sort_values('timestamp').reset_index(drop=True)
        
        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(`timestamp`) = :date_str")
                result = conn.execute(delete_query, {"date_str": date_str})
                print(f"已删除{self.table_name}表中{date_str}的旧数据，共{result.rowcount}条")
            
            # 获取全市场美股日线数据
            print(f"正在获取{date_str}全市场美股日线数据...")
            df = self.fetch_us_daily_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 run_full_collection(self, start_date=None, end_date=None):
        """
        执行全量历史数据采集（按股票逐个采集）
        
        Args:
            start_date: 开始日期，格式YYYY-MM-DD，默认为180天前
            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:
                # 默认采集近180天历史数据
                start_date_obj = end_date_obj - timedelta(days=180)
            else:
                start_date_obj = datetime.strptime(start_date, '%Y-%m-%d')
            
            start_date_str = start_date_obj.strftime('%Y%m%d')
            end_date_str = end_date_obj.strftime('%Y%m%d')
            
            print(f"日期范围: {start_date_obj.strftime('%Y-%m-%d')} 至 {end_date_obj.strftime('%Y-%m-%d')}")
            print("=" * 60)
            
            # 获取美股代码列表
            stock_codes = self.fetch_us_stock_list()
            
            if not stock_codes:
                print("没有找到美股代码，请先在us_code表中添加美股代码")
                return
            
            print(f"开始采集{len(stock_codes)}只美股的历史数据...")
            print("=" * 60)
            
            # 按股票逐个采集
            success_count = 0
            fail_count = 0
            total_records = 0
            
            for i, symbol in enumerate(stock_codes, 1):
                try:
                    # 获取该股票历史数据
                    df = self.fetch_us_daily_data(
                        ts_code=symbol,
                        start_date=start_date_str,
                        end_date=end_date_str
                    )
                    
                    if not df.empty:
                        # 删除该股票在该日期范围的旧数据
                        with self.engine.begin() as conn:
                            delete_query = text(
                                f"DELETE FROM {self.table_name} "
                                f"WHERE symbol = :symbol "
                                f"AND `timestamp` BETWEEN :start_date AND :end_date"
                            )
                            conn.execute(delete_query, {
                                "symbol": symbol,
                                "start_date": start_date_obj.strftime('%Y-%m-%d'),
                                "end_date": end_date_obj.strftime('%Y-%m-%d')
                            })
                        
                        # 转换数据格式
                        result_df = self.transform_data(df)
                        
                        if not result_df.empty:
                            # 保存到数据库
                            result_df.to_sql(self.table_name, self.engine, if_exists='append', index=False)
                            total_records += len(result_df)
                            success_count += 1
                    
                    # 每10只股票显示一次进度
                    if i % 10 == 0:
                        print(f"进度: {i}/{len(stock_codes)} ({success_count}成功, {fail_count}失败, 共{total_records}条记录)")
                    
                except Exception as e:
                    fail_count += 1
                    if i % 10 == 0:
                        print(f"采集 {symbol} 失败: {e}")
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("全量历史数据采集完成")
            print("=" * 60)
            print(f"日期范围: {start_date_obj.strftime('%Y-%m-%d')} 至 {end_date_obj.strftime('%Y-%m-%d')}")
            print(f"成功采集: {success_count}只")
            print(f"失败: {fail_count}只")
            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 run_date_range_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:
                # 默认采集近30天历史数据
                start_date_obj = end_date_obj - timedelta(days=30)
            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')
                trade_date = current_date.strftime('%Y%m%d')
                
                print(f"\n[{processed_days + 1}/{total_days}] 采集日期: {date_str}")
                
                try:
                    # 获取全市场数据
                    df = self.fetch_us_daily_data(trade_date=trade_date)
                    
                    if not df.empty:
                        # 转换数据格式
                        result_df = self.transform_data(df)
                        
                        if not result_df.empty:
                            # 删除该日期旧数据
                            with self.engine.begin() as conn:
                                delete_query = text(f"DELETE FROM {self.table_name} WHERE DATE(`timestamp`) = :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_us_daily_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': 全量覆盖采集（按股票逐个采集）
            - 'range': 日期范围采集（按日期逐日采集全市场，推荐用于历史数据）
        date: 指定日期（仅mode='daily'时有效），格式YYYY-MM-DD
        start_date: 开始日期（mode='full'或'range'时有效），格式YYYY-MM-DD
        end_date: 结束日期（mode='full'或'range'时有效），格式YYYY-MM-DD
    
    注意：
        - 美股数据保存到us_day_data表（美股专用表）
        - symbol格式与Tushare一致（如AAPL、TSLA等）
        - 每天美东收盘后（北京时间早上5-6点）更新
    """
    collector = USDailyDataCollector(db_url, tushare_token)
    
    if mode == 'daily':
        collector.run_daily_collection(date)
    elif mode == 'full':
        collector.run_full_collection(start_date, end_date)
    elif mode == 'range':
        collector.run_date_range_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. 每日增量采集（推荐每天早上6点后使用，一次性获取全市场数据）
    # collect_us_daily_data(db_url, tushare_token, mode='daily')
    
    # 2. 日期范围采集（推荐用于历史数据，按日期逐日采集全市场）
    # collect_us_daily_data(db_url, tushare_token, mode='range',
    #                       start_date='2024-01-01', end_date='2024-12-31')
    #10月28日的待采集
    # 3. 采集指定日期的数据
    collect_us_daily_data(db_url, tushare_token, mode='daily', date='2025-10-28')
    
    # 4. 全量覆盖采集（按股票逐个采集，较慢，不推荐）
    # collect_us_daily_data(db_url, tushare_token, mode='full',
    #                       start_date='2024-01-01', end_date='2024-12-31')

