"""
Migrate Qdrant local → Qdrant Cloud.
Usage: python migrate_qdrant.py
"""
import os
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from tqdm import tqdm
from dotenv import load_dotenv

load_dotenv()

COLLECTION_NAME = "rag_docs"
QDRANT_PATH = "qdrant_storage"
BATCH_SIZE = 20

QDRANT_URL = os.getenv("QDRANT_URL")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")

if not QDRANT_URL or not QDRANT_API_KEY:
    raise ValueError("Thiếu QDRANT_URL hoặc QDRANT_API_KEY trong .env")


def migrate():
    print("Kết nối local Qdrant...")
    local = QdrantClient(path=QDRANT_PATH)

    print("Kết nối Qdrant Cloud...")
    cloud = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY, timeout=60)

    # Lấy info collection local
    collection_info = local.get_collection(COLLECTION_NAME)
    vector_size = collection_info.config.params.vectors.size
    print(f"Collection: {COLLECTION_NAME}, vector dim: {vector_size}")

    # Tạo collection trên cloud
    existing = [c.name for c in cloud.get_collections().collections]
    if COLLECTION_NAME in existing:
        print(f"Collection '{COLLECTION_NAME}' đã tồn tại trên cloud, xóa và tạo lại...")
        cloud.delete_collection(COLLECTION_NAME)

    cloud.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE),
    )
    print("Đã tạo collection trên cloud.")

    # Đọc tất cả points từ local
    print("Đọc points từ local...")
    offset = None
    all_points = []
    while True:
        result, next_offset = local.scroll(
            collection_name=COLLECTION_NAME,
            limit=BATCH_SIZE,
            offset=offset,
            with_vectors=True,
            with_payload=True,
        )
        all_points.extend(result)
        if next_offset is None:
            break
        offset = next_offset

    print(f"Tổng {len(all_points)} points cần migrate.")

    # Upload lên cloud theo batch
    for i in tqdm(range(0, len(all_points), BATCH_SIZE), desc="Uploading"):
        batch = all_points[i:i + BATCH_SIZE]
        points = [
            PointStruct(id=p.id, vector=p.vector, payload=p.payload)
            for p in batch
        ]
        cloud.upsert(collection_name=COLLECTION_NAME, points=points)

    count = cloud.count(COLLECTION_NAME).count
    print(f"Migration xong! Cloud có {count} points.")


if __name__ == "__main__":
    migrate()