#!/usr/bin/env python3
"""
测试股东数数据获取功能
"""
import asyncio
import sys
import os
import json
from datetime import datetime, timedelta

# 添加项目根目录到Python路径
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'backend'))

from tushare_legacy_client import TushareLegacyClient as TushareClient


async def test_holder_number_data():
    """测试获取股东数数据"""
    print("🧪 测试股东数数据获取...")
    print("=" * 50)
    
    # 从环境变量或配置文件读取 token
    base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
    config_path = os.path.join(base_dir, 'config', 'config.json')
    
    token = os.environ.get('TUSHARE_TOKEN')
    if not token and os.path.exists(config_path):
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
            token = config.get('data_sources', {}).get('tushare', {}).get('api_key')
    
    if not token:
        print("❌ 未找到 Tushare token")
        print("请设置环境变量 TUSHARE_TOKEN 或在 config/config.json 中配置")
        return
    
    print(f"✅ Token 已加载: {token[:10]}...")
    
    # 测试股票代码
    test_ts_code = "000001.SZ"  # 平安银行
    years = 5  # 查询最近5年的数据
    
    # 计算日期范围
    end_date = datetime.now().strftime("%Y%m%d")
    start_date = (datetime.now() - timedelta(days=years * 365)).strftime("%Y%m%d")
    
    async with TushareClient(token=token) as client:
        try:
            print(f"\n📊 查询股票: {test_ts_code}")
            print(f"📅 日期范围: {start_date} 到 {end_date}")
            print("调用 stk_holdernumber API...")
            
            # 调用 stk_holdernumber API
            data = await client.query(
                api_name="stk_holdernumber",
                params={
                    "ts_code": test_ts_code,
                    "start_date": start_date,
                    "end_date": end_date,
                    "limit": 5000
                }
            )
            
            if data:
                print(f"✅ 成功获取 {len(data)} 条记录")
                print("\n返回的数据字段:")
                if data:
                    for key in data[0].keys():
                        print(f"  - {key}")
                
                print("\n股东数数据:")
                print("-" * 60)
                for row in data[:10]:  # 只显示前10条
                    end_date_val = row.get('end_date', 'N/A')
                    holder_num = row.get('holder_num', 'N/A')
                    print(f"  日期: {end_date_val}, 股东数: {holder_num}")
                
                if len(data) > 10:
                    print(f"  ... 还有 {len(data) - 10} 条记录")
                
                print("\n完整数据示例（第一条）:")
                print(json.dumps(data[0], indent=2, ensure_ascii=False))
                
                # 检查是否有 holder_num 字段
                if data and 'holder_num' in data[0]:
                    print("\n✅ 成功获取 holder_num 字段数据")
                else:
                    print("\n⚠️  未找到 holder_num 字段")
                    
            else:
                print("⚠️  未返回数据")
                
        except Exception as e:
            print(f"❌ 错误: {e}")
            import traceback
            traceback.print_exc()


if __name__ == "__main__":
    print("🚀 开始测试股东数数据获取功能\n")
    asyncio.run(test_holder_number_data())
    print("\n" + "=" * 50)
    print("✅ 测试完成")