import faiss       # Facebook AI Similarity Search
import numpy as np # for working with arrays of numbers


def build_index(chunks):
    """
    Builds a FAISS index from all chunk embeddings.

    Input  : chunks → list of chunk dicts, each with an "embedding" field
    Output : a FAISS index ready for searching
    """

    # Stack all embeddings into a 2D array
    # Shape: (number_of_chunks, 384)
    # astype("float32") because FAISS requires 32-bit floats
    embeddings = np.array(
        [chunk["embedding"] for chunk in chunks]
    ).astype("float32")

    dimension = embeddings.shape[1]       # size of each vector (384 for MiniLM)
    index = faiss.IndexFlatL2(dimension)  # create the index with L2 distance
    index.add(embeddings)                 # add all chunk vectors to the index

    print(f"FAISS index built with {index.ntotal} vectors.")
    return index


def search_index(index, chunks, query_embedding, top_k=3):
    """
    Searches the FAISS index for the most relevant chunks.

    Input:
      index           → the FAISS index we built
      chunks          → our original list of chunk dicts
      query_embedding → the vector of the user question
      top_k           → how many results to return (default: 3)

    Output : list of the top_k most relevant chunk dicts
    """

    # FAISS expects a 2D array even for a single query
    # So we wrap it in another list and convert to float32
    query_vec = np.array([query_embedding]).astype("float32")

    # Search the index
    # D → distances  (how far each result is from the query)
    # I → indices    (which chunk numbers are the closest matches)
    D, I = index.search(query_vec, top_k)

    results = []
    for idx in I[0]:        # I[0] because we have just one query
        if idx != -1:       # FAISS returns -1 if no result was found
            results.append(chunks[idx])

    return results