#!/usr/bin/env python3
"""
多 AI 模型调用脚本 (Universal Architect Engine v4.0)

支持双中转站 API:
  - 中转站 GPT:    gpt-5.2-codex 等
  - 中转站 Gemini: gemini-2.5-pro, gemini-2.5-flash 等

环境变量:
  AI_API_KEY       - 中转站 API 密钥 (GPT 模型)
  AI_BASE_URL      - 中转站地址 (默认: https://api.gbro.site)
  GEMINI_API_KEY   - 中转站 API 密钥 (Gemini 模型，可与 AI_API_KEY 相同)
  GEMINI_BASE_URL  - Gemini 中转站地址 (默认: 与 AI_BASE_URL 相同)

使用方法:
  python multi_ai.py --model gpt-5.2-codex --task audit --file code.py
  python multi_ai.py --model gemini-2.5-pro --task frontend --prompt "生成登录组件"
  python multi_ai.py --list-models
  python multi_ai.py --test
"""

import os
import sys
import argparse
import requests
import json

# ============================================
# 配置
# ============================================

# 默认中转站地址
DEFAULT_BASE_URL = "https://api.gbro.site"

# 模型到 API 源的映射
MODEL_API_SOURCE = {
    # GPT 模型 → GPT 中转站
    "gpt-5.2-codex": "gpt",
    "gpt-5.2": "gpt",
    "gpt-5.1": "gpt",
    "gpt-5.1-codex": "gpt",
    "gpt-5.1-codex-mini": "gpt",
    "gpt-5.1-codex-max": "gpt",
    
    # Gemini 模型 → Gemini 中转站
    "gemini-2.5-pro": "gemini",
    "gemini-2.5-flash": "gemini",
    "gemini-3-pro-preview": "gemini",
    "gemini-3-pro-preview-thinking": "gemini",
    "gemini-2.5-flash-image": "gemini",
    "gemini-2.5-flash-image-preview": "gemini",
    
    # Claude 模型 → GPT 中转站
    "claude-sonnet-4-5-20250929": "gpt",
    "claude-opus-4-5": "gpt",
    "claude-haiku-4-5-20251001": "gpt",
}

# 模型别名映射
MODELS = {
    # ===== OpenAI/GPT 系列 =====
    "gpt-5.2-codex": "gpt-5.2-codex",
    "gpt-5.2": "gpt-5.2",
    "gpt-5.1": "gpt-5.1",
    "codex": "gpt-5.2-codex",
    "openai": "gpt-5.2-codex",
    "gpt": "gpt-5.2-codex",
    
    # ===== Gemini 系列 =====
    "gemini-2.5-pro": "gemini-2.5-pro",
    "gemini-2.5-flash": "gemini-2.5-flash",
    "gemini-3-pro": "gemini-3-pro-preview",
    "gemini-3-pro-preview": "gemini-3-pro-preview",
    "gemini": "gemini-2.5-flash",
    "flash": "gemini-2.5-flash",
    
    # ===== Claude 系列 =====
    "claude-sonnet": "claude-sonnet-4-5-20250929",
    "claude-opus": "claude-opus-4-5",
    "claude-haiku": "claude-haiku-4-5-20251001",
    "claude": "claude-sonnet-4-5-20250929",
    
    # ===== 任务别名 =====
    "audit": "gpt-5.2-codex",           # 代码审计用 GPT
    "frontend": "gemini-2.5-flash",      # 前端生成用 Gemini
    "design": "gemini-2.5-flash",        # 设计系统用 Gemini
    "backend": "gpt-5.2-codex",          # 后端逻辑用 GPT
    "fast": "gemini-2.5-flash",          # 快速任务用 Gemini Flash
}

# 任务提示词模板
TASK_PROMPTS = {
    "audit": """你是一位首席架构师，请对以下代码进行深度审计。

审计维度:
1. 🔐 Security: 注入预防、鉴权逻辑、敏感数据泄露
2. ⚡ Performance: N+1 查询、冗余渲染、索引建议
3. 🔧 Maintainability: 代码耦合度、模块化清晰度

请输出结构化报告，区分:
- 🚫 BLOCKER (必须修改)
- 💡 ADVICE (改进建议)

代码:
```
{content}
```

请用中文回复。""",

    "frontend": """你是一位高级前端架构师，请根据以下需求生成现代化前端代码。

要求:
1. 使用 React + TypeScript + Tailwind CSS
2. 组件命名具备清晰的业务语义
3. 包含完整的类型定义
4. 预留 OpenTelemetry 埋点

需求:
{content}

请用中文注释。""",

    "design-system": """你是一位 UI/UX 设计系统专家，请根据以下项目信息生成 Design Tokens。

项目信息:
{content}

请输出 JSON 格式的 Design Tokens，包括:
1. colors (primary, secondary, semantic)
2. spacing (unit, scale)
3. typography (fontFamily, fontSize)
4. borderRadius
5. shadows

同时提供关键 UI 组件的设计规范描述。""",

    "backend": """你是一位后端架构专家，请根据以下需求设计后端架构。

要求:
1. 遵循 SOLID 原则和 Clean Architecture
2. 考虑高并发和幂等性设计
3. 预留分布式追踪埋点
4. 包含异常处理框架

需求:
{content}

请用中文注释。""",

    "contract": """你是一位全栈架构师，请根据以下 API 设计生成前后端握手契约。

要求:
1. TypeScript Interfaces 或 Protobuf 格式
2. 包含 DTO (Data Transfer Object) 定义
3. 包含错误码枚举
4. 确保前端组件与后端 100% 类型对齐

API 设计:
{content}

请输出完整的 Handshake Contract。""",

    "general": """{content}"""
}


