import pickle

import numpy as np
import faiss
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()
client = OpenAI()

MODEL = "text-embedding-3-small"


def embed(text):
    res = client.embeddings.create(
        model=MODEL,
        input=text
    )
    return res.data[0].embedding


index = faiss.read_index("index/index.faiss")

with open("index/meta.pkl", "rb") as f:
    chunks, meta = pickle.load(f)


def search(query, k=5):
    qvec = np.array([embed(query)]).astype("float32")

    distances, ids = index.search(qvec, k)

    for i in ids[0]:
        print("\n----")
        print(meta[i]["source"])
        print(chunks[i][:500])


if __name__ == "__main__":
    import sys
    query = " ".join(sys.argv[1:])
    search(query)