#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
演示改进后的流式处理超时机制

此脚本展示TravelContentCreator中AI_Agent类中优化后的流式处理超时机制:
1. 更智能的流块超时检测 - 只在收到初始块后才检测流块超时
2. 首次响应超时检测 - 检测是否在合理时间内收到第一个响应块
3. 全局超时保护 - 防止整体处理时间过长
"""

import os
import sys
import asyncio
import time
import logging
import random
from pathlib import Path
from datetime import datetime

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

from core.ai_agent import AI_Agent, Timeout

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%H:%M:%S'
)

# 示例提示词
SYSTEM_PROMPT = "你是一个专业的旅游内容创作助手，请根据用户的提示生成相关内容。"
# 普通提示词
NORMAL_PROMPT = "请生成一段关于杭州西湖的简短介绍，不超过200字。"
# 长提示词，可能需要更长处理时间
LONG_PROMPT = """请按照以下格式为我创建一份详尽的云南七日旅游攻略：
1. 每日行程计划（早中晚三餐和活动）
2. 每个景点的历史背景和特色
3. 当地特色美食推荐和品尝地点
4. 交通建议（包括景点间如何移动）
5. 住宿推荐（每个地区的优质酒店或民宿）
6. 注意事项（天气、海拔、装备等）
7. 购物指南（值得购买的纪念品和特产）

