import faiss
import numpy as np
import os


class FaissDB:
    def __init__(self, dimension=512):
        self.dimension = dimension
        self.index = faiss.IndexFlatIP(dimension)

    def create(self):
        self.index = faiss.IndexFlatIP(self.dimension)

    def load(self, path):
        if os.path.exists(path):
            self.index = faiss.read_index(path)
            return True
        return False

    def save(self, path):
        faiss.write_index(self.index, path)

    def add(self, embeddings):
        embeddings = np.asarray(embeddings, dtype=np.float32)
        self.index.add(embeddings)

    def search(self, query_embedding, k=10):
        embeddings = np.asarray(
            query_embedding, dtype=np.float32).reshape(1, -1)
        scores, indices = self.index.search(embeddings, k)
        return scores[0], indices[0]

    def size(self):
        return self.index.ntotal