#!/usr/bin/env python3
"""
Ollama 服务测试
"""

import asyncio
import sys
import os

# 添加项目根目录到 Python 路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.services.ollama_service import OllamaService


async def test_ollama_health():
    """测试 Ollama 健康状态"""
    print("🧪 测试 Ollama 健康状态")
    print("-" * 30)
    
    service = OllamaService()
    
    try:
        is_healthy = await service.health_check()
        if is_healthy:
            print("✅ Ollama 服务健康")
        else:
            print("❌ Ollama 服务不健康")
            
        # 获取详细健康信息
        health_info = await service.get_health_info()
        print(f"📊 健康信息: {health_info}")
        
    except Exception as e:
        print(f"❌ 健康检查异常: {e}")


async def test_list_models():
    """测试获取模型列表"""
    print("\n🧪 测试获取模型列表")
    print("-" * 30)
    
    service = OllamaService()
    
    try:
        models = await service.list_models()
        print(f"✅ 获取到 {len(models)} 个模型:")
        
        for model in models:
            name = model.get('name', 'unknown')
            size = model.get('size', 0)
            size_mb = size / (1024 * 1024) if size > 0 else 0
            print(f"  📦 {name} ({size_mb:.1f} MB)")
            
    except Exception as e:
        print(f"❌ 获取模型列表异常: {e}")


async def test_simple_chat():
    """测试简单聊天"""
    print("\n🧪 测试简单聊天")
    print("-" * 30)
    
    service = OllamaService()
    
    try:
        messages = [
            {"role": "user", "content": "你好，请简单介绍一下你自己"}
        ]
        
        print("📝 用户: 你好，请简单介绍一下你自己")
        print("🤖 AI: ", end="", flush=True)
        
        response = await service.chat(messages)
        print(response)
        print("✅ 简单聊天测试通过")
        
    except Exception as e:
        print(f"❌ 简单聊天异常: {e}")


async def test_streaming_chat():
    """测试流式聊天"""
    print("\n🧪 测试流式聊天")
    print("-" * 30)
    
    service = OllamaService()
    
    try:
        messages = [
            {"role": "user", "content": "请用一句话介绍人工智能"}
        ]
        
        print("📝 用户: 请用一句话介绍人工智能")
        print("🤖 AI: ", end="", flush=True)
        
        chunk_count = 0
        async for chunk in service.chat_streaming(messages):
            print(chunk, end="", flush=True)
            chunk_count += 1
        
        print(f"\n✅ 流式聊天测试通过，收到 {chunk_count} 个数据块")
        
    except Exception as e:
        print(f"❌ 流式聊天异常: {e}")


async def test_router_model():
    """测试路由模型"""
    print("\n🧪 测试路由模型")
    print("-" * 30)
    
    service = OllamaService()
    
    test_cases = [
        "我想查询订单 ORD12345",
        "你好，今天天气怎么样？",
        "帮我退货",
        "谢谢你的帮助"
    ]
    
    try:
        for user_input in test_cases:
            print(f"📝 用户: {user_input}")
            
            result = await service.analyze_intent(user_input)
            intent = result.get('intent', 'unknown')
            confidence = result.get('confidence', 0.0)
            
            print(f"🎯 意图: {intent} (置信度: {confidence:.2f})")
            print("-" * 20)
        
        print("✅ 路由模型测试通过")
        
    except Exception as e:
        print(f"❌ 路由模型异常: {e}")


async def test_model_availability():
    """测试模型可用性"""
    print("\n🧪 测试模型可用性")
    print("-" * 30)
    
    service = OllamaService()
    
    # 测试聊天模型
    chat_available = await service.is_model_available(service.chat_model)
    chat_icon = "✅" if chat_available else "❌"
    print(f"{chat_icon} 聊天模型 ({service.chat_model}): {'可用' if chat_available else '不可用'}")
    
    # 测试路由模型
    router_available = await service.is_model_available(service.router_model)
    router_icon = "✅" if router_available else "❌"
    print(f"{router_icon} 路由模型 ({service.router_model}): {'可用' if router_available else '不可用'}")
    
    if not chat_available or not router_available:
        print("\n💡 如果模型不可用，请运行以下命令下载:")
        if not chat_available:
            print(f"   ollama pull {service.chat_model}")
        if not router_available:
            print(f"   ollama pull {service.router_model}")


async def test_error_handling():
    """测试错误处理"""
    print("\n🧪 测试错误处理")
    print("-" * 30)
    
    service = OllamaService()
    
    try:
        # 测试不存在的模型
        print("📝 测试不存在的模型...")
        old_model = service.chat_model
        service.chat_model = "nonexistent-model"
        
        try:
            await service.chat([{"role": "user", "content": "test"}])
            print("❌ 应该抛出异常但没有")
        except Exception as e:
            print(f"✅ 正确处理了模型不存在的错误: {type(e).__name__}")
        finally:
            service.chat_model = old_model
        
        # 测试空消息
        print("📝 测试空消息...")
        try:
            await service.chat([])
            print("❌ 应该抛出异常但没有")
        except Exception as e:
            print(f"✅ 正确处理了空消息错误: {type(e).__name__}")
        
        print("✅ 错误处理测试通过")
        
    except Exception as e:
        print(f"❌ 错误处理测试异常: {e}")


async def main():
    """主测试函数"""
    print("🚀 开始测试 Ollama 服务")
    print("=" * 50)
    
    try:
        # 基础测试
        await test_ollama_health()
        await test_list_models()
        await test_model_availability()
        
        # 功能测试
        await test_simple_chat()
        await test_streaming_chat()
        await test_router_model()
        
        # 错误处理测试
        await test_error_handling()
        
        print("\n🎊 所有 Ollama 测试通过！")
        
    except Exception as e:
        print(f"\n💥 测试过程中出现错误: {e}")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())