#!/usr/bin/env python3
"""
Test script for Gemini streaming JSON parsing fix.

This script tests the robust error handling added to handle malformed JSON chunks
during streaming responses from Gemini models.
"""

from xoneaiagents import Agent

def test_gemini_streaming_robustness():
    """Test Gemini streaming with robust error handling."""
    print("🧪 Testing Gemini Streaming Robustness Fix")
    print("=" * 60)
    
    try:
        # Create agent with Gemini model (using a lightweight model for testing)
        agent = Agent(
            instructions="You are a helpful assistant. Be concise.",
            llm="gemini/gemini-2.5-flash",  # Using flash for faster testing
            stream=True,
            verbose=True  # Enable verbose to see the error handling in action
        )
        
        print("✅ Agent created successfully")
        print(f"📊 Model: {agent.llm}")
        print(f"📊 Stream enabled: {agent.stream}")
        print(f"📊 Verbose enabled: {agent.verbose}")
        print()
        
        # Test streaming with a simple prompt that might cause chunking issues
        print("🔄 Testing streaming response...")
        prompt = "Explain what real-time streaming is in AI applications, focusing on the benefits and challenges."
        
        chunk_count = 0
        response_content = ""
        
        try:
            for chunk in agent.start(prompt):
                if chunk:
                    response_content += chunk
                    chunk_count += 1
                    print(chunk, end="", flush=True)
                    
        except Exception as streaming_error:
            print(f"\n❌ Streaming error occurred: {streaming_error}")
            print("🔄 This error should now be handled gracefully with fallback to non-streaming mode")
            return False
            
        print("\n\n" + "="*60)
        print("✅ Streaming completed successfully!")
        print(f"📊 Total chunks received: {chunk_count}")
        print(f"📊 Total response length: {len(response_content)} characters")
        
        if chunk_count > 1:
            print("✅ SUCCESS: Streaming worked with multiple chunks")
        else:
            print("⚠️  WARNING: Only received 1 chunk (may have fallen back to non-streaming)")
        
        return True
        
    except Exception as e:
        print(f"❌ Test failed with error: {e}")
        return False

if __name__ == "__main__":
    print("🚀 Starting Gemini Streaming Robustness Test")
    print("This test validates the JSON parsing error fixes for Gemini streaming")
    print()
    
    success = test_gemini_streaming_robustness()
    
    print(f"\n{'='*60}")
    if success:
        print("🎉 TEST PASSED: Gemini streaming robustness fix is working!")
    else:
        print("💥 TEST FAILED: Issues detected with streaming robustness")
    
    print()
    print("📝 Key improvements tested:")
    print("  • Graceful handling of malformed JSON chunks")
    print("  • Automatic fallback to non-streaming on repeated errors")
    print("  • Better error logging and categorization")
    print("  • Chunk-level error recovery")
    
    # Exit with appropriate status code for CI integration
    import sys
    sys.exit(0 if success else 1)