import asyncio
from qdrant_client import QdrantClient
from qdrant_client.http import models
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def fix_qdrant_dimensions():
    """Fix Qdrant collection dimensions to match Snowflake Arctic Embed model (1024 dimensions)"""
    
    # Connect to Qdrant
    client = QdrantClient(host="localhost", port=6333)
    
    # Collections that need to be updated
    collections = ["chunks", "entities", "relationships"]
    
    for collection_name in collections:
        try:
            # Check if collection exists
            existing_collections = client.get_collections().collections
            collection_exists = any(c.name == collection_name for c in existing_collections)
            
            if collection_exists:
                # Delete the existing collection with wrong dimensions
                logger.info(f"Deleting collection '{collection_name}' with wrong dimensions...")
                client.delete_collection(collection_name=collection_name)
                logger.info(f"Collection '{collection_name}' deleted successfully")
            
            # Recreate collection with correct dimensions (1024 for Snowflake Arctic Embed)
            logger.info(f"Creating collection '{collection_name}' with 1024 dimensions...")
            client.create_collection(
                collection_name=collection_name,
                vectors_config=models.VectorParams(
                    size=1024,  # Snowflake Arctic Embed dimensions
                    distance=models.Distance.COSINE
                )
            )
            logger.info(f"Collection '{collection_name}' created successfully with 1024 dimensions")
            
        except Exception as e:
            logger.error(f"Error processing collection '{collection_name}': {e}")
    
    logger.info("Qdrant dimension fix completed!")

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