def get_config():
    """获取配置"""
    # GPT 中转站配置
    gpt_key = os.environ.get("AI_API_KEY")
    gpt_url = os.environ.get("AI_BASE_URL") or DEFAULT_BASE_URL
    
    # Gemini 中转站配置（可单独设置，也可复用 GPT 配置）
    gemini_key = os.environ.get("GEMINI_API_KEY") or gpt_key
    gemini_url = os.environ.get("GEMINI_BASE_URL") or gpt_url
    
    return {
        "gpt_key": gpt_key,
        "gpt_url": gpt_url,
        "gemini_key": gemini_key,
        "gemini_url": gemini_url
    }


def call_model(model: str, prompt: str, temperature: float = 0.3) -> str:
    """调用模型"""
    config = get_config()
    
    # 解析模型名称
    actual_model = MODELS.get(model, model)
    
    # 确定使用哪个 API 源
    api_source = MODEL_API_SOURCE.get(actual_model, "gpt")
    
    if api_source == "gemini":
        return call_proxy_api(
            actual_model, prompt, 
            config["gemini_key"], config["gemini_url"], 
            temperature
        )
    else:
        return call_proxy_api(
            actual_model, prompt, 
            config["gpt_key"], config["gpt_url"], 
            temperature
        )


def call_proxy_api(model: str, prompt: str, api_key: str, base_url: str, temperature: float) -> str:
    """调用中转站 API (OpenAI 兼容格式)"""
    if not api_key:
        print("❌ 未设置 API Key 环境变量")
        sys.exit(1)
    
    url = f"{base_url}/v1/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    data = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": temperature
    }
    
    response = requests.post(url, headers=headers, json=data, timeout=180)
    response.raise_for_status()
    
    result = response.json()
    return result["choices"][0]["message"]["content"]


def list_models():
    """列出可用模型"""
    config = get_config()
    
    print("=" * 60)
    print("📋 可用模型列表")
    print("=" * 60)
    
    # 显示 GPT 中转站模型
    print("\n🔶 GPT 中转站模型:")
    print(f"   地址: {config['gpt_url']}")
    if config["gpt_key"]:
        url = f"{config['gpt_url']}/v1/models"
        headers = {"Authorization": f"Bearer {config['gpt_key']}"}
        
        try:
            response = requests.get(url, headers=headers, timeout=30)
            response.raise_for_status()
            
            models = response.json().get("data", [])
            gpt_models = [m["id"] for m in models if "gpt" in m["id"].lower() or "codex" in m["id"].lower()]
            
            for m in sorted(gpt_models)[:10]:
                print(f"   - {m}")
            if len(gpt_models) > 10:
                print(f"   ... 还有 {len(gpt_models) - 10} 个")
                
        except Exception as e:
            print(f"   ❌ 获取失败: {e}")
    else:
        print("   ⚠️ 未设置 AI_API_KEY")
    
    # 显示 Gemini 中转站模型
    print("\n🔷 Gemini 中转站模型:")
    print(f"   地址: {config['gemini_url']}")
    if config["gemini_key"]:
        url = f"{config['gemini_url']}/v1/models"
        headers = {"Authorization": f"Bearer {config['gemini_key']}"}
        
        try:
            response = requests.get(url, headers=headers, timeout=30)
            response.raise_for_status()
            
            models = response.json().get("data", [])
            gemini_models = [m["id"] for m in models if "gemini" in m["id"].lower()]
            
            for m in sorted(gemini_models)[:10]:
                print(f"   - {m}")
            if len(gemini_models) > 10:
                print(f"   ... 还有 {len(gemini_models) - 10} 个")
                
        except Exception as e:
            print(f"   ❌ 获取失败: {e}")
    else:
        print("   ⚠️ 未设置 GEMINI_API_KEY")
    
    print("\n" + "=" * 60)
    print("环境变量配置:")
    print(f"  AI_API_KEY:      {'✅ 已设置' if config['gpt_key'] else '❌ 未设置'} (GPT)")
    print(f"  GEMINI_API_KEY:  {'✅ 已设置' if config['gemini_key'] else '❌ 未设置'} (Gemini)")
    print("=" * 60)


