import os
import sys
import chromadb
from chromadb.utils import embedding_functions
from qdrant_client import QdrantClient
from qdrant_client.http import models
import time

def migrate(qdrant_url, qdrant_api_key):
    # 1. Load Local Chroma
    kb_path = "data/knowledge_base"
    print(f"📂 Loading local ChromaDB from {kb_path}...")
    
    if not os.path.exists(kb_path):
        print("❌ Local Knowledge Base not found!")
        return

    chroma_client = chromadb.PersistentClient(path=kb_path)
    
    # Use standard EF matching the build script
    ef = embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name="all-MiniLM-L6-v2"
    )
    
    col = chroma_client.get_collection("medical_knowledge", embedding_function=ef)
    count = col.count()
    print(f"✅ Found {count} documents in ChromaDB.")

    # 2. Connect to Qdrant Cloud
    print(f"☁️ Connecting to Qdrant Cloud: {qdrant_url}...")
    qdrant_client = QdrantClient(
        url=qdrant_url, 
        api_key=qdrant_api_key,
        timeout=60 # Extended timeout for uploads
    )
    
    # Check/Create Collection
    collection_name = "medical_knowledge"
    try:
        qdrant_client.get_collection(collection_name)
        print(f"✅ Qdrant Collection '{collection_name}' exists.")
    except:
        print(f"⚠️ Creating new collection '{collection_name}' with Quantization...")
        qdrant_client.create_collection(
            collection_name=collection_name,
            vectors_config=models.VectorParams(
                size=384,  # all-MiniLM-L6-v2 dimension
                distance=models.Distance.COSINE,
                quantization_config=models.ScalarQuantization(
                    scalar=models.ScalarQuantizationConfig(
                        type=models.ScalarType.INT8,
                        quantile=0.99,
                        always_ram=True
                    )
                )
            )
        )

    # 3. Migrate Data in Batches
    batch_size = 100
    total_migrated = 0
    
    print("🚀 Starting Migration...")
    
    # Fetch all data (Chroma get allows large fetch? Yes, usually)
    # Ideally use offset/limit pagination
    limit = 1000
    offset = 0
    
    while True:
        results = col.get(
            include=['documents', 'metadatas', 'embeddings'],
            limit=limit,
            offset=offset
        )
        
        ids = results['ids']
        if not ids:
            break
            
        points = []
        for i, doc_id in enumerate(ids):
            points.append(models.PointStruct(
                id=i + offset, # Use integer ID based on offset? No, Qdrant allows UUID or Int. original ID is better?
                # Chroma IDs might be strings. Qdrant supports UUID strings.
                # Let's map to UUID if needed, or use integer offset as ID.
                # Integer IDs are efficient in Qdrant.
                vector=results['embeddings'][i],
                payload={
                    "page_content": results['documents'][i],
                    **results['metadatas'][i]
                }
            ))
            
        qdrant_client.upsert(
            collection_name=collection_name,
            points=points
        )
        
        total_migrated += len(points)
        print(f"  Processed {total_migrated}/{count}...")
        offset += limit
        
    print(f"🎉 Migration Complete! {total_migrated} vectors uploaded to Qdrant.")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python migrate_to_qdrant.py <QDRANT_URL> <QDRANT_API_KEY>")
        sys.exit(1)
        
    migrate(sys.argv[1], sys.argv[2])