#!/usr/bin/env python3
"""
Test script to demonstrate the LlamaIndex Memory integration with the party planner agent.
"""

import asyncio
from llama_index.core.memory import Memory
from llama_index.core.memory.memory import StaticMemoryBlock, FactExtractionMemoryBlock
from llama_index.core.llms import ChatMessage
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI

async def test_memory_integration():
    """Test the LlamaIndex Memory integration."""
    print("🧠 Testing LlamaIndex Memory Integration")
    print("=" * 50)
    
    # Initialize the same LLM as in the app
    llm = HuggingFaceInferenceAPI(model="meta-llama/Llama-3.2-11B-Vision-Instruct")
    
    # Create memory blocks (same as in app.py)
    memory_blocks = [
        StaticMemoryBlock(
            name="assistant_info",
            static_content="You are Alfred, a sophisticated gala assistant. You help users find information about gala guests including their names, relationships, descriptions, and contact details. You have access to a comprehensive guest database.",
            priority=0,
        ),
        FactExtractionMemoryBlock(
            name="extracted_facts",
            llm=llm,
            max_facts=30,
            priority=1,
        ),
    ]
    
    # Create memory instance
    memory = Memory.from_defaults(
        session_id="test_session",
        token_limit=8000,
        memory_blocks=memory_blocks,
        insert_method="system",
    )
    
    print("✅ Memory instance created successfully")
    print(f"Session ID: {memory.session_id}")
    print(f"Token limit: {memory.token_limit}")
    print(f"Number of memory blocks: {len(memory.memory_blocks)}")
    print()
    
    # Test adding messages to memory
    test_messages = [
        ChatMessage(role="user", content="What is the email of Lady Ada Lovelace?"),
        ChatMessage(role="assistant", content="Lady Ada Lovelace's email is ada.lovelace@example.com. She is known as the first computer programmer and will be attending the gala."),
        ChatMessage(role="user", content="Who are the scientists attending the gala?"),
        ChatMessage(role="assistant", content="Several scientists are attending including Marie Curie (physicist), Charles Darwin (naturalist), and Nikola Tesla (inventor/electrical engineer)."),
    ]
    
    print("📝 Adding test messages to memory...")
    memory.put_messages(test_messages)
    print(f"✅ Added {len(test_messages)} messages to memory")
    print()
    
    # Test retrieving memory
    print("🔍 Retrieving memory contents:")
    print("-" * 30)
    chat_history = memory.get()
    
    for i, message in enumerate(chat_history, 1):
        print(f"Message {i} ({message.role}):")
        print(f"Content: {message.content[:200]}{'...' if len(message.content) > 200 else ''}")
        print()
    
    # Test memory with additional context
    print("🧪 Testing memory with additional context:")
    print("-" * 30)
    
    # Add more messages to trigger fact extraction
    additional_messages = [
        ChatMessage(role="user", content="Tell me about Marie Curie's research"),
        ChatMessage(role="assistant", content="Marie Curie was a pioneering physicist and chemist who conducted groundbreaking research on radioactivity. She was the first woman to win a Nobel Prize and the only person to win Nobel Prizes in two different scientific fields."),
        ChatMessage(role="user", content="What about her contact information?"),
        ChatMessage(role="assistant", content="Marie Curie's email is marie.curie@example.com. She will be presenting her research findings at the gala."),
    ]
    
    memory.put_messages(additional_messages)
    print(f"✅ Added {len(additional_messages)} more messages")
    
    # Retrieve updated memory
    updated_chat_history = memory.get()
    print(f"📊 Total messages in memory: {len(updated_chat_history)}")
    
    # Show the system message with memory blocks
    if updated_chat_history and updated_chat_history[0].role == "system":
        print("\n🎯 System message with memory blocks:")
        print(updated_chat_history[0].content[:500] + "..." if len(updated_chat_history[0].content) > 500 else updated_chat_history[0].content)
    
    print("\n🎉 Memory integration test completed!")

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