"""
测试 FlaskAPIDataSource 的交易日历获取功能
"""

import sys
from pathlib import Path

# 添加项目根目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from datasource.flask_api_source import get_flask_api_source


def test_trading_calendar():
    """测试交易日历获取"""
    print("=" * 60)
    print("测试 FlaskAPIDataSource 交易日历获取")
    print("=" * 60)
    
    source = get_flask_api_source()
    
    # 测试 1: 获取服务信息
    print("\n[测试 1] 获取服务信息")
    info = source.get_service_info()
    if 'error' not in info:
        print(f"✓ 服务名称: {info.get('name', 'N/A')}")
        print(f"✓ API 版本: {info.get('version', 'N/A')}")
    else:
        print(f"✗ 服务信息获取失败: {info['error']}")
    
    # 测试 2: 获取日历信息
    print("\n[测试 2] 获取日历信息")
    cal_info = source.get_calendar_info()
    if 'error' not in cal_info:
        print(f"✓ pandas_market_calendars 已安装: {cal_info.get('pandas_market_calendars_installed')}")
        print(f"✓ 支持的市场: {list(cal_info.get('supported_markets', {}).keys())}")
    else:
        print(f"✗ 日历信息获取失败: {cal_info['error']}")
    
    # 测试 3: A 股交易日历
    print("\n[测试 3] A 股交易日历 (2024-01)")
    dates_a = source.get_trading_calendar('A', '2024-01-01', '2024-01-31')
    if dates_a is not None:
        print(f"✓ 返回 {len(dates_a)} 个交易日")
        print(f"  首个交易日: {dates_a[0].strftime('%Y-%m-%d')}")
        print(f"  最后交易日: {dates_a[-1].strftime('%Y-%m-%d')}")
    else:
        print("✗ A 股交易日历获取失败")
    
    # 测试 4: 美股交易日历
    print("\n[测试 4] 美股交易日历 (2024-01)")
    dates_us = source.get_trading_calendar('US', '2024-01-01', '2024-01-15')
    if dates_us is not None:
        print(f"✓ 返回 {len(dates_us)} 个交易日")
        print(f"  首个交易日: {dates_us[0].strftime('%Y-%m-%d')}")
        print(f"  最后交易日: {dates_us[-1].strftime('%Y-%m-%d')}")
        # 验证马丁·路德·金日（1月15日）是否被排除
        mlk_day = '2024-01-15'
        if mlk_day not in [d.strftime('%Y-%m-%d') for d in dates_us]:
            print(f"  ✓ 正确识别马丁·路德·金日休市")
    else:
        print("✗ 美股交易日历获取失败")
    
    # 测试 5: 港股交易日历
    print("\n[测试 5] 港股交易日历 (2024-01)")
    dates_hk = source.get_trading_calendar('HK', '2024-01-01', '2024-01-15')
    if dates_hk is not None:
        print(f"✓ 返回 {len(dates_hk)} 个交易日")
        print(f"  首个交易日: {dates_hk[0].strftime('%Y-%m-%d')}")
        print(f"  最后交易日: {dates_hk[-1].strftime('%Y-%m-%d')}")
    else:
        print("✗ 港股交易日历获取失败")
    
    # 测试 6: 验证 OHLCV 功能仍然正常
    print("\n[测试 6] 验证 OHLCV 数据获取")
    df = source.fetch('518880.SH', '2024-01-01', '2024-01-10')
    if df is not None and len(df) > 0:
        print(f"✓ 获取 {len(df)} 条 OHLCV 数据")
    else:
        print("✗ OHLCV 数据获取失败")
    
    # 总结
    print("\n" + "=" * 60)
    success_count = sum([
        dates_a is not None,
        dates_us is not None,
        dates_hk is not None,
        df is not None
    ])
    print(f"测试完成: {success_count}/4 通过")
    print("=" * 60)


if __name__ == '__main__':
    test_trading_calendar()