import json
import faiss
import numpy as np
from modules.embed import embed_text

CHUNK_SIZE = 300  # number of words per chunk
OVERLAP = 50

FAISS_INDEX_PATH = "faiss_index.index"
METADATA_PATH = "faiss_metadata.json"
KB_PATH = "knowledge_base.txt"

def chunk_text(text, chunk_size, overlap):
    words = text.split()
    chunks = []
    for i in range(0, len(words), chunk_size - overlap):
        chunk = " ".join(words[i:i + chunk_size])
        chunks.append(chunk)
    return chunks

def build_faiss_index():
    with open(KB_PATH, "r") as f:
        kb_text = f.read()

    chunks = chunk_text(kb_text, CHUNK_SIZE, OVERLAP)
    embeddings = [embed_text(chunk) for chunk in chunks]

    dimension = embeddings[0].shape[0]
    index = faiss.IndexFlatL2(dimension)
    index.add(np.array(embeddings))

    metadata = {str(i): chunk for i, chunk in enumerate(chunks)}
    faiss.write_index(index, FAISS_INDEX_PATH)
    with open(METADATA_PATH, "w") as f:
        json.dump(metadata, f)

if __name__ == "__main__":
    build_faiss_index()
    print("FAISS index built successfully.")