#!/usr/bin/env python3
"""
捕获所有提供商的原始输出脚本

这个脚本会调用每个提供商的API，并保存原始响应数据到JSON文件中。
用于分析不同提供商的返回格式，以便设计统一的返回体格式。

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

2. 运行脚本：
   python capture_raw_outputs.py
"""

import os
import json
import time
import sys
from datetime import datetime
from pathlib import Path

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

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


def save_raw_output(provider_name, data, metadata=None):
    """保存原始输出到文件"""
    output_dir = Path(__file__).parent
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    filename = f"{provider_name}_raw_{timestamp}.json"
    filepath = output_dir / filename
    
    output_data = {
        "provider": provider_name,
        "timestamp": timestamp,
        "metadata": metadata or {},
        "raw_data": data
    }
    
    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(output_data, f, ensure_ascii=False, indent=2)
    
    print(f"✅ {provider_name}: 原始输出已保存到 {filename}")
    return filepath


def capture_zhipu():
    """捕获智谱AI的原始输出"""
    print("🔍 开始捕获智谱AI原始输出...")
    try:
        fetcher = FetchZhipu()
        
        # 为了获取原始HTML，我们需要修改或直接调用内部方法
        # 这里我们先尝试调用fetch_models，然后尝试获取更多原始数据
        result = fetcher.fetch_models()
        
        metadata = {
            "result_type": type(result).__name__,
            "categories_count": len(result) if isinstance(result, dict) else "N/A",
            "total_models": sum(len(v) for v in result.values()) if isinstance(result, dict) else "N/A"
        }
        
        save_raw_output("zhipu", result, metadata)
        return True
    except Exception as e:
        print(f"❌ 智谱AI捕获失败: {e}")
        save_raw_output("zhipu", {"error": str(e)}, {"error": True})
        return False


def capture_openrouter():
    """捕获OpenRouter的原始输出"""
    print("🔍 开始捕获OpenRouter原始输出...")
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        print("⚠️  OPENROUTER_API_KEY环境变量未设置，跳过OpenRouter")
        save_raw_output("openrouter", {"error": "API_KEY_NOT_SET"}, {"skipped": True})
        return False
    
    try:
        fetcher = FetchOpenRouter(api_key=api_key)
        
        # 为了获取原始响应，我们需要修改代码或使用requests直接调用
        # 这里我们先使用现有的fetch_models方法
        result = fetcher.fetch_models()
        
        metadata = {
            "result_type": type(result).__name__,
            "models_count": len(result) if isinstance(result, list) else "N/A"
        }
        
        save_raw_output("openrouter", result, metadata)
        return True
    except Exception as e:
        print(f"❌ OpenRouter捕获失败: {e}")
        save_raw_output("openrouter", {"error": str(e)}, {"error": True})
        return False


def capture_modelscope():
    """捕获ModelScope的原始输出"""
    print("🔍 开始捕获ModelScope原始输出...")
    try:
        fetcher = FetchModelScope()
        result = fetcher.fetch_models()
        
        metadata = {
            "result_type": type(result).__name__,
            "models_count": len(result) if isinstance(result, list) else "N/A"
        }
        
        save_raw_output("modelscope", result, metadata)
        return True
    except Exception as e:
        print(f"❌ ModelScope捕获失败: {e}")
        save_raw_output("modelscope", {"error": str(e)}, {"error": True})
        return False


def capture_siliconflow():
    """捕获SiliconFlow的原始输出"""
    print("🔍 开始捕获SiliconFlow原始输出...")
    api_key = os.getenv("SILICONFLOW_API_KEY")
    if not api_key:
        print("⚠️  SILICONFLOW_API_KEY环境变量未设置，跳过SiliconFlow")
        save_raw_output("siliconflow", {"error": "API_KEY_NOT_SET"}, {"skipped": True})
        return False
    
    try:
        fetcher = FetchSiliconFlow()
        result = fetcher.fetch_models()
        
        metadata = {
            "result_type": type(result).__name__,
            "models_count": len(result) if isinstance(result, list) else "N/A"
        }
        
        save_raw_output("siliconflow", result, metadata)
        return True
    except Exception as e:
        print(f"❌ SiliconFlow捕获失败: {e}")
        save_raw_output("siliconflow", {"error": str(e)}, {"error": True})
        return False


