import json
import chromadb

from tqdm import tqdm
from sentence_transformers import SentenceTransformer

# =====================================================
# LOAD STATUTES
# =====================================================

with open(
    "/content/drive/MyDrive/all_statutes/all_statutes.json",
    "r",
    encoding="utf-8"
) as f:


    statutes = json.load(f)

print("Total Records:", len(statutes))

# =====================================================
# LOAD EMBEDDING MODEL
# =====================================================

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

# =====================================================
# CREATE CHROMA DB
# =====================================================

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

collection_name = "indian_statutes"

# Delete old collection if exists
try:
    client.delete_collection(
        collection_name
    )
except:
    pass

collection = client.create_collection(
    collection_name
)

# =====================================================
# PREPARE DATA
# =====================================================

ids = []
documents = []
metadatas = []

for record in statutes:

    ids.append(
        record["chunk_id"]
    )

    documents.append(
        record["retrieval_text"]
    )

    meta = {

        "act_short":
            record["metadata"]["act_short"],

        "act_name":
            record["metadata"]["act_name"],

        "section_number":
            str(
                record["metadata"]["section_number"]
            ),

        "title":
            record["metadata"]["title"]
    }

    metadatas.append(meta)

print("Prepared:", len(ids))

# =====================================================
# EMBED + STORE
# =====================================================

BATCH_SIZE = 100

for i in tqdm(
    range(
        0,
        len(documents),
        BATCH_SIZE
    )
):

    batch_docs = documents[
        i:i+BATCH_SIZE
    ]

    batch_ids = ids[
        i:i+BATCH_SIZE
    ]

    batch_meta = metadatas[
        i:i+BATCH_SIZE
    ]

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

    collection.add(
        ids=batch_ids,
        documents=batch_docs,
        metadatas=batch_meta,
        embeddings=embeddings.tolist()
    )

print()
print("=" * 60)
print("DONE")
print("Documents Stored:",
      collection.count())
print("=" * 60)