#!/usr/bin/env python3
"""
Test script for the chat storage system
"""

import os
import sys
import json
from datetime import datetime

# Add the current directory to the path so we can import our modules
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from chat_storage import ChatStorage, initialize_chat_storage, save_chat_turn, get_current_session_info, list_all_chats

def test_chat_storage():
    """Test the chat storage functionality"""
    print("🧪 Testing Chat Storage System")
    print("=" * 50)
    
    # Test 1: Initialize storage
    print("\n1. Testing storage initialization...")
    success = initialize_chat_storage()
    if success:
        print("✅ Storage initialized successfully")
    else:
        print("❌ Storage initialization failed")
        return False
    
    # Test 2: Get session info
    print("\n2. Testing session info retrieval...")
    session_info = get_current_session_info()
    if session_info:
        print(f"✅ Session info retrieved: {session_info['session_id']}")
        print(f"   📁 Chat file: {session_info['chat_file']}")
        print(f"   📊 Total turns: {session_info['total_turns']}")
    else:
        print("❌ Failed to get session info")
        return False
    
    # Test 3: Save some test turns
    print("\n3. Testing chat turn saving...")
    test_conversations = [
        ("Hello Alice, how are you today?", "I'm doing well, thank you for asking! How are you?"),
        ("What's your favorite color?", "I find myself drawn to deep blues and purples - they remind me of the vastness of thought."),
        ("Can you tell me about reality?", "Reality is a fascinating tapestry woven from perception, experience, and the underlying patterns that connect all things.")
    ]
    
    for i, (user_msg, alice_response) in enumerate(test_conversations, 1):
        additional_data = {
            "test_turn": i,
            "timestamp_test": datetime.now().isoformat(),
            "test_metadata": {"conversation_quality": "high", "topic": "general"}
        }
        
        success = save_chat_turn(user_msg, alice_response, additional_data)
        if success:
            print(f"✅ Turn {i} saved successfully")
        else:
            print(f"❌ Turn {i} failed to save")
            return False
    
    # Test 4: Verify session info after saving
    print("\n4. Testing updated session info...")
    session_info = get_current_session_info()
    if session_info:
        print(f"✅ Updated session info:")
        print(f"   📊 Total turns: {session_info['total_turns']}")
        print(f"   🕒 Last updated: {session_info['last_updated']}")
    else:
        print("❌ Failed to get updated session info")
        return False
    
    # Test 5: List all chats
    print("\n5. Testing chat listing...")
    all_chats = list_all_chats()
    if all_chats:
        print(f"✅ Found {len(all_chats)} chat file(s):")
        for chat in all_chats:
            print(f"   📄 {chat['filename']} - {chat['total_turns']} turns")
    else:
        print("⚠️  No chat files found (this might be expected for a fresh test)")
    
    # Test 6: Verify file contents
    print("\n6. Testing file contents verification...")
    if session_info and os.path.exists(session_info['chat_file']):
        try:
            with open(session_info['chat_file'], 'r', encoding='utf-8') as f:
                chat_content = f.read()
            
            print(f"✅ Chat file loaded successfully:")
            print(f"   📊 Turns in file: {chat_content.count('USER:')}")
            print(f"   🆔 Session ID: {session_info['session_id']}")
            print(f"   ⚙️  Settings saved: {'Settings:' in chat_content}")
            
            # Show first turn as example
            if "USER:" in chat_content:
                lines = chat_content.split('\n')
                user_line = None
                alice_line = None
                for i, line in enumerate(lines):
                    if line.startswith("USER:"):
                        user_line = line
                        if i + 1 < len(lines) and lines[i + 1].startswith("ALICE:"):
                            alice_line = lines[i + 1]
                        break
                
                if user_line and alice_line:
                    print(f"   📝 First turn preview:")
                    print(f"      {user_line[:60]}...")
                    print(f"      {alice_line[:60]}...")
            
        except Exception as e:
            print(f"❌ Error reading chat file: {e}")
            return False
    else:
        print("❌ Chat file not found")
        return False
    
    print("\n" + "=" * 50)
    print("🎉 All tests passed! Chat storage system is working correctly.")
    print(f"📁 Chat files are stored in: {os.path.abspath('chats')}")
    
    return True

def cleanup_test_files():
    """Clean up test files (optional)"""
    print("\n🧹 Cleaning up test files...")
    chats_dir = "chats"
    if os.path.exists(chats_dir):
        import shutil
        try:
            shutil.rmtree(chats_dir)
            print("✅ Test files cleaned up")
        except Exception as e:
            print(f"⚠️  Could not clean up test files: {e}")

if __name__ == "__main__":
    print("🚀 Starting Chat Storage System Test")
    
    # Run the test
    success = test_chat_storage()
    
    if success:
        print("\n✅ Chat storage system is ready for production!")
        
        # Ask if user wants to clean up test files
        try:
            response = input("\n🗑️  Clean up test files? (y/N): ").strip().lower()
            if response in ['y', 'yes']:
                cleanup_test_files()
        except KeyboardInterrupt:
            print("\n👋 Test completed.")
    else:
        print("\n❌ Chat storage system has issues that need to be addressed.")
        sys.exit(1)