#!/usr/bin/env python3
"""
真实环境原始输出测试脚本

这个脚本用于测试所有提供商在真实环境下的原始输出，
验证各个 Fetcher 能够正确获取和解析模型列表。

使用方法：
1. 设置环境变量（如果需要）：
   set OPENROUTER_API_KEY=your_key
   set SILICONFLOW_API_KEY=your_key
   set IFLOW_API_KEY=your_key

2. 直接运行脚本：
   python TestRawData.py

3. 或使用 pytest 运行：
   pytest TestRawData.py -v
"""

import os
import sys
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Any, Optional

# 添加项目 src 目录到 Python 路径
project_root = Path(__file__).parent.parent.parent
src_path = project_root / "src"
sys.path.insert(0, str(src_path))

from oiiai import (
    FetchZhipu,
    FetchOpenRouter,
    FetchModelScope,
    FetchSiliconFlow,
    FetchIFlow,
)


class RawDataResult:
    """测试结果数据类"""
    
    def __init__(self, provider: str):
        self.provider = provider
        self.success = False
        self.models: List[str] = []
        self.model_count = 0
        self.error: Optional[str] = None
        self.raw_type: Optional[str] = None
        self.duration_ms: float = 0
        self.metadata: Dict[str, Any] = {}
    
    def __repr__(self):
        status = "✅" if self.success else "❌"
        return f"{status} {self.provider}: {self.model_count} models"


def _fetch_zhipu_raw() -> RawDataResult:
    """测试智谱AI原始输出"""
    result = RawDataResult("Zhipu")
    
    import time
    start = time.time()
    
    try:
        fetcher = FetchZhipu()
        models = fetcher.fetch_models()
        
        result.duration_ms = (time.time() - start) * 1000
        result.raw_type = type(models).__name__
        result.models = models
        result.model_count = len(models)
        result.success = len(models) > 0
        
        # 额外元数据
        result.metadata = {
            "sample_models": models[:5] if models else [],
            "has_glm_models": any("glm" in m.lower() for m in models) if models else False,
        }
        
        # 如果没有获取到模型，设置错误信息
        if not result.success:
            result.error = "未能获取到任何模型，可能是页面结构已变化"
        
    except Exception as e:
        result.duration_ms = (time.time() - start) * 1000
        result.error = str(e)
        result.success = False
    
    return result


def _fetch_openrouter_raw() -> RawDataResult:
    """测试OpenRouter原始输出"""
    result = RawDataResult("OpenRouter")
    
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        result.error = "OPENROUTER_API_KEY 环境变量未设置"
        result.metadata["skipped"] = True
        return result
    
    import time
    start = time.time()
    
    try:
        fetcher = FetchOpenRouter(api_key=api_key)
        models = fetcher.fetch_models()
        
        result.duration_ms = (time.time() - start) * 1000
        result.raw_type = type(models).__name__
        result.models = models
        result.model_count = len(models)
        result.success = len(models) > 0
        
        result.metadata = {
            "sample_models": models[:5] if models else [],
        }
        
    except Exception as e:
        result.duration_ms = (time.time() - start) * 1000
        result.error = str(e)
        result.success = False
    
    return result


def _fetch_modelscope_raw() -> RawDataResult:
    """测试ModelScope原始输出"""
    result = RawDataResult("ModelScope")
    
    import time
    start = time.time()
    
    try:
        fetcher = FetchModelScope()
        models = fetcher.fetch_models()
        
        result.duration_ms = (time.time() - start) * 1000
        result.raw_type = type(models).__name__
        result.models = models
        result.model_count = len(models)
        result.success = len(models) > 0
        
        result.metadata = {
            "sample_models": models[:5] if models else [],
        }
        
    except Exception as e:
        result.duration_ms = (time.time() - start) * 1000
        result.error = str(e)
        result.success = False
    
    return result


def _fetch_siliconflow_raw() -> RawDataResult:
    """测试SiliconFlow原始输出"""
    result = RawDataResult("SiliconFlow")
    
    api_key = os.getenv("SILICONFLOW_API_KEY")
    if not api_key:
        result.error = "SILICONFLOW_API_KEY 环境变量未设置"
        result.metadata["skipped"] = True
        return result
    
    import time
    start = time.time()
    
    try:
        fetcher = FetchSiliconFlow()
        models = fetcher.fetch_models()
        
        result.duration_ms = (time.time() - start) * 1000
        result.raw_type = type(models).__name__
        result.models = models
        result.model_count = len(models)
        result.success = len(models) > 0
        
        result.metadata = {
            "sample_models": models[:5] if models else [],
        }
        
    except Exception as e:
        result.duration_ms = (time.time() - start) * 1000
        result.error = str(e)
        result.success = False
    
    return result


def _fetch_iflow_raw() -> RawDataResult:
    """测试IFlow原始输出"""
    result = RawDataResult("IFlow")
    
    api_key = os.getenv("IFLOW_API_KEY")
    if not api_key:
        result.error = "IFLOW_API_KEY 环境变量未设置"
        result.metadata["skipped"] = True
        return result
    
    import time
    start = time.time()
    
    try:
        fetcher = FetchIFlow()
        models = fetcher.fetch_models()
        
        result.duration_ms = (time.time() - start) * 1000
        result.raw_type = type(models).__name__
        result.models = models
        result.model_count = len(models)
        result.success = len(models) > 0
        
        result.metadata = {
            "sample_models": models[:5] if models else [],
        }
        
    except Exception as e:
        result.duration_ms = (time.time() - start) * 1000
        result.error = str(e)
        result.success = False
    
    return result


