#!/usr/bin/env python3
"""
Test 2-phase retrieval strategy
"""

import asyncio
import httpx
import json
import logging
import random
import time

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

async def test_2phase_retrieval():
    """Test that agent uses 2-phase retrieval for content-focused queries"""
    
    session_id = f"2phase-test-{random.randint(1000000000, 9999999999)}"
    base_url = "http://127.0.0.1:8000"
    
    # Test query that should trigger 2-phase retrieval
    query = "如何测试电动汽车的充电性能？请详细说明测试方法和步骤。"
    
    logger.info("🎯 2-PHASE RETRIEVAL TEST")
    logger.info("=" * 80)
    logger.info(f"📝 Session: {session_id}")
    logger.info(f"📝 Query: {query}")
    logger.info("-" * 60)
    
    # Create the request payload
    payload = {
        "messages": [
            {
                "role": "user",
                "content": query
            }
        ],
        "session_id": session_id
    }
    
    # Track tool usage
    metadata_tools = 0
    content_tools = 0
    total_tools = 0
    
    timeout = httpx.Timeout(120.0)  # 2 minute timeout
    
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            logger.info("✅ Streaming response started")
            
            async with client.stream(
                "POST", 
                f"{base_url}/api/chat",
                json=payload,
                headers={"Content-Type": "application/json"}
            ) as response:
                
                # Check if the response started successfully
                if response.status_code != 200:
                    error_body = await response.aread()
                    logger.error(f"❌ HTTP {response.status_code}: {error_body.decode()}")
                    return
                
                # Process the streaming response
                current_event_type = None
                
                async for line in response.aiter_lines():
                    if not line.strip():
                        continue
                    
                    if line.startswith("event: "):
                        current_event_type = line[7:]  # Remove "event: " prefix
                        continue
                    
                    if line.startswith("data: "):
                        data_str = line[6:]  # Remove "data: " prefix
                        
                        if data_str == "[DONE]":
                            logger.info("✅ Stream completed with [DONE]")
                            break
                        
                        try:
                            event_data = json.loads(data_str)
                            event_type = current_event_type or "unknown"
                            
                            if event_type == "tool_start":
                                total_tools += 1
                                tool_name = event_data.get("name", "unknown")
                                args = event_data.get("args", {})
                                query_arg = args.get("query", "")[:50] + "..." if len(args.get("query", "")) > 50 else args.get("query", "")
                                
                                if tool_name == "retrieve_standard_regulation":
                                    metadata_tools += 1
                                    logger.info(f"📋 Phase 1 Tool {metadata_tools}: {tool_name}")
                                    logger.info(f"   Query: {query_arg}")
                                elif tool_name == "retrieve_doc_chunk_standard_regulation":
                                    content_tools += 1
                                    logger.info(f"📄 Phase 2 Tool {content_tools}: {tool_name}")
                                    logger.info(f"   Query: {query_arg}")
                                else:
                                    logger.info(f"🔧 Tool {total_tools}: {tool_name}")
                            
                            elif event_type == "tool_result":
                                tool_name = event_data.get("name", "unknown")
                                results_count = len(event_data.get("results", []))
                                took_ms = event_data.get("took_ms", 0)
                                logger.info(f"✅ Tool completed: {tool_name} ({results_count} results, {took_ms}ms)")
                            
                            elif event_type == "tokens":
                                # Don't log every token, just count them
                                pass
                            
                            # Reset event type for next event
                            current_event_type = None
                            
                            # Break after many tools to avoid too much output
                            if total_tools > 20:
                                logger.info("   ⚠️ Breaking after 20 tools...")
                                break
                                
                        except json.JSONDecodeError as e:
                            logger.warning(f"⚠️ Failed to parse event: {e}")
                            current_event_type = None
                    
    except Exception as e:
        logger.error(f"❌ Request failed: {e}")
        return
    
    # Results
    logger.info("=" * 80)
    logger.info("📊 2-PHASE RETRIEVAL ANALYSIS")
    logger.info("=" * 80)
    logger.info(f"Phase 1 (Metadata) tools: {metadata_tools}")
    logger.info(f"Phase 2 (Content) tools: {content_tools}")
    logger.info(f"Total tools executed: {total_tools}")
    logger.info("-" * 60)
    
    # Success criteria
    success_criteria = [
        (metadata_tools > 0, f"Phase 1 metadata retrieval: {'✅' if metadata_tools > 0 else '❌'} ({metadata_tools} tools)"),
        (content_tools > 0, f"Phase 2 content retrieval: {'✅' if content_tools > 0 else '❌'} ({content_tools} tools)"),
        (total_tools >= 2, f"Multi-tool execution: {'✅' if total_tools >= 2 else '❌'} ({total_tools} tools)")
    ]
    
    logger.info("✅ SUCCESS CRITERIA:")
    all_passed = True
    for passed, message in success_criteria:
        logger.info(f"   {message}")
        if not passed:
            all_passed = False
    
    if all_passed:
        logger.info("🎉 2-PHASE RETRIEVAL TEST PASSED!")
        logger.info("   ✅ Agent correctly uses both metadata and content retrieval tools")
    else:
        logger.info("❌ 2-PHASE RETRIEVAL TEST FAILED!")
        if metadata_tools == 0:
            logger.info("   ❌ No metadata retrieval tools used")
        if content_tools == 0:
            logger.info("   ❌ No content retrieval tools used - this is the main issue!")

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