#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
测试分批次数据获取功能
验证A股和加密货币数据的大时间范围获取
"""

import sys
import os
sys.path.append('web')

from datetime import datetime, timedelta
from cn_stock_data import ChinaStockData
import ccxt

def test_a_stock_batch_data():
    """测试A股分批次数据获取"""
    print("=== 测试A股分批次数据获取 ===")
    
    china_stock = ChinaStockData()
    
    # 测试获取更长时间范围的数据
    end_date = datetime.now()
    start_date = end_date - timedelta(days=180)  # 6个月数据
    
    print(f"测试时间范围: {start_date.strftime('%Y-%m-%d')} 到 {end_date.strftime('%Y-%m-%d')}")
    
    # 测试不同时间周期
    test_cases = [
        ('600519', '1d', '日线数据'),
        ('600519', '1h', '1小时数据'),
        ('600519', '15m', '15分钟数据'),
    ]
    
    for symbol, timeframe, description in test_cases:
        print(f"\n测试 {description}: {symbol} {timeframe}")
        
        try:
            df = china_stock.get_kl_data(
                symbol=symbol,
                timeframe=timeframe,
                start_date=start_date.strftime('%Y-%m-%d'),
                end_date=end_date.strftime('%Y-%m-%d'),
                limit=5000
            )
            
            if df is not None:
                print(f"✅ 成功获取 {len(df)} 条记录")
                print(f"   时间范围: {df['date'].min()} 到 {df['date'].max()}")
                print(f"   数据列: {list(df.columns)}")
            else:
                print(f"❌ 获取失败")
                
        except Exception as e:
            print(f"❌ 错误: {e}")

def test_crypto_batch_data():
    """测试加密货币分批次数据获取"""
    print("\n=== 测试加密货币分批次数据获取 ===")
    
    # 初始化交易所
    exchange = ccxt.binance({
        'enableRateLimit': True,
    })
    
    # 测试获取更长时间范围的数据
    end_time = datetime.now()
    start_time = end_time - timedelta(days=30)  # 30天数据
    
    print(f"测试时间范围: {start_time} 到 {end_time}")
    
    # 转换为时间戳
    start_timestamp = int(start_time.timestamp() * 1000)
    end_timestamp = int(end_time.timestamp() * 1000)
    
    # 测试不同时间周期
    test_cases = [
        ('BTC/USDT:USDT', '1d', '日线数据'),
        ('BTC/USDT:USDT', '1h', '1小时数据'),
        ('BTC/USDT:USDT', '5m', '5分钟数据'),
    ]
    
    for symbol, timeframe, description in test_cases:
        print(f"\n测试 {description}: {symbol} {timeframe}")
        
        try:
            # 模拟分批次获取逻辑
            all_ohlcv = []
            current_since = start_timestamp
            request_count = 0
            max_requests = 10
            
            batch_size = 500 if timeframe in ['1m', '5m'] else 1000
            
            while request_count < max_requests and current_since < end_timestamp:
                request_count += 1
                print(f"  批次 {request_count}: 获取数据...")
                
                ohlcv = exchange.fetch_ohlcv(symbol, timeframe, since=current_since, limit=batch_size)
                
                if not ohlcv or len(ohlcv) == 0:
                    break
                
                all_ohlcv.extend(ohlcv)
                last_timestamp = ohlcv[-1][0]
                
                if last_timestamp >= end_timestamp:
                    break
                
                if len(ohlcv) < batch_size:
                    break
                
                current_since = last_timestamp + 1
                
                # 防止请求过频
                import time
                time.sleep(0.3)
            
            if all_ohlcv:
                print(f"✅ 成功获取 {len(all_ohlcv)} 条记录 (共 {request_count} 个批次)")
                
                # 时间范围检查
                first_time = datetime.fromtimestamp(all_ohlcv[0][0] / 1000)
                last_time = datetime.fromtimestamp(all_ohlcv[-1][0] / 1000)
                print(f"   时间范围: {first_time} 到 {last_time}")
            else:
                print(f"❌ 获取失败")
                
        except Exception as e:
            print(f"❌ 错误: {e}")

def test_data_quality():
    """测试数据质量"""
    print("\n=== 测试数据质量 ===")
    
    china_stock = ChinaStockData()
    
    # 获取一小段数据进行质量检查
    df = china_stock.get_kl_data(
        symbol='600519',
        timeframe='1d',
        limit=100
    )
    
    if df is not None:
        print(f"数据行数: {len(df)}")
        print(f"数据列: {list(df.columns)}")
        
        # 检查缺失值
        missing_values = df.isnull().sum()
        print(f"缺失值统计:")
        for col, count in missing_values.items():
            if count > 0:
                print(f"  {col}: {count}")
        
        # 检查数据类型
        print(f"数据类型:")
        for col, dtype in df.dtypes.items():
            print(f"  {col}: {dtype}")
        
        # 检查时间连续性
        if len(df) > 1:
            time_diffs = df['date'].diff().dropna()
            print(f"时间间隔统计:")
            print(f"  最小间隔: {time_diffs.min()}")
            print(f"  最大间隔: {time_diffs.max()}")
            print(f"  平均间隔: {time_diffs.mean()}")
        
        # 检查价格合理性
        price_cols = ['open', 'high', 'low', 'close']
        for col in price_cols:
            if col in df.columns:
                print(f"{col} 价格范围: {df[col].min():.2f} - {df[col].max():.2f}")
        
        print("✅ 数据质量检查完成")
    else:
        print("❌ 无法获取数据进行质量检查")

if __name__ == '__main__':
    print("开始测试分批次数据获取功能...\n")
    
    # 测试A股数据
    test_a_stock_batch_data()
    
    # 测试加密货币数据
    test_crypto_batch_data()
    
    # 测试数据质量
    test_data_quality()
    
    print("\n测试完成！")