# coding:utf-8

"""
A股实时行情数据采集工具（基于Tushare）
功能：从Tushare获取A股实时行情数据，存储到Redis
API：ts.realtime_quote() - 新浪实时行情数据
注意：
  1. 免费接口，无积分限制
  2. 单次最多支持50个股票代码
  3. 实时数据，交易时间段可频繁调用
  4. 数据结构与雪球格式对齐，便于业务兼容
数据存储：
  - 存储位置：Redis Hash
  - Key：xq_stock_changes_latest（与雪球脚本保持一致）
  - Field：股票代码（如SZ000978）
  - Value：JSON格式的实时行情数据
采集策略：
  - 从gp_code_all表获取所有A股代码
  - 每49只股票调用一次API（避免超过50个限制）
  - 批量写入Redis，替换旧数据
"""

import tushare as ts
import pandas as pd
import redis
import json
from datetime import datetime
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
from sqlalchemy import create_engine

# Redis配置
try:
    from src.valuation_analysis.config import REDIS_CONFIG
except ImportError:
    REDIS_CONFIG = {
        'host': 'localhost',
        'port': 6379,
        'db': 0,
        'password': None
    }

REDIS_KEY = 'xq_stock_changes_latest'  # 与雪球脚本保持一致


class RealtimeQuoteCollector:
    """A股实时行情数据采集器（基于Tushare）"""

    def __init__(self, db_url, tushare_token, batch_size=49):
        """
        初始化采集器
        
        Args:
            db_url: 数据库连接URL
            tushare_token: Tushare API Token
            batch_size: 每批采集的股票数量，默认49（API最大支持50）
        """
        self.db_url = db_url
        self.batch_size = batch_size
        
        # 初始化Tushare
        ts.set_token(tushare_token)
        
        # 初始化Redis连接
        self.redis_conn = self._get_redis_conn()
        
        print("=" * 60)
        print("A股实时行情数据采集工具（Tushare版）")
        print(f"Redis Key: {REDIS_KEY}")
        print(f"批次大小: {batch_size}只股票/次")
        print("=" * 60)

    def _get_redis_conn(self):
        """获取Redis连接"""
        pool = redis.ConnectionPool(
            host=REDIS_CONFIG['host'],
            port=REDIS_CONFIG['port'],
            db=REDIS_CONFIG.get('db', 0),
            password=REDIS_CONFIG.get('password', None),
            decode_responses=True
        )
        return redis.Redis(connection_pool=pool)

    def convert_db_code_to_tushare(self, db_code):
        """
        将数据库格式的股票代码转换为Tushare格式
        数据库格式：SH600000、SZ000001、BJ430047
        Tushare格式：600000.SH、000001.SZ、430047.BJ
        
        Args:
            db_code: 数据库格式的股票代码
            
        Returns:
            Tushare格式的股票代码
        """
        if db_code.startswith('SH'):
            return f"{db_code[2:]}.SH"
        elif db_code.startswith('SZ'):
            return f"{db_code[2:]}.SZ"
        elif db_code.startswith('BJ'):
            return f"{db_code[2:]}.BJ"
        else:
            return db_code

    def fetch_stock_codes(self):
        """
        从数据库获取所有A股股票代码
        
        Returns:
            股票代码列表（数据库格式，如SH600000）
        """
        try:
            engine = create_engine(self.db_url)
            query = "SELECT gp_code FROM gp_code_all WHERE gp_code IS NOT NULL ORDER BY gp_code"
            df = pd.read_sql(query, engine)
            engine.dispose()
            
            stock_codes = df['gp_code'].tolist()
            print(f"从数据库获取到{len(stock_codes)}只A股代码")
            return stock_codes
        except Exception as e:
            print(f"从数据库获取股票代码失败: {e}")
            return []

    def fetch_realtime_quote(self, ts_codes):
        """
        获取实时行情数据
        
        Args:
            ts_codes: Tushare格式的股票代码列表（如['600000.SH', '000001.SZ']）
            
        Returns:
            实时行情DataFrame
        """
        try:
            # 将列表转换为逗号分隔的字符串
            ts_code_str = ','.join(ts_codes)
            
            # 调用Tushare实时行情接口
            df = ts.realtime_quote(ts_code=ts_code_str)
            
            return df
        except Exception as e:
            print(f"获取实时行情失败: {e}")
            return pd.DataFrame()

    def transform_to_redis_format(self, row, db_code):
        """
        将Tushare实时行情数据转换为Redis存储格式（与雪球格式对齐）
        
        Args:
            row: Tushare实时行情数据的一行（Series）
            db_code: 数据库格式的股票代码（如SH600000）
            
        Returns:
            dict: 转换后的数据字典
        """
        # 创建与雪球格式对齐的数据结构
        result = {
            # 基础信息
            "symbol": db_code,  # 股票代码（数据库格式）
            "name": row.get('NAME', ''),  # 股票名称
            
            # 价格信息
            "current": row.get('PRICE', None),  # 当前价
            "open": row.get('OPEN', None),  # 开盘价
            "high": row.get('HIGH', None),  # 最高价
            "low": row.get('LOW', None),  # 最低价
            "pre_close": row.get('PRE_CLOSE', None),  # 昨收价（对应last_close）
            
            # 涨跌信息
            "chg": None,  # 涨跌额（需要计算）
            "percent": None,  # 涨跌幅（需要计算）
            
            # 成交信息
            "volume": row.get('VOLUME', None),  # 成交量
            "amount": row.get('AMOUNT', None),  # 成交额
            
            # 估值信息（实时行情接口无，设为None）
            "pe_ttm": None,
            "pb": None,
            "pb_ttm": None,
            "ps": None,
            "pcf": None,
            "dividend_yield": None,
            "roe_ttm": None,
            
            # 市值信息（实时行情接口无，设为None）
            "market_capital": None,
            "float_market_capital": None,
            "total_shares": None,
            "float_shares": None,
            
            # 交易信息（实时行情接口无，设为None）
            "turnover_rate": None,
            "volume_ratio": None,
            "amplitude": None,
            
            # 其他信息
            "fetch_time": datetime.now().strftime('%Y-%m-%d %H:%M:%S'),  # 采集时间
            "type": 11,  # 默认类型
            "has_follow": False,
            "tick_size": 0.01,
            "lot_size": 100,
        }
        
        # 计算涨跌额和涨跌幅
        if result['current'] is not None and result['pre_close'] is not None:
            try:
                current = float(result['current'])
                pre_close = float(result['pre_close'])
                
                # 涨跌额
                result['chg'] = round(current - pre_close, 2)
                
                # 涨跌幅（百分比）
                if pre_close != 0:
                    result['percent'] = round((current - pre_close) / pre_close * 100, 2)
            except:
                pass
        
        return result

    def run_collection(self):
        """
        执行实时行情采集
        """
        print("=" * 60)
        print("开始执行A股实时行情采集")
        print("=" * 60)
        
        start_time = time.time()
        
        try:
            # 获取所有股票代码
            db_codes = self.fetch_stock_codes()
            
            if not db_codes:
                print("未获取到股票代码，终止采集")
                return
            
            total_codes = len(db_codes)
            print(f"共需采集 {total_codes} 只股票")
            print(f"将分 {(total_codes + self.batch_size - 1) // self.batch_size} 批次进行")
            print("=" * 60)
            
            # 先采集所有数据到内存，避免中途出错导致数据不完整
            print("开始采集数据（先采集到内存，完成后统一保存）...")
            all_data = {}  # 存储所有采集到的数据 {db_code: data_dict}
            batch_num = 0
            total_collected = 0
            
            for i in range(0, total_codes, self.batch_size):
                batch_num += 1
                batch_db_codes = db_codes[i:i + self.batch_size]
                
                # 转换为Tushare格式
                batch_ts_codes = [self.convert_db_code_to_tushare(code) for code in batch_db_codes]
                
                try:
                    # 获取实时行情
                    df = self.fetch_realtime_quote(batch_ts_codes)
                    
                    if df.empty:
                        print(f"  批次 {batch_num} 未获取到数据")
                        continue
                    
                    # 转换数据并存储到内存
                    # 创建TS_CODE到DB_CODE的映射
                    ts_to_db_map = {
                        self.convert_db_code_to_tushare(db_code): db_code
                        for db_code in batch_db_codes
                    }
                    
                    for _, row in df.iterrows():
                        # 根据TS_CODE匹配对应的DB_CODE
                        ts_code = row.get('TS_CODE', '')
                        if ts_code in ts_to_db_map:
                            db_code = ts_to_db_map[ts_code]
                            
                            # 转换为Redis格式
                            data = self.transform_to_redis_format(row, db_code)
                            
                            # 存储到内存字典
                            all_data[db_code] = data
                            total_collected += 1
                    
                    # 显示进度
                    print(f"  累计进度: {total_collected}/{total_codes} ({total_collected*100//total_codes}%)")
                    
                except Exception as e:
                    print(f"  批次 {batch_num} 采集失败: {e}")
                    continue
                
                # API限频控制
                time.sleep(0.1)
            
            # 采集完成后，一次性保存到Redis
            print("\n" + "=" * 60)
            print(f"数据采集完成，共采集到 {total_collected} 只股票的数据")
            print("开始保存到Redis...")
            print("=" * 60)
            
            if all_data:
                # 使用pipeline批量写入，提高性能
                pipe = self.redis_conn.pipeline()
                
                for db_code, data in all_data.items():
                    pipe.hset(REDIS_KEY, db_code, json.dumps(data, ensure_ascii=False))
                
                # 执行批量写入（这会替换Redis中的旧数据）
                pipe.execute()
                print(f"成功保存 {len(all_data)} 条记录到Redis")
            else:
                print("警告：未采集到任何数据，Redis数据保持不变")
            
            end_time = time.time()
            
            # 输出统计信息
            print("\n" + "=" * 60)
            print("实时行情采集完成")
            print("=" * 60)
            print(f"总股票数: {total_codes}")
            print(f"成功采集: {total_collected}")
            print(f"成功保存: {len(all_data)}")
            print(f"失败: {total_codes - total_collected}")
            print(f"耗时: {end_time - start_time:.2f} 秒")
            print(f"Redis Key: {REDIS_KEY}")
            print("=" * 60)
            
        except Exception as e:
            print(f"采集过程发生错误: {e}")
            import traceback
            traceback.print_exc()

    def get_stock_realtime_info(self, stock_code):
        """
        从Redis获取单只股票的实时行情（兼容接口）
        
        Args:
            stock_code: 股票代码，支持多种格式（SH600000、600000.SH等）
            
        Returns:
            dict or None
        """
        # 统一转换为数据库格式
        if '.' in stock_code:
            code, market = stock_code.split('.')
            if market == 'SH':
                db_code = f'SH{code}'
            elif market == 'SZ':
                db_code = f'SZ{code}'
            elif market == 'BJ':
                db_code = f'BJ{code}'
            else:
                db_code = stock_code
        elif stock_code.startswith(('SH', 'SZ', 'BJ')):
            db_code = stock_code
        else:
            # 默认不处理
            db_code = stock_code
        
        # 从Redis获取
        value = self.redis_conn.hget(REDIS_KEY, db_code)
        if not value:
            return None
        
        try:
            data = json.loads(value)
            return data
        except Exception as e:
            print(f"解析Redis数据失败: {e}")
            return None


def collect_realtime_quote(db_url, tushare_token, batch_size=49):
    """
    A股实时行情采集入口函数
    
    Args:
        db_url: 数据库连接URL
        tushare_token: Tushare API Token
        batch_size: 每批采集的股票数量，默认49
    """
    collector = RealtimeQuoteCollector(db_url, tushare_token, batch_size)
    collector.run_collection()


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
    
    # 执行采集
    collect_realtime_quote(db_url, tushare_token, batch_size=49)