def test_models():
    """测试主要模型是否可用"""
    config = get_config()
    
    print("=" * 60)
    print("🔍 测试模型可用性")
    print("=" * 60)
    
    # 测试 GPT 中转站
    print("\n--- GPT 中转站 ---")
    if config["gpt_key"]:
        for model, name in [("gpt-5.2-codex", "GPT-5.2 Codex")]:
            print(f"\n📝 测试 {name}...")
            try:
                result = call_proxy_api(model, "说你好", config["gpt_key"], config["gpt_url"], 0.1)
                print(f"   ✅ 可用: {result[:30]}...")
            except requests.exceptions.HTTPError as e:
                error_msg = str(e)
                if hasattr(e, 'response'):
                    error_msg = e.response.text[:100]
                print(f"   ❌ 失败: {error_msg}")
            except Exception as e:
                print(f"   ❌ 失败: {e}")
    else:
        print("   ⚠️ 未设置 AI_API_KEY")
    
    # 测试 Gemini 中转站
    print("\n--- Gemini 中转站 ---")
    if config["gemini_key"]:
        for model, name in [
            ("gemini-2.5-flash", "Gemini 2.5 Flash"),
            ("gemini-2.5-pro", "Gemini 2.5 Pro"),
        ]:
            print(f"\n📝 测试 {name}...")
            try:
                result = call_proxy_api(model, "说你好", config["gemini_key"], config["gemini_url"], 0.1)
                print(f"   ✅ 可用: {result[:30]}...")
            except requests.exceptions.HTTPError as e:
                error_msg = str(e)
                if hasattr(e, 'response'):
                    try:
                        err_json = e.response.json()
                        error_msg = err_json.get("error", {}).get("message", str(e))[:80]
                    except:
                        error_msg = e.response.text[:80]
                print(f"   ❌ 失败: {error_msg}")
            except Exception as e:
                print(f"   ❌ 失败: {e}")
    else:
        print("   ⚠️ 未设置 GEMINI_API_KEY")
    
    print("\n" + "=" * 60)
    print("测试完成")
    print("=" * 60)


def main():
    parser = argparse.ArgumentParser(
        description='多 AI 模型调用工具 (Universal Architect Engine v4.0)',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument('--model', '-m', default='gpt-5.2-codex',
                        help='模型名称 (gpt-5.2-codex, gemini-2.5-pro, etc.)')
    parser.add_argument('--task', '-t', default='general',
                        choices=['audit', 'frontend', 'design-system', 'backend', 'contract', 'general'],
                        help='任务类型')
    parser.add_argument('--file', '-f', help='输入文件')
    parser.add_argument('--prompt', '-p', help='直接输入提示词')
    parser.add_argument('--output', '-o', help='输出文件')
    parser.add_argument('--temperature', type=float, default=0.3, help='温度参数')
    parser.add_argument('--list-models', action='store_true', help='列出可用模型')
    parser.add_argument('--test', action='store_true', help='测试主要模型可用性')
    parser.add_argument('--json', '-j', action='store_true', help='JSON 格式输出')
    
    args = parser.parse_args()
    
    # 列出模型
    if args.list_models:
        list_models()
        return
    
    # 测试模型
    if args.test:
        test_models()
        return
    
    # 获取输入内容
    if args.file:
        with open(args.file, 'r', encoding='utf-8') as f:
            content = f.read()
    elif args.prompt:
        content = args.prompt
    elif not sys.stdin.isatty():
        content = sys.stdin.read()
    else:
        parser.print_help()
        return
    
    # 构建提示词
    task_template = TASK_PROMPTS.get(args.task, TASK_PROMPTS["general"])
    full_prompt = task_template.format(content=content)
    
    # 输出状态
    actual_model = MODELS.get(args.model, args.model)
    print(f"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", file=sys.stderr)
    print(f"[Status: Executing {args.task} via {actual_model}]", file=sys.stderr)
    print(f"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", file=sys.stderr)
    
    try:
        result = call_model(args.model, full_prompt, args.temperature)
        
        if args.json:
            output = json.dumps({
                "model": actual_model,
                "task": args.task,
                "result": result
            }, ensure_ascii=False, indent=2)
        else:
            output = result
        
        # 输出结果
        if args.output:
            with open(args.output, 'w', encoding='utf-8') as f:
                f.write(output)
            print(f"✅ 结果已保存到: {args.output}", file=sys.stderr)
        else:
            print()
            print(output)
        
        print(f"\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", file=sys.stderr)
        print(f"[Task Complete]", file=sys.stderr)
        print(f"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", file=sys.stderr)
        
    except requests.exceptions.HTTPError as e:
        print(f"❌ API 错误: {e}")
        if hasattr(e, 'response'):
            print(f"   响应: {e.response.text}")
        sys.exit(1)
    except Exception as e:
        print(f"❌ 失败: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()