请确保涵盖昆明、大理、丽江和香格里拉等主要旅游地点，并考虑季节特点和当地习俗。提供实用且详细的信息。
"""

def print_separator(title):
    """打印分隔线和标题"""
    print("\n" + "="*60)
    print(f" {title} ".center(60, "="))
    print("="*60 + "\n")

def print_with_timestamp(message, end='\n'):
    """打印带有时间戳的消息"""
    timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
    print(f"[{timestamp}] {message}", end=end, flush=True)

def demo_sync_stream():
    """演示同步流式响应方法的改进超时机制"""
    print_separator("同步流式响应的超时机制")
    
    # 创建AI_Agent实例，设置较短的超时便于测试
    agent = AI_Agent(
        base_url=os.environ.get("API_BASE", "https://api.openai.com/v1"),
        model_name=os.environ.get("MODEL_NAME", "gpt-3.5-turbo"),
        api=os.environ.get("API_KEY", "your_api_key"),
        timeout=15,  # 15秒请求超时
        stream_chunk_timeout=5,  # 5秒流块超时
        max_retries=1  # 最小重试减少测试时间
    )
    
    print_with_timestamp("开始生成内容(同步流式方法)...")
    start_time = time.time()
    
    try:
        # 使用同步流式方法
        result = agent.generate_text_stream(
            SYSTEM_PROMPT,
            NORMAL_PROMPT,
            temperature=0.7,
            top_p=0.9,
            presence_penalty=0.0
        )
        
        end_time = time.time()
        
        print_with_timestamp(f"生成完成! 耗时: {end_time - start_time:.2f}秒")
        print_with_timestamp(f"内容长度: {len(result)} 字符")
        
        # 检查是否包含超时或错误提示
        if "[注意:" in result:
            print_with_timestamp("检测到警告信息:")
            warning_start = result.find("[注意:")
            warning_end = result.find("]", warning_start)
            if warning_end != -1:
                print_with_timestamp(f"警告内容: {result[warning_start:warning_end+1]}")
        
    except Exception as e:
        print_with_timestamp(f"生成过程中出错: {type(e).__name__} - {e}")
    
    # 关闭agent
    agent.close()

def demo_callback_stream():
    """演示回调流式响应方法的改进超时机制"""
    print_separator("回调流式响应的超时机制")
    
    # 创建AI_Agent实例
    agent = AI_Agent(
        base_url=os.environ.get("API_BASE", "https://api.openai.com/v1"),
        model_name=os.environ.get("MODEL_NAME", "gpt-3.5-turbo"),
        api=os.environ.get("API_KEY", "your_api_key"),
        timeout=15,
        stream_chunk_timeout=5,
        max_retries=1
    )
    
    # 定义回调函数
    def callback(chunk, accumulated=None):
        """处理流式响应的回调函数"""
        if chunk:
            # 实时打印内容，不换行
            print_with_timestamp(f"收到块({len(chunk)}字符): {chunk[:10]}...", end="\r")
    
    print_with_timestamp("开始生成内容(回调流式方法)...")
    start_time = time.time()
    
    try:
        # 使用回调流式方法
        result = agent.generate_text_stream_with_callback(
            SYSTEM_PROMPT,
            LONG_PROMPT,  # 使用更长的提示
            temperature=0.7,
            top_p=0.9,
            presence_penalty=0.0,
            callback=callback,
            accumulate=True  # 启用累积模式
        )
        
        end_time = time.time()
        print_with_timestamp(f"\n生成完成! 耗时: {end_time - start_time:.2f}秒")
        print_with_timestamp(f"内容长度: {len(result)} 字符")
        
        # 检查是否包含超时或错误提示
        if "[注意:" in result:
            print_with_timestamp("检测到警告信息:")
            warning_start = result.find("[注意:")
            warning_end = result.find("]", warning_start)
            if warning_end != -1:
                print_with_timestamp(f"警告内容: {result[warning_start:warning_end+1]}")
    
    except Exception as e:
        print_with_timestamp(f"生成过程中出错: {type(e).__name__} - {e}")
    
    # 关闭agent
    agent.close()

async def demo_async_stream():
    """演示异步流式响应方法的改进超时机制"""
    print_separator("异步流式响应的超时机制")
    
    # 创建AI_Agent实例
    agent = AI_Agent(
        base_url=os.environ.get("API_BASE", "https://api.openai.com/v1"),
        model_name=os.environ.get("MODEL_NAME", "gpt-3.5-turbo"),
        api=os.environ.get("API_KEY", "your_api_key"),
        timeout=15,
        stream_chunk_timeout=5,
        max_retries=1
    )
    
    print_with_timestamp("开始生成内容(异步流式方法)...")
    start_time = time.time()
    full_response = ""
    chunk_count = 0
    
    try:
        # 使用异步流式方法
        async_stream = agent.async_generate_text_stream(
            SYSTEM_PROMPT,
            NORMAL_PROMPT,
            temperature=0.7,
            top_p=0.9,
            presence_penalty=0.0
        )
        
        # 异步迭代流
        async for content in async_stream:
            # 累积完整响应
            full_response += content
            chunk_count += 1
            
            # 每收到5个块打印一次进度
            if chunk_count % 5 == 0:
                print_with_timestamp(f"已收到 {chunk_count} 个块，当前内容长度: {len(full_response)}字符")
        
    except Exception as e:
        print_with_timestamp(f"生成过程中出错: {type(e).__name__} - {e}")
    
    end_time = time.time()
    print_with_timestamp(f"生成完成! 耗时: {end_time - start_time:.2f}秒")
    print_with_timestamp(f"共收到 {chunk_count} 个块，内容长度: {len(full_response)}字符")
    
    # 检查是否包含超时或错误提示
    if "[注意:" in full_response:
        print_with_timestamp("检测到警告信息:")
        warning_start = full_response.find("[注意:")
        warning_end = full_response.find("]", warning_start)
        if warning_end != -1:
            print_with_timestamp(f"警告内容: {full_response[warning_start:warning_end+1]}")
    
    # 关闭agent
    agent.close()

async def simulate_delayed_response():
    """模拟延迟响应的场景"""
    print_separator("模拟延迟响应")
    
    # 创建带有非常短超时的AI_Agent实例
    agent = AI_Agent(
        base_url=os.environ.get("API_BASE", "https://api.openai.com/v1"),
        model_name=os.environ.get("MODEL_NAME", "gpt-3.5-turbo"),
        api=os.environ.get("API_KEY", "your_api_key"),
        timeout=10,
        stream_chunk_timeout=3,  # 非常短的流块超时
        max_retries=1
    )
    
    # 自定义处理器模拟延迟
    class DelayedProcessor:
        def __init__(self):
            self.chunks = []
            self.total_chars = 0
        
        async def process_chunk(self, chunk):
            self.chunks.append(chunk)
            self.total_chars += len(chunk)
            # 打印进度
            print_with_timestamp(f"收到块 #{len(self.chunks)}, 内容: {chunk[:10]}...", end="\r")
            
            # 随机延迟模拟处理时间
            delay = random.uniform(0.1, 4.0)  # 0.1-4秒，有时会超过超时设置
            if len(self.chunks) % 5 == 0:
                print_with_timestamp(f"\n模拟处理延迟: {delay:.2f}秒")
            await asyncio.sleep(delay)
            
            # 返回处理后的块
            return f"[处理完成: {chunk}]"
    
    processor = DelayedProcessor()
    print_with_timestamp("开始生成内容(带处理延迟)...")
    start_time = time.time()
    full_response = ""
    
    try:
        # 使用异步流式方法
        async_stream = agent.async_generate_text_stream(
            SYSTEM_PROMPT,
            NORMAL_PROMPT,
            temperature=0.7,
            top_p=0.9,
            presence_penalty=0.0
        )
        
        # 异步迭代流，添加处理延迟
        async for content in async_stream:
            # 模拟处理延迟
            processed = await processor.process_chunk(content)
            full_response += processed
        
    except Exception as e:
        print_with_timestamp(f"\n生成过程中出错: {type(e).__name__} - {e}")
    
    end_time = time.time()
    print_with_timestamp(f"\n生成完成! 耗时: {end_time - start_time:.2f}秒")
    print_with_timestamp(f"共处理 {len(processor.chunks)} 个块，总字符数: {processor.total_chars}")
    
    # 检查是否包含超时或错误提示
    if "[注意:" in full_response:
        print_with_timestamp("检测到警告信息:")
        warning_start = full_response.find("[注意:")
        warning_end = full_response.find("]", warning_start)
        if warning_end != -1:
            print_with_timestamp(f"警告内容: {full_response[warning_start:warning_end+1]}")
    
    # 关闭agent
    agent.close()

async def main():
    """主函数"""
    print_with_timestamp("测试改进后的流式处理超时机制...")
    
    # 1. 测试同步流式响应
    demo_sync_stream()
    
    # 2. 测试回调流式响应
    demo_callback_stream()
    
    # 3. 测试异步流式响应
    await demo_async_stream()
    
    # 4. 模拟延迟响应
    await simulate_delayed_response()
    
    print_with_timestamp("\n所有测试完成!")

if __name__ == "__main__":
    # 运行异步主函数
    asyncio.run(main()) 