import { ChromaClient, IncludeEnum } from "chromadb";
import { MetadataMode } from "../../Node.js";
import { VectorStoreBase, VectorStoreQueryMode } from "./types.js";
import { metadataDictToNode, nodeToMetadata } from "./utils.js";
const DEFAULT_TEXT_KEY = "text";
export class ChromaVectorStore extends VectorStoreBase {
    storesText = true;
    flatMetadata = true;
    textKey;
    chromaClient;
    collection = null;
    collectionName;
    constructor(init){
        super(init.embedModel);
        this.collectionName = init.collectionName;
        this.chromaClient = new ChromaClient(init.chromaClientParams);
        this.textKey = init.textKey ?? DEFAULT_TEXT_KEY;
    }
    client() {
        return this.chromaClient;
    }
    async getCollection() {
        if (!this.collection) {
            const coll = await this.chromaClient.getOrCreateCollection({
                name: this.collectionName
            });
            this.collection = coll;
        }
        return this.collection;
    }
    getDataToInsert(nodes) {
        const metadatas = nodes.map((node)=>nodeToMetadata(node, true, this.textKey, this.flatMetadata));
        return {
            embeddings: nodes.map((node)=>node.getEmbedding()),
            ids: nodes.map((node)=>node.id_),
            metadatas,
            documents: nodes.map((node)=>node.getContent(MetadataMode.NONE))
        };
    }
    async add(nodes) {
        if (!nodes || nodes.length === 0) {
            return [];
        }
        const dataToInsert = this.getDataToInsert(nodes);
        const collection = await this.getCollection();
        await collection.add(dataToInsert);
        return nodes.map((node)=>node.id_);
    }
    async delete(refDocId, deleteOptions) {
        const collection = await this.getCollection();
        await collection.delete({
            ids: [
                refDocId
            ],
            where: deleteOptions?.where,
            whereDocument: deleteOptions?.whereDocument
        });
    }
    async query(query, options) {
        if (query.docIds) {
            throw new Error("ChromaDB does not support querying by docIDs");
        }
        if (query.mode != VectorStoreQueryMode.DEFAULT) {
            throw new Error("ChromaDB does not support querying by mode");
        }
        // fixme: type is broken
        let chromaWhere = {};
        if (query.filters?.filters) {
            query.filters.filters.map((filter)=>{
                const filterKey = filter.key;
                const filterValue = filter.value;
                if (filterKey in chromaWhere || "$or" in chromaWhere) {
                    if (!chromaWhere["$or"]) {
                        chromaWhere = {
                            $or: [
                                {
                                    ...chromaWhere
                                },
                                {
                                    [filterKey]: filterValue
                                }
                            ]
                        };
                    } else {
                        chromaWhere["$or"].push({
                            [filterKey]: filterValue
                        });
                    }
                } else {
                    chromaWhere[filterKey] = filterValue;
                }
            });
        }
        const collection = await this.getCollection();
        const queryResponse = await collection.query({
            queryEmbeddings: query.queryEmbedding ?? undefined,
            queryTexts: query.queryStr ?? undefined,
            nResults: query.similarityTopK,
            where: Object.keys(chromaWhere).length ? chromaWhere : undefined,
            whereDocument: options?.whereDocument,
            //ChromaDB doesn't return the result embeddings by default so we need to include them
            include: [
                IncludeEnum.Distances,
                IncludeEnum.Metadatas,
                IncludeEnum.Documents,
                IncludeEnum.Embeddings
            ]
        });
        const vectorStoreQueryResult = {
            nodes: queryResponse.ids[0].map((id, index)=>{
                const text = queryResponse.documents[0][index];
                const metaData = queryResponse.metadatas[0][index] ?? {};
                const node = metadataDictToNode(metaData);
                node.setContent(text);
                return node;
            }),
            similarities: queryResponse.distances[0].map((distance)=>1 - distance),
            ids: queryResponse.ids[0]
        };
        return vectorStoreQueryResult;
    }
}
