#!/usr/bin/env python3
"""
Clear all databases - PostgreSQL (vector) and Neo4j (graph)
"""

import asyncio
import asyncpg
from neo4j import AsyncGraphDatabase

# Configuration
SUPABASE_DB_URL = "postgresql://postgres:CzkaYmRvc26Y@localhost:5435/postgres"
NEO4J_URL = "bolt://localhost:7687"
NEO4J_USER = "neo4j"
NEO4J_PASSWORD = "langmem_neo4j_password"

async def clear_postgresql():
    """Clear PostgreSQL database completely"""
    print("🧹 Clearing PostgreSQL database...")
    
    try:
        conn = await asyncpg.connect(SUPABASE_DB_URL)
        
        # Drop all tables and extensions
        await conn.execute("DROP SCHEMA public CASCADE;")
        await conn.execute("CREATE SCHEMA public;")
        await conn.execute("GRANT ALL ON SCHEMA public TO postgres;")
        await conn.execute("GRANT ALL ON SCHEMA public TO public;")
        
        print("   ✅ PostgreSQL database cleared completely")
        
        await conn.close()
        return True
        
    except Exception as e:
        print(f"   ❌ Error clearing PostgreSQL: {e}")
        return False

async def clear_neo4j():
    """Clear Neo4j database completely"""
    print("🧹 Clearing Neo4j database...")
    
    try:
        driver = AsyncGraphDatabase.driver(NEO4J_URL, auth=(NEO4J_USER, NEO4J_PASSWORD))
        
        async with driver.session() as session:
            # Delete all nodes and relationships
            await session.run("MATCH (n) DETACH DELETE n")
            
            # Verify it's empty
            result = await session.run("MATCH (n) RETURN count(n) as count")
            record = await result.single()
            node_count = record['count']
            
            print(f"   ✅ Neo4j database cleared completely (nodes: {node_count})")
            
        await driver.close()
        return True
        
    except Exception as e:
        print(f"   ❌ Error clearing Neo4j: {e}")
        return False

async def restart_langmem_api():
    """Restart LangMem API to recreate tables"""
    print("🔄 Restarting LangMem API to recreate tables...")
    
    import subprocess
    try:
        # Restart the API container
        result = subprocess.run(
            ["docker", "compose", "restart", "langmem-api"],
            cwd="/home/klas/langmem-project",
            capture_output=True,
            text=True
        )
        
        if result.returncode == 0:
            print("   ✅ LangMem API restarted successfully")
            
            # Wait for API to be ready
            await asyncio.sleep(3)
            
            # Check API health
            import httpx
            async with httpx.AsyncClient() as client:
                try:
                    response = await client.get("http://localhost:8765/health", timeout=10.0)
                    if response.status_code == 200:
                        data = response.json()
                        print(f"   ✅ API health status: {data['status']}")
                        return True
                    else:
                        print(f"   ⚠️  API health check returned: {response.status_code}")
                        return False
                except Exception as e:
                    print(f"   ⚠️  API health check failed: {e}")
                    return False
        else:
            print(f"   ❌ Failed to restart API: {result.stderr}")
            return False
            
    except Exception as e:
        print(f"   ❌ Error restarting API: {e}")
        return False

async def main():
    """Main function to clear all databases"""
    print("🚀 Clearing All LangMem Databases")
    print("=" * 50)
    
    # Clear PostgreSQL
    postgres_cleared = await clear_postgresql()
    
    # Clear Neo4j
    neo4j_cleared = await clear_neo4j()
    
    # Restart API to recreate tables
    api_restarted = await restart_langmem_api()
    
    # Summary
    print("\n" + "=" * 50)
    print("📊 Database Clear Summary:")
    print(f"   PostgreSQL: {'✅ CLEARED' if postgres_cleared else '❌ FAILED'}")
    print(f"   Neo4j: {'✅ CLEARED' if neo4j_cleared else '❌ FAILED'}")
    print(f"   API Restart: {'✅ SUCCESS' if api_restarted else '❌ FAILED'}")
    
    if all([postgres_cleared, neo4j_cleared, api_restarted]):
        print("\n🎉 All databases cleared successfully!")
        print("   Ready for fresh data storage")
        return True
    else:
        print("\n⚠️  Some operations failed - check logs above")
        return False

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