def run_all_tests() -> Dict[str, RawDataResult]:
    """运行所有提供商的测试"""
    results = {}
    
    print("=" * 60)
    print("真实环境原始输出测试")
    print(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 60)
    
    # 环境变量检查
    print("\n📋 环境变量状态:")
    env_vars = ["OPENROUTER_API_KEY", "SILICONFLOW_API_KEY", "IFLOW_API_KEY"]
    for var in env_vars:
        status = "✅ 已设置" if os.getenv(var) else "❌ 未设置"
        print(f"  {var}: {status}")
    
    print("\n" + "-" * 60)
    print("开始测试各提供商...")
    print("-" * 60)
    
    # 测试各提供商
    fetch_functions = [
        ("Zhipu", _fetch_zhipu_raw),
        ("OpenRouter", _fetch_openrouter_raw),
        ("ModelScope", _fetch_modelscope_raw),
        ("SiliconFlow", _fetch_siliconflow_raw),
        ("IFlow", _fetch_iflow_raw),
    ]
    
    for name, fetch_func in fetch_functions:
        print(f"\n🔍 测试 {name}...")
        result = fetch_func()
        results[name] = result
        
        if result.success:
            print(f"  ✅ 成功: {result.model_count} 个模型 ({result.duration_ms:.0f}ms)")
            if result.metadata.get("sample_models"):
                print(f"  📝 示例: {result.metadata['sample_models'][:3]}")
        elif result.metadata.get("skipped"):
            print(f"  ⏭️  跳过: {result.error}")
        else:
            print(f"  ❌ 失败: {result.error}")
    
    return results


def print_summary(results: Dict[str, RawDataResult]):
    """打印测试汇总"""
    print("\n" + "=" * 60)
    print("测试汇总")
    print("=" * 60)
    
    successful = sum(1 for r in results.values() if r.success)
    skipped = sum(1 for r in results.values() if r.metadata.get("skipped"))
    failed = len(results) - successful - skipped
    
    print(f"\n📊 结果统计:")
    print(f"  ✅ 成功: {successful}")
    print(f"  ⏭️  跳过: {skipped}")
    print(f"  ❌ 失败: {failed}")
    print(f"  📋 总计: {len(results)}")
    
    print(f"\n📝 详细结果:")
    for name, result in results.items():
        if result.success:
            print(f"  ✅ {name}: {result.model_count} 模型 ({result.duration_ms:.0f}ms)")
        elif result.metadata.get("skipped"):
            print(f"  ⏭️  {name}: 跳过 - {result.error}")
        else:
            print(f"  ❌ {name}: 失败 - {result.error}")
    
    # 总模型数
    total_models = sum(r.model_count for r in results.values() if r.success)
    print(f"\n🎯 成功获取的模型总数: {total_models}")


def main():
    """主函数"""
    try:
        results = run_all_tests()
        print_summary(results)
        
        # 返回是否有成功的测试
        successful = sum(1 for r in results.values() if r.success)
        return successful > 0
        
    except KeyboardInterrupt:
        print("\n\n⚠️  用户中断操作")
        return False
    except Exception as e:
        print(f"\n❌ 脚本执行出错: {e}")
        import traceback
        traceback.print_exc()
        return False


# pytest 兼容的测试类
class TestRawDataPytest:
    """pytest 兼容的测试类"""
    
    def test_zhipu_fetcher(self):
        """测试 Zhipu fetcher 能够获取模型列表"""
        result = _fetch_zhipu_raw()
        # Zhipu 不需要 API key，但页面结构可能变化
        # 如果失败，跳过测试而不是断言失败
        if not result.success:
            import pytest
            pytest.skip(f"Zhipu 测试跳过: {result.error}")
    
    def test_modelscope_fetcher(self):
        """测试 ModelScope fetcher 能够获取模型列表"""
        result = _fetch_modelscope_raw()
        # ModelScope 不需要 API key，应该总是能成功
        assert result.success, f"ModelScope 测试失败: {result.error}"
        assert result.model_count > 0, "ModelScope 应该返回至少一个模型"
    
    def test_openrouter_fetcher(self):
        """测试 OpenRouter fetcher（需要 API key）"""
        result = _fetch_openrouter_raw()
        if result.metadata.get("skipped"):
            import pytest
            pytest.skip("OPENROUTER_API_KEY 未设置")
        assert result.success, f"OpenRouter 测试失败: {result.error}"
    
    def test_siliconflow_fetcher(self):
        """测试 SiliconFlow fetcher（需要 API key）"""
        result = _fetch_siliconflow_raw()
        if result.metadata.get("skipped"):
            import pytest
            pytest.skip("SILICONFLOW_API_KEY 未设置")
        assert result.success, f"SiliconFlow 测试失败: {result.error}"
    
    def test_iflow_fetcher(self):
        """测试 IFlow fetcher（需要 API key）"""
        result = _fetch_iflow_raw()
        if result.metadata.get("skipped"):
            import pytest
            pytest.skip("IFLOW_API_KEY 未设置")
        assert result.success, f"IFlow 测试失败: {result.error}"


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)
