#!/usr/bin/env python3
"""
Script to apply performance optimizations to LangMem system
"""

import asyncio
import os
import sys
import logging
from pathlib import Path

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

import asyncpg
from neo4j import AsyncGraphDatabase
from performance_optimizer import PerformanceOptimizer, ConnectionManager

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Configuration
SUPABASE_DB_URL = os.getenv("SUPABASE_DB_URL", "postgresql://postgres:your_password@localhost:5435/postgres")
NEO4J_URL = os.getenv("NEO4J_URL", "bolt://localhost:7687")
NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "password")

async def run_optimizations():
    """Run all performance optimizations"""
    print("🚀 Starting LangMem Performance Optimization")
    print("=" * 50)
    
    # Create connection manager
    conn_manager = ConnectionManager(max_connections=10)
    
    # Initialize database connections
    try:
        print("📡 Connecting to databases...")
        
        # PostgreSQL connection
        db_pool = await conn_manager.create_optimized_pool(SUPABASE_DB_URL)
        print("✅ PostgreSQL connected")
        
        # Neo4j connection
        neo4j_driver = AsyncGraphDatabase.driver(
            NEO4J_URL,
            auth=(NEO4J_USER, NEO4J_PASSWORD)
        )
        print("✅ Neo4j connected")
        
        # Initialize optimizer
        optimizer = PerformanceOptimizer(db_pool, neo4j_driver)
        
        # Run optimizations
        print("\n🛠️  Running optimizations...")
        
        # 1. Create database indexes
        print("📊 Creating database indexes...")
        await optimizer.create_database_indexes()
        
        # 2. Create Neo4j indexes
        print("🔗 Creating Neo4j indexes...")
        await optimizer.create_neo4j_indexes()
        
        # 3. Optimize memory table
        print("📈 Optimizing memory table...")
        await optimizer.optimize_memory_table()
        
        # 4. Optimize embeddings storage
        print("🧠 Optimizing embeddings storage...")
        await optimizer.optimize_embeddings_storage()
        
        # 5. Get performance stats
        print("\n📊 Getting performance statistics...")
        stats = await optimizer.get_query_performance_stats()
        
        print(f"Database size: {stats.get('database_size', 'N/A')}")
        print(f"Table size: {stats.get('table_size', 'N/A')}")
        
        if stats.get('table_stats'):
            table_stats = stats['table_stats']
            print(f"Live tuples: {table_stats.get('live_tuples', 'N/A')}")
            print(f"Dead tuples: {table_stats.get('dead_tuples', 'N/A')}")
        
        # 6. Show index usage
        print("\n🔍 Index usage statistics:")
        for index in stats.get('index_stats', []):
            print(f"  {index['indexname']}: {index['scans']} scans")
        
        print("\n✅ All optimizations completed successfully!")
        
    except Exception as e:
        logger.error(f"❌ Optimization failed: {e}")
        raise
    
    finally:
        # Close connections
        if 'db_pool' in locals():
            await db_pool.close()
        if 'neo4j_driver' in locals():
            await neo4j_driver.close()

async def cleanup_old_data():
    """Clean up old data to improve performance"""
    print("\n🧹 Cleaning up old data...")
    
    try:
        # Simple cleanup without full optimizer
        db_pool = await asyncpg.create_pool(SUPABASE_DB_URL, min_size=1, max_size=5)
        
        async with db_pool.acquire() as conn:
            # Delete test data older than 1 day
            result = await conn.execute("""
                DELETE FROM memories 
                WHERE created_at < NOW() - INTERVAL '1 day'
                AND (
                    user_id LIKE '%test%' OR 
                    user_id LIKE '%claude_test%' OR
                    content LIKE '%test%'
                )
            """)
            
            deleted_count = int(result.split()[-1]) if result else 0
            print(f"🗑️  Deleted {deleted_count} test memories")
        
        await db_pool.close()
        
    except Exception as e:
        logger.error(f"❌ Cleanup failed: {e}")

async def main():
    """Main function"""
    try:
        await run_optimizations()
        await cleanup_old_data()
        
        print("\n" + "=" * 50)
        print("🎯 Performance optimization complete!")
        print("💡 Your LangMem system should now be faster and more efficient.")
        
    except Exception as e:
        print(f"\n❌ Failed to complete optimizations: {e}")
        return 1
    
    return 0

if __name__ == "__main__":
    exit_code = asyncio.run(main())
    sys.exit(exit_code)