#!/usr/bin/env python3
"""
Clean up all Supabase vecs tables and start fresh
"""

import psycopg2
import vecs

def cleanup_all_tables():
    """Clean up all tables in vecs schema"""
    print("=" * 60)
    print("SUPABASE VECS SCHEMA CLEANUP")
    print("=" * 60)
    
    try:
        # Connect to database directly
        conn = psycopg2.connect("postgresql://supabase_admin:CzkaYmRvc26Y@localhost:5435/postgres")
        cur = conn.cursor()
        
        print("🔍 Finding all tables in vecs schema...")
        cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'vecs';")
        tables = cur.fetchall()
        table_names = [t[0] for t in tables]
        print(f"Found tables: {table_names}")
        
        if table_names:
            print(f"\n🗑️  Dropping {len(table_names)} tables...")
            for table_name in table_names:
                try:
                    cur.execute(f'DROP TABLE IF EXISTS vecs."{table_name}" CASCADE;')
                    print(f"   ✅ Dropped: {table_name}")
                except Exception as e:
                    print(f"   ❌ Failed to drop {table_name}: {e}")
            
            # Commit the drops
            conn.commit()
            print("✅ All table drops committed")
        else:
            print("ℹ️  No tables found in vecs schema")
        
        # Verify cleanup
        cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'vecs';")
        remaining_tables = cur.fetchall()
        print(f"\n📋 Remaining tables: {[t[0] for t in remaining_tables]}")
        
        cur.close()
        conn.close()
        
        print("\n🧪 Testing fresh vecs client connection...")
        connection_string = "postgresql://supabase_admin:CzkaYmRvc26Y@localhost:5435/postgres"
        db = vecs.create_client(connection_string)
        
        collections = db.list_collections()
        print(f"Collections after cleanup: {[c.name for c in collections]}")
        
        print("\n🎯 Testing fresh collection creation...")
        test_collection = db.get_or_create_collection(name="test_fresh_start", dimension=1536)
        print(f"✅ Successfully created: {test_collection.name} with dimension {test_collection.dimension}")
        
        # Test basic operations
        print("🧪 Testing basic vector operations...")
        import numpy as np
        
        # Insert test vector
        test_vector = np.random.random(1536).tolist()
        test_id = "test_vector_1"
        test_metadata = {"content": "Fresh start test", "user_id": "test"}
        
        test_collection.upsert([(test_id, test_vector, test_metadata)])
        print("✅ Vector upserted successfully")
        
        # Search test
        results = test_collection.query(data=test_vector, limit=1, include_metadata=True)
        print(f"✅ Search successful, found {len(results)} results")
        if results:
            print(f"   Result: ID={results[0][0]}, Score={results[0][1]}")
        
        # Clean up test collection
        db.delete_collection("test_fresh_start")
        print("✅ Test collection cleaned up")
        
        print("\n" + "=" * 60)
        print("🎉 CLEANUP SUCCESSFUL - VECS IS READY!")
        print("=" * 60)
        
    except Exception as e:
        print(f"❌ Cleanup failed: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    cleanup_all_tables()