#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
测试流式检索的真实进度反馈
"""

import asyncio
import json
import time
import traceback
from datetime import datetime

import httpx

async def test_stream_real():
    """测试真实的流式进度"""

    url = "http://localhost:8080/retrieve/stream"
    headers = {
        "Accept": "text/event-stream",
        "Content-Type": "application/json"
    }

    data = {
        "query": "数据管理实践的白皮书是什么",
        "mode": "0",
        "save_output": False
    }

    print(f"\n{'='*60}")
    print(f"开始流式检索测试 - {datetime.now()}")
    print(f"查询: {data['query']}")
    print(f"{'='*60}\n")

    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            async with client.stream("POST", url, json=data, headers=headers) as response:
                print(f"状态码: {response.status_code}")
                print(f"响应头: {dict(response.headers)}\n")

                if response.status_code != 200:
                    content = await response.aread()
                    print(f"错误响应: {content.decode()}")
                    return

                print("接收流式数据:\n")
                buffer = ""
                event_count = 0

                async for chunk in response.aiter_text():
                    buffer += chunk

                    # 处理完整的SSE消息
                    while "\n\n" in buffer:
                        message, buffer = buffer.split("\n\n", 1)

                        if message.startswith("data: "):
                            data_str = message[6:]

                            if data_str == "[DONE]":
                                print(f"\n流式传输完成")
                                break

                            try:
                                event = json.loads(data_str)
                                event_count += 1

                                # 打印事件详情
                                timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3]
                                event_type = event.get("type", "unknown")
                                progress = event.get("progress", 0)
                                elapsed = event.get("elapsed", 0)

                                print(f"[{timestamp}] Event #{event_count}")
                                print(f"  类型: {event_type}")
                                print(f"  进度: {progress}%")
                                print(f"  耗时: {elapsed}s")

                                # 打印特定类型的详情
                                if event_type == "sub_queries" and "data" in event:
                                    queries = event["data"]
                                    print(f"  子查询数: {len(queries)}")
                                    for i, q in enumerate(queries[:3], 1):
                                        print(f"    {i}. {q}")

                                elif event_type == "documents" and "data" in event:
                                    doc_data = event["data"]
                                    print(f"  文档数: {doc_data.get('count', 0)}")
                                    if "sources" in doc_data:
                                        sources = doc_data["sources"][:3]
                                        print(f"  来源: {', '.join(sources)}")

                                elif event_type == "iteration" and "data" in event:
                                    iter_data = event["data"]
                                    print(f"  当前轮次: {iter_data.get('current')}/{iter_data.get('max')}")

                                elif event_type == "sufficiency_check" and "data" in event:
                                    check = event["data"]
                                    print(f"  充分性: {check.get('is_sufficient')}")
                                    print(f"  置信度: {check.get('confidence')}%")

                                elif event_type == "answer" and "data" in event:
                                    answer_data = event["data"]
                                    content = answer_data.get("content", "")[:200]
                                    print(f"  答案预览: {content}...")

                                print()  # 空行分隔

                            except json.JSONDecodeError as e:
                                print(f"JSON解析错误: {e}")
                                print(f"原始数据: {data_str}")

    except Exception as e:
        print(f"\n测试失败: {e}")
        traceback.print_exc()

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