#%%
import os
import json
import chromadb
from sentence_transformers import SentenceTransformer
from tqdm import tqdm

CHUNKS_DIR = "/content/drive/MyDrive/updatedparentchunk/chunks/parent_child"

client = chromadb.PersistentClient(
    path="/content/drive/MyDrive/chroma_bge_v2"
)

collection = client.get_or_create_collection(
    "sci_judgments_bge_v2"
)

model = SentenceTransformer(
    "BAAI/bge-small-en-v1.5"
)

def flatten_metadata(meta):
    flat = {}

    for k, v in meta.items():

        if isinstance(v, (str, int, float, bool)):
            flat[k] = v

        elif isinstance(v, (list, dict)):
            flat[k] = json.dumps(v)

        else:
            flat[k] = str(v)

    return flat


files = [
    f for f in os.listdir(CHUNKS_DIR)
    if f.endswith(".json")
]

print("Files:", len(files))

all_ids = []
all_docs = []
all_meta = []

for filename in tqdm(files):

    with open(
        os.path.join(CHUNKS_DIR, filename),
        encoding="utf-8"
    ) as f:

        data = json.load(f)

    parent = data["parent"]

    all_ids.append(parent["chunk_id"])
    all_docs.append(parent["text"])
    all_meta.append(
        flatten_metadata(parent["metadata"])
    )

    for child in data["children"]:

        all_ids.append(child["chunk_id"])
        all_docs.append(child["text"])
        all_meta.append(
            flatten_metadata(child["metadata"])
        )

print("Total chunks:", len(all_ids))

BATCH_SIZE = 512

for start in tqdm(
    range(0, len(all_ids), BATCH_SIZE),
    desc="Embedding"
):

    end = min(
        start + BATCH_SIZE,
        len(all_ids)
    )

    batch_docs = all_docs[start:end]

    embeddings = model.encode(
        batch_docs,
        batch_size=128,
        normalize_embeddings=True,
        show_progress_bar=False
    )

    collection.add(
        ids=all_ids[start:end],
        embeddings=embeddings.tolist(),
        documents=batch_docs,
        metadatas=all_meta[start:end]
    )

print("\nDONE")
print("Collection count:", collection.count())