#!/usr/bin/env python3
"""
Test script to validate the new SQLAlchemy persistence system.
"""

import asyncio
import sys
from pathlib import Path

import pytest

# Add src to path for imports
sys.path.insert(0, str(Path(__file__).parent / "src"))

from broadie.persistence.database import initialize_database
from broadie.persistence.repository import PersistenceRepository


@pytest.mark.asyncio
async def test_persistence():
    """Test the new persistence system."""
    print("🧪 Testing new SQLAlchemy persistence system...")

    try:
        # Initialize database
        print("📊 Initializing database...")
        await initialize_database()
        print("✅ Database initialized successfully")

        # Create repository
        repo = PersistenceRepository()

        # Test thread message saving
        print("💬 Testing message saving...")
        thread_id = "test-thread-123"
        message_data = {
            "role": "user",
            "content": "Hello, this is a test message!",
            "user_id": "test-user-456",
            "message_id": "test-msg-789",
        }

        message_id = await repo.save_thread_message(thread_id, message_data)
        print(f"✅ Saved message with ID: {message_id}")

        # Test thread history retrieval
        print("📜 Testing history retrieval...")
        history = await repo.get_thread_history(thread_id, limit=10)
        print(f"✅ Retrieved {len(history)} messages from history")

        if history:
            print(f"   First message: {history[0].value['content']}")

        # Test user facts
        print("🧠 Testing user facts...")
        fact_data = {"type": "preference", "content": "User prefers technical explanations", "importance": 0.8}

        fact_id = await repo.remember_user_fact("test-user-456", fact_data)
        print(f"✅ Saved user fact with ID: {fact_id}")

        # Test fact recall
        recalled_facts = await repo.recall_user_facts("test-user-456", "technical", limit=5)
        print(f"✅ Recalled {len(recalled_facts)} matching facts")

        # Test checkpoint saving
        print("🎯 Testing checkpoint saving...")
        checkpoint_id = await repo.save_checkpoint(
            thread_id=thread_id,
            checkpoint_type="llm_response",
            state_data={"test_state": "This is test state data", "step": 1},
            message_id=message_id,
            run_id="test-run-999",
        )
        print(f"✅ Saved checkpoint with ID: {checkpoint_id}")

        # Test analytics
        print("📈 Testing analytics...")
        thread_analytics = await repo.get_thread_analytics(thread_id)
        print(f"✅ Thread analytics: {thread_analytics.get('total_messages', 0)} messages")

        user_analytics = await repo.get_user_analytics("test-user-456")
        print(f"✅ User analytics: {user_analytics.get('total_threads', 0)} threads")

        print("\n🎉 All persistence tests passed!")
        return True

    except Exception as e:
        print(f"❌ Test failed with error: {e}")
        import traceback

        traceback.print_exc()
        return False


@pytest.mark.asyncio
async def test_compatibility():
    """Test that the new PersistenceMixin maintains compatibility."""
    print("\n🔄 Testing PersistenceMixin compatibility...")

    try:
        from broadie.mixins import PersistenceMixin

        # Create mixin instance
        mixin = PersistenceMixin()

        # Test the same interface as before
        thread_id = "compat-test-thread"
        message = {"role": "user", "content": "Testing compatibility", "user_id": "compat-user"}

        # Save message (should work exactly like before)
        message_id = await mixin.save_thread_message(thread_id, message)
        print(f"✅ PersistenceMixin saved message: {message_id}")

        # Get history (should work exactly like before)
        history = await mixin.get_thread_history(thread_id)
        print(f"✅ PersistenceMixin retrieved {len(history)} messages")

        # Test user facts (should work exactly like before)
        fact_id = await mixin.remember_user_fact("compat-user", {"test": "fact"})
        print(f"✅ PersistenceMixin saved user fact: {fact_id}")

        facts = await mixin.recall_user_facts("compat-user", "test")
        print(f"✅ PersistenceMixin recalled {len(facts)} facts")

        # Test new checkpoint functionality
        checkpoint_id = await mixin.save_agent_checkpoint(
            thread_id=thread_id, checkpoint_type="test_checkpoint", state_data={"test": "checkpoint"}
        )
        print(f"✅ PersistenceMixin saved checkpoint: {checkpoint_id}")

        print("🎉 PersistenceMixin compatibility test passed!")
        return True

    except Exception as e:
        print(f"❌ Compatibility test failed: {e}")
        import traceback

        traceback.print_exc()
        return False


if __name__ == "__main__":

    async def main():
        success = await test_persistence()
        if success:
            success = await test_compatibility()

        if success:
            print("\n✨ All tests passed! The new persistence system is working correctly.")
            print("🔄 The system maintains 100% compatibility while adding new capabilities:")
            print("   • SQLAlchemy-backed storage for any database")
            print("   • LangGraph state checkpoints")
            print("   • Enhanced analytics and querying")
            print("   • Better performance and reliability")
        else:
            print("\n❌ Tests failed. Please check the implementation.")
            sys.exit(1)

    asyncio.run(main())