def capture_iflow():
    """捕获IFlow的原始输出"""
    print("🔍 开始捕获IFlow原始输出...")
    api_key = os.getenv("IFLOW_API_KEY")
    if not api_key:
        print("⚠️  IFLOW_API_KEY环境变量未设置，跳过IFlow")
        save_raw_output("iflow", {"error": "API_KEY_NOT_SET"}, {"skipped": True})
        return False
    
    try:
        fetcher = FetchIFlow()
        result = fetcher.fetch_models()
        
        metadata = {
            "result_type": type(result).__name__,
            "models_count": len(result) if isinstance(result, list) else "N/A"
        }
        
        save_raw_output("iflow", result, metadata)
        return True
    except Exception as e:
        print(f"❌ IFlow捕获失败: {e}")
        save_raw_output("iflow", {"error": str(e)}, {"error": True})
        return False


def capture_direct_api_calls():
    """直接调用API获取最原始的响应"""
    print("\n🔧 开始直接API调用获取原始响应...")
    import requests
    
    results = {}
    
    # 1. ModelScope直接调用
    print("  调用ModelScope API...")
    try:
        response = requests.get(
            "https://modelscope.cn/api/v1/models?PageSize=50&PageNumber=1",
            timeout=30
        )
        results["modelscope_direct"] = {
            "status_code": response.status_code,
            "headers": dict(response.headers),
            "data": response.json() if response.status_code == 200 else response.text
        }
        print("  ✅ ModelScope直接调用成功")
    except Exception as e:
        results["modelscope_direct"] = {"error": str(e)}
        print(f"  ❌ ModelScope直接调用失败: {e}")
    
    # 2. 智谱AI网页抓取
    print("  抓取智谱AI网页...")
    try:
        response = requests.get(
            "https://open.bigmodel.cn/dev/howuse/model",
            timeout=30
        )
        results["zhipu_direct"] = {
            "status_code": response.status_code,
            "headers": dict(response.headers),
            "html_length": len(response.text),
            "sample_html": response.text[:1000] if response.status_code == 200 else ""
        }
        print("  ✅ 智谱AI网页抓取成功")
    except Exception as e:
        results["zhipu_direct"] = {"error": str(e)}
        print(f"  ❌ 智谱AI网页抓取失败: {e}")
    
    # 保存直接API调用的结果
    save_raw_output("direct_api_calls", results, {
        "timestamp": datetime.now().strftime("%Y%m%d_%H%M%S"),
        "note": "直接API调用结果，包含原始HTTP响应"
    })
    
    return results


def main():
    """主函数"""
    print("=" * 60)
    print("开始捕获所有提供商的原始输出")
    print("=" * 60)
    
    # 检查环境变量
    print("\n📋 环境变量检查:")
    env_vars = ["OPENROUTER_API_KEY", "SILICONFLOW_API_KEY", "IFLOW_API_KEY"]
    for var in env_vars:
        value = os.getenv(var)
        status = "✅ 已设置" if value else "❌ 未设置"
        print(f"  {var}: {status}")
    
    print("\n" + "=" * 60)
    
    # 捕获各个提供商的输出
    results = {
        "zhipu": capture_zhipu(),
        "openrouter": capture_openrouter(),
        "modelscope": capture_modelscope(),
        "siliconflow": capture_siliconflow(),
        "iflow": capture_iflow(),
    }
    
    # 直接API调用
    direct_results = capture_direct_api_calls()
    
    # 汇总结果
    print("\n" + "=" * 60)
    print("捕获完成！汇总结果:")
    print("=" * 60)
    
    successful = sum(1 for v in results.values() if v)
    total = len(results)
    
    print(f"\n📊 提供商捕获结果: {successful}/{total} 成功")
    for provider, success in results.items():
        status = "✅ 成功" if success else "❌ 失败/跳过"
        print(f"  {provider}: {status}")
    
    print(f"\n📁 所有输出文件已保存到: {Path(__file__).parent}")
    print("\n📝 下一步建议:")
    print("1. 检查生成的JSON文件，分析不同提供商的原始输出格式")
    print("2. 设计统一的返回体格式")
    print("3. 修改各个提供商的fetch_models方法，返回统一格式")
    print("4. 更新测试用例以验证统一格式")
    
    return successful > 0  # 只要有一个成功就返回True


if __name__ == "__main__":
    try:
        success = main()
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\n\n⚠️  用户中断操作")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 脚本执行出错: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
