import chromadb
from tqdm import tqdm  # Optional: For progress bar

# Connect to source and destination local persistent clients
source_client = chromadb.PersistentClient(
    path="../vedam_ai/chromadb-store"
)
destination_client = chromadb.PersistentClient(path="./chromadb-store")

source_collection_name = "chathusloki"
destination_collection_name = "chathusloki"

# Get the source collection
source_collection = source_client.get_collection(source_collection_name)

# Retrieve all data from the source collection
source_data = source_collection.get(
    include=["documents", "metadatas", "embeddings"]
)

# Create or get the destination collection
if destination_client.get_or_create_collection(destination_collection_name):
    print("Deleting existing collection", destination_collection_name)
    destination_client.delete_collection(destination_collection_name)

destination_collection = destination_client.get_or_create_collection(
    destination_collection_name,
    metadata=source_collection.metadata,  # Copy metadata if needed
)

# Add data to the destination collection in batches
BATCH_SIZE = 500
total_records = len(source_data["ids"])

print(f"Copying {total_records} records in batches of {BATCH_SIZE}...")

for i in tqdm(range(0, total_records, BATCH_SIZE)):
    batch_ids = source_data["ids"][i:i + BATCH_SIZE]
    batch_docs = source_data["documents"][i:i + BATCH_SIZE]
    batch_metas = source_data["metadatas"][i:i + BATCH_SIZE]
    batch_embeds = (
        source_data["embeddings"][i:i + BATCH_SIZE]
        if "embeddings" in source_data and source_data["embeddings"] is not None
        else None
    )

    destination_collection.add(
        ids=batch_ids,
        documents=batch_docs,
        metadatas=batch_metas,
        embeddings=batch_embeds,
    )

print("✅ Collection copied successfully!")
print("Total records in source collection = ", source_collection.count())
print("Total records in destination collection = ", destination_collection.count())