import { fs, path } from "@llamaindex/env";
import { getTopKEmbeddings, getTopKEmbeddingsLearner, getTopKMMREmbeddings } from "../../embeddings/utils.js";
import { exists } from "../FileSystem.js";
import { DEFAULT_PERSIST_DIR } from "../constants.js";
import { VectorStoreBase, VectorStoreQueryMode } from "./types.js";
const LEARNER_MODES = new Set([
    VectorStoreQueryMode.SVM,
    VectorStoreQueryMode.LINEAR_REGRESSION,
    VectorStoreQueryMode.LOGISTIC_REGRESSION
]);
const MMR_MODE = VectorStoreQueryMode.MMR;
class SimpleVectorStoreData {
    embeddingDict = {};
    textIdToRefDocId = {};
}
export class SimpleVectorStore extends VectorStoreBase {
    storesText = false;
    data;
    persistPath;
    constructor(init){
        super(init?.embedModel);
        this.data = init?.data || new SimpleVectorStoreData();
    }
    static async fromPersistDir(persistDir = DEFAULT_PERSIST_DIR, embedModel) {
        const persistPath = path.join(persistDir, "vector_store.json");
        return await SimpleVectorStore.fromPersistPath(persistPath, embedModel);
    }
    get client() {
        return null;
    }
    async get(textId) {
        return this.data.embeddingDict[textId];
    }
    async add(embeddingResults) {
        for (const node of embeddingResults){
            this.data.embeddingDict[node.id_] = node.getEmbedding();
            if (!node.sourceNode) {
                continue;
            }
            this.data.textIdToRefDocId[node.id_] = node.sourceNode?.nodeId;
        }
        if (this.persistPath) {
            await this.persist(this.persistPath);
        }
        return embeddingResults.map((result)=>result.id_);
    }
    async delete(refDocId) {
        const textIdsToDelete = Object.keys(this.data.textIdToRefDocId).filter((textId)=>this.data.textIdToRefDocId[textId] === refDocId);
        for (const textId of textIdsToDelete){
            delete this.data.embeddingDict[textId];
            delete this.data.textIdToRefDocId[textId];
        }
        if (this.persistPath) {
            await this.persist(this.persistPath);
        }
        return Promise.resolve();
    }
    async query(query) {
        if (!(query.filters == null)) {
            throw new Error("Metadata filters not implemented for SimpleVectorStore yet.");
        }
        const items = Object.entries(this.data.embeddingDict);
        let nodeIds, embeddings;
        if (query.docIds) {
            const availableIds = new Set(query.docIds);
            const queriedItems = items.filter((item)=>availableIds.has(item[0]));
            nodeIds = queriedItems.map((item)=>item[0]);
            embeddings = queriedItems.map((item)=>item[1]);
        } else {
            // No docIds specified, so use all available items
            nodeIds = items.map((item)=>item[0]);
            embeddings = items.map((item)=>item[1]);
        }
        const queryEmbedding = query.queryEmbedding;
        let topSimilarities, topIds;
        if (LEARNER_MODES.has(query.mode)) {
            [topSimilarities, topIds] = getTopKEmbeddingsLearner(queryEmbedding, embeddings, query.similarityTopK, nodeIds);
        } else if (query.mode === MMR_MODE) {
            const mmrThreshold = query.mmrThreshold;
            [topSimilarities, topIds] = getTopKMMREmbeddings(queryEmbedding, embeddings, null, query.similarityTopK, nodeIds, mmrThreshold);
        } else if (query.mode === VectorStoreQueryMode.DEFAULT) {
            [topSimilarities, topIds] = getTopKEmbeddings(queryEmbedding, embeddings, query.similarityTopK, nodeIds);
        } else {
            throw new Error(`Invalid query mode: ${query.mode}`);
        }
        return Promise.resolve({
            similarities: topSimilarities,
            ids: topIds
        });
    }
    async persist(persistPath = path.join(DEFAULT_PERSIST_DIR, "vector_store.json")) {
        await SimpleVectorStore.persistData(persistPath, this.data);
    }
    static async persistData(persistPath, data) {
        const dirPath = path.dirname(persistPath);
        if (!await exists(dirPath)) {
            await fs.mkdir(dirPath);
        }
        await fs.writeFile(persistPath, JSON.stringify(data));
    }
    static async fromPersistPath(persistPath, embedModel) {
        const dirPath = path.dirname(persistPath);
        if (!await exists(dirPath)) {
            await fs.mkdir(dirPath, {
                recursive: true
            });
        }
        let dataDict = {};
        try {
            const fileData = await fs.readFile(persistPath);
            dataDict = JSON.parse(fileData.toString());
        } catch (e) {
            console.error(`No valid data found at path: ${persistPath} starting new store.`);
            // persist empty data, to ignore this error in the future
            await SimpleVectorStore.persistData(persistPath, new SimpleVectorStoreData());
        }
        const data = new SimpleVectorStoreData();
        data.embeddingDict = dataDict.embeddingDict ?? {};
        data.textIdToRefDocId = dataDict.textIdToRefDocId ?? {};
        const store = new SimpleVectorStore({
            data,
            embedModel
        });
        store.persistPath = persistPath;
        return store;
    }
    static fromDict(saveDict, embedModel) {
        const data = new SimpleVectorStoreData();
        data.embeddingDict = saveDict.embeddingDict;
        data.textIdToRefDocId = saveDict.textIdToRefDocId;
        return new SimpleVectorStore({
            data,
            embedModel
        });
    }
    toDict() {
        return {
            embeddingDict: this.data.embeddingDict,
            textIdToRefDocId: this.data.textIdToRefDocId
        };
    }
}
