#!/usr/bin/env python3
"""
Script to reset Qdrant collections with correct dimensions (1024) for Snowflake Arctic Embed
"""

import os
import sys
from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance

def reset_qdrant_collections():
    """Reset all Qdrant collections with correct 1024 dimensions"""
    
    # Connect to Qdrant
    client = QdrantClient(url="http://localhost:6333")
    
    # Collections used by LightRAG
    collections = [
        "rag_storage_chunks",
        "rag_storage_entities", 
        "rag_storage_relationships"
    ]
    
    for collection_name in collections:
        try:
            # Check if collection exists
            try:
                collection_info = client.get_collection(collection_name)
                current_dim = collection_info.config.params.vectors.size
                print(f"Collection '{collection_name}': current dimension = {current_dim}")
                
                # Delete if wrong dimension
                if current_dim != 1024:
                    print(f"  Deleting collection (wrong dimension: {current_dim})")
                    client.delete_collection(collection_name)
                else:
                    print(f"  ✓ Correct dimension (1024)")
                    continue
                    
            except Exception as e:
                print(f"  Collection '{collection_name}' doesn't exist or error: {e}")
            
            # Create collection with correct dimension
            print(f"  Creating collection with dimension 1024")
            client.create_collection(
                collection_name=collection_name,
                vectors_config=VectorParams(size=1024, distance=Distance.COSINE),
            )
            print(f"  ✓ Created successfully")
            
        except Exception as e:
            print(f"  Error processing '{collection_name}': {e}")
    
    print("\nQdrant collections reset complete!")

if __name__ == "__main__":
    reset_qdrant_collections()