#!/usr/bin/env python3
"""Qwen 端到端测试脚本"""

import os
import sys

# 设置环境变量
os.environ['QWEN_API_KEY'] = 'sk-426efxxxxxxxxxxxxxxxxxxxxxxxxxaba'
os.environ['DASHSCOPE_BASE_URL'] = 'https://dashscope-intl.aliyuncs.com/api/v1'

# 添加项目根目录到 Python 路径（仅用于简化测试）
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def test_direct_api_call():
    """测试直接调用 DashScope API"""
    print("=" * 60)
    print("测试 1: 直接调用 DashScope API")
    print("=" * 60)

    try:
        import dashscope
        from dashscope import Generation

        # 配置
        dashscope.api_key = os.environ['QWEN_API_KEY']
        dashscope.base_http_api_url = os.environ['DASHSCOPE_BASE_URL']

        # 测试文本生成
        response = Generation.call(
            model='qwen-turbo',
            prompt='请用一句话介绍什么是 SQL',
            result_format='message'
        )

        if response.status_code == 200:
            content = response.output.choices[0].message.content
            print(f"✅ 测试通过")
            print(f"   响应: {content[:80]}...")
            print(f"   Tokens: {response.usage.input_tokens} (in) + {response.usage.output_tokens} (out)")
            return True
        else:
            print(f"❌ 测试失败: {response.message}")
            return False

    except Exception as e:
        print(f"❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_json_format():
    """测试 JSON 格式返回"""
    print("\n" + "=" * 60)
    print("测试 2: JSON 格式生成")
    print("=" * 60)

    try:
        import dashscope
        from dashscope import Generation
        import json

        # 配置
        dashscope.api_key = os.environ['QWEN_API_KEY']
        dashscope.base_http_api_url = os.environ['DASHSCOPE_BASE_URL']

        # 测试 JSON 生成
        prompt = '''生成一个 SQL 查询，查询所有用户。
请严格按照以下 JSON 格式返回:
{"sql": "...", "description": "..."}
只返回 JSON，不要其他文字。'''

        response = Generation.call(
            model='qwen-turbo',
            prompt=prompt,
            result_format='message'
        )

        if response.status_code == 200:
            content = response.output.choices[0].message.content

            # 尝试解析 JSON
            try:
                # 提取 JSON 部分
                import re
                json_match = re.search(r'\{.*\}', content, re.DOTALL)
                if json_match:
                    json_data = json.loads(json_match.group())
                    print(f"✅ 测试通过")
                    print(f"   SQL: {json_data.get('sql', 'N/A')}")
                    print(f"   描述: {json_data.get('description', 'N/A')}")
                    return True
                else:
                    print(f"⚠️  返回内容不包含 JSON")
                    print(f"   响应: {content[:100]}...")
                    return False
            except json.JSONDecodeError:
                print(f"⚠️  JSON 解析失败")
                print(f"   响应: {content[:100]}...")
                return False
        else:
            print(f"❌ 测试失败: {response.message}")
            return False

    except Exception as e:
        print(f"❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_different_models():
    """测试不同的 Qwen 模型"""
    print("\n" + "=" * 60)
    print("测试 3: 测试不同模型")
    print("=" * 60)

    models_to_test = [
        ('qwen-turbo', 'Qwen Turbo'),
        ('qwen-plus', 'Qwen Plus'),
    ]

    success_count = 0

    try:
        import dashscope
        from dashscope import Generation

        # 配置
        dashscope.api_key = os.environ['QWEN_API_KEY']
        dashscope.base_http_api_url = os.environ['DASHSCOPE_BASE_URL']

        for model_id, model_name in models_to_test:
            print(f"\n  测试 {model_name} ({model_id})...")

            try:
                response = Generation.call(
                    model=model_id,
                    prompt='请说"Hello"',
                    result_format='message'
                )

                if response.status_code == 200:
                    content = response.output.choices[0].message.content
                    print(f"    ✅ {model_name}: {content[:30]}...")
                    success_count += 1
                else:
                    print(f"    ❌ {model_name}: {response.message}")

            except Exception as e:
                print(f"    ❌ {model_name}: {e}")

        print(f"\n  总结: {success_count}/{len(models_to_test)} 个模型测试通过")
        return success_count > 0

    except Exception as e:
        print(f"❌ 测试失败: {e}")
        return False


def main():
    print("\n" + "=" * 60)
    print("Qwen 端到端测试套件")
    print("=" * 60)
    print(f"API Key: {os.environ['QWEN_API_KEY'][:20]}...")
    print(f"Endpoint: {os.environ['DASHSCOPE_BASE_URL']}")
    print()

    # 运行所有测试
    test1 = test_direct_api_call()
    test2 = test_json_format()
    test3 = test_different_models()

    # 总结
    print("\n" + "=" * 60)
    print("测试总结")
    print("=" * 60)
    print(f"{'✅' if test1 else '❌'} 测试 1: 直接 API 调用")
    print(f"{'✅' if test2 else '❌'} 测试 2: JSON 格式生成")
    print(f"{'✅' if test3 else '❌'} 测试 3: 多模型测试")

    total_passed = sum([test1, test2, test3])
    print(f"\n总计: {total_passed}/3 个测试通过")

    if total_passed == 3:
        print("\n🎉 所有测试通过！Qwen 集成准备就绪。")
        return 0
    else:
        print("\n⚠️  部分测试失败，请检查上述错误信息。")
        return 1


if __name__ == "__main__":
    sys.exit(main())