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

"""
A股数据获取测试脚本
"""

import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), 'web'))

from web.cn_stock_data import ChinaStockData
import pandas as pd

def test_a_stock_data():
    """测试A股数据获取功能"""
    print("开始测试A股数据获取功能...")
    
    # 初始化A股数据获取器
    china_stock = ChinaStockData()
    
    # 测试获取热门股票列表
    print("\n1. 测试获取热门股票列表:")
    popular_stocks = china_stock.get_popular_stocks()
    print(f"热门股票数量: {len(popular_stocks)}")
    for i, stock in enumerate(popular_stocks[:5]):
        print(f"  {i+1}. {stock['symbol']} - {stock['name']}")
    
    # 测试获取股票K线数据
    print("\n2. 测试获取股票K线数据:")
    test_symbols = ['000001', '600519', '000858']  # 平安银行、贵州茅台、五粮液
    
    for symbol in test_symbols:
        print(f"\n测试股票: {symbol}")
        
        # 测试日线数据
        print("  获取日线数据...")
        try:
            df_daily = china_stock.get_kl_data(symbol, '1d', limit=100)
            if df_daily is not None:
                print(f"    成功获取 {len(df_daily)} 条日线数据")
                print(f"    时间范围: {df_daily['date'].min()} 到 {df_daily['date'].max()}")
                print(f"    最新价格: {df_daily['close'].iloc[-1]:.2f}")
            else:
                print("    获取日线数据失败")
        except Exception as e:
            print(f"    获取日线数据出错: {e}")
        
        # 测试分钟数据
        print("  获取5分钟数据...")
        try:
            df_5m = china_stock.get_kl_data(symbol, '5m', limit=50)
            if df_5m is not None:
                print(f"    成功获取 {len(df_5m)} 条5分钟数据")
                print(f"    时间范围: {df_5m['date'].min()} 到 {df_5m['date'].max()}")
            else:
                print("    获取5分钟数据失败")
        except Exception as e:
            print(f"    获取5分钟数据出错: {e}")
    
    # 测试获取股票列表
    print("\n3. 测试获取股票列表:")
    try:
        stock_list = china_stock.get_stock_list()
        if stock_list:
            print(f"成功获取 {len(stock_list)} 只股票")
            print("前5只股票:")
            for i, stock in enumerate(stock_list[:5]):
                print(f"  {i+1}. {stock['symbol']} - {stock['name']} - 价格: {stock['price']} - 涨跌幅: {stock['change_pct']}%")
        else:
            print("获取股票列表失败")
    except Exception as e:
        print(f"获取股票列表出错: {e}")
    
    print("\nA股数据获取测试完成!")

if __name__ == '__main__':
    test_a_stock_data()