#!/usr/bin/env python3
"""
测试员工数数据获取功能
"""
import asyncio
import sys
import os
import json

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

from app.services.tushare_client import TushareClient


async def test_employees_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"  # 平安银行
    
    async with TushareClient(token=token) as client:
        try:
            print(f"\n📊 查询股票: {test_ts_code}")
            print("调用 stock_company API...")
            
            # 调用 stock_company API
            data = await client.query(
                api_name="stock_company",
                params={"ts_code": test_ts_code, "limit": 10}
            )
            
            if data:
                print(f"✅ 成功获取 {len(data)} 条记录")
                print("\n返回的数据字段:")
                if data:
                    for key in data[0].keys():
                        print(f"  - {key}")
                
                print("\n员工数相关字段:")
                for row in data:
                    if 'employees' in row:
                        print(f"  ✅ employees: {row.get('employees')}")
                    if 'employee' in row:
                        print(f"  ✅ employee: {row.get('employee')}")
                
                print("\n完整数据示例:")
                print(json.dumps(data[0], indent=2, ensure_ascii=False))
            else:
                print("⚠️  未返回数据")
                
        except Exception as e:
            print(f"❌ 错误: {e}")
            import traceback
            traceback.print_exc()


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