#!/usr/bin/env python3
"""
OpenRouter API 客户端
===================

支持 GPT-4o, Claude, Qwen 等多个模型

作者: xialab
日期: 2026-04-27
"""

import os
import sys
from pathlib import Path
from typing import List, Dict, Tuple
import time

# 加载环境变量
from dotenv import load_dotenv
PROJECT_ROOT = Path("/home/xzq/gene-family-mas-v2")
load_dotenv(PROJECT_ROOT / ".env")

# 尝试导入 openai
try:
    from openai import OpenAI
except ImportError:
    print("[错误] 请先安装 openai: pip install openai")
    sys.exit(1)


class OpenRouterClient:
    """OpenRouter API 客户端"""
    
    # 可用模型列表
    MODELS = {
        "gpt-4o": "openai/gpt-4o",
        "gpt-4o-mini": "openai/gpt-4o-mini",
        "claude-3.5-sonnet": "anthropic/claude-3.5-sonnet",
        "claude-3-opus": "anthropic/claude-3-opus",
        "qwen-2.5-72b": "qwen/qwen-2.5-72b-instruct",
        "gemini-1.5-pro": "google/gemini-pro-1.5",
        "llama-3.1-70b": "meta-llama/llama-3.1-70b-instruct",
    }
    
    def __init__(self):
        self.api_key = os.getenv("OPENROUTER_API_KEY")
        self.base_url = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
        self.referer = os.getenv("OPENROUTER_REFERER", "https://localhost")
        self.title = os.getenv("OPENROUTER_TITLE", "GeneFamily-RAG")
        
        if not self.api_key:
            raise ValueError("OPENROUTER_API_KEY not found in environment")
        
        self.client = OpenAI(
            base_url=self.base_url,
            api_key=self.api_key,
            default_headers={
                "HTTP-Referer": self.referer,
                "X-Title": self.title,
            }
        )
        print(f"[✓] OpenRouter client initialized")
        print(f"    Base URL: {self.base_url}")
    
    def query(self, question: str, model: str = "gpt-4o", temperature: float = 0.3) -> Tuple[str, float, Dict]:
        """
        调用OpenRouter API
        
        Args:
            question: 问题文本
            model: 模型名称 (gpt-4o, claude-3.5-sonnet, 等)
            temperature: 温度参数
            
        Returns:
            (answer, latency, metadata)
        """
        start_time = time.time()
        
        # 获取完整模型ID
        model_id = self.MODELS.get(model, model)
        
        try:
            response = self.client.chat.completions.create(
                model=model_id,
                messages=[
                    {
                        "role": "system",
                        "content": "You are a gene family analysis expert. Answer accurately based on your knowledge."
                    },
                    {
                        "role": "user",
                        "content": question
                    }
                ],
                temperature=temperature,
                max_tokens=2048,
            )
            
            latency = time.time() - start_time
            answer = response.choices[0].message.content
            
            metadata = {
                "model": model,
                "model_id": model_id,
                "prompt_tokens": response.usage.prompt_tokens,
                "completion_tokens": response.usage.completion_tokens,
                "total_tokens": response.usage.total_tokens,
            }
            
            return answer, latency, metadata
            
        except Exception as e:
            latency = time.time() - start_time
            return f"[错误] {str(e)}", latency, {"error": str(e)}
    
    def query_multiple(self, question: str, models: List[str] = None) -> Dict[str, Dict]:
        """
        使用多个模型查询同一问题
        
        Args:
            question: 问题文本
            models: 模型列表，默认 [gpt-4o, claude-3.5-sonnet, qwen-2.5-72b]
            
        Returns:
            各模型的回答结果
        """
        if models is None:
            models = ["gpt-4o", "claude-3.5-sonnet", "qwen-2.5-72b"]
        
        results = {}
        print(f"\n问题: {question}\n")
        print("=" * 70)
        
        for model in models:
            print(f"\n[{model}] 查询中...")
            answer, latency, metadata = self.query(question, model)
            
            results[model] = {
                "answer": answer,
                "latency": latency,
                "metadata": metadata
            }
            
            print(f"  耗时: {latency:.2f}s")
            print(f"  Token: {metadata.get('total_tokens', 'N/A')}")
            print(f"  回答: {answer[:100]}...")
        
        return results


def test_single():
    """测试单个问题"""
    client = OpenRouterClient()
    
    question = "拟南芥中SPL基因家族有哪些成员？"
    
    # 使用 gpt-4o-mini (更便宜)
    answer, latency, metadata = client.query(question, model="gpt-4o-mini")
    
    print(f"\n问题: {question}")
    print(f"模型: {metadata['model_id']}")
    print(f"耗时: {latency:.2f}s")
    print(f"Token: {metadata['total_tokens']}")
    print(f"\n回答:\n{answer}")


def test_multiple():
    """测试多个模型对比"""
    client = OpenRouterClient()
    
    question = "SPL基因家族在植物发育中起什么作用？"
    models = ["gpt-4o", "claude-3.5-sonnet"]
    
    results = client.query_multiple(question, models)
    
    # 保存结果
    import json
    output_file = PROJECT_ROOT / "evaluation" / "results" / "openrouter_comparison.json"
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    
    print(f"\n结果已保存: {output_file}")


def run_benchmark():
    """运行完整基准测试"""
    client = OpenRouterClient()
    
    # 测试问题
    questions = [
        "拟南芥中SPL基因家族有哪些成员？",
        "SPL转录因子如何调控植物开花时间？",
        "SPL基因家族在单子叶和双子叶植物中的扩张模式",
    ]
    
    models = ["gpt-4o", "claude-3.5-sonnet", "qwen-2.5-72b"]
    
    all_results = []
    
    for q in questions:
        print(f"\n{'='*70}")
        print(f"问题: {q}")
        print('='*70)
        
        for model in models:
            answer, latency, metadata = client.query(q, model)
            
            result = {
                "question": q,
                "model": model,
                "answer": answer,
                "latency": latency,
                "tokens": metadata.get('total_tokens', 0)
            }
            all_results.append(result)
            
            print(f"\n[{model}] {latency:.2f}s | {metadata.get('total_tokens', 0)} tokens")
    
    # 保存结果
    import json
    from datetime import datetime
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = PROJECT_ROOT / "evaluation" / "results" / f"openrouter_benchmark_{timestamp}.json"
    
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)
    
    print(f"\n{'='*70}")
    print(f"基准测试完成! 结果保存: {output_file}")
    
    # 统计
    avg_latency = sum(r['latency'] for r in all_results) / len(all_results)
    total_tokens = sum(r['tokens'] for r in all_results)
    print(f"平均响应时间: {avg_latency:.2f}s")
    print(f"总Token消耗: {total_tokens}")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='OpenRouter API 测试')
    parser.add_argument('--mode', choices=['single', 'multiple', 'benchmark'], 
                        default='single', help='测试模式')
    
    args = parser.parse_args()
    
    if args.mode == 'single':
        test_single()
    elif args.mode == 'multiple':
        test_multiple()
    elif args.mode == 'benchmark':
        run_benchmark()