import { CloseVectorEmbeddings, CloseVectorEmbeddingsParams } from 'closevector-common';
import { encryptToken, chunkArray, END_POINT } from './lib'

export { CloseVectorEmbeddings, CloseVectorEmbeddingsParams };
export class CloseVectorEmbeddingsWeb implements CloseVectorEmbeddings {
    batchSize = 512;
    stripNewLines = true;
    timeout?: number;

    config: {
        key: string;
        secret: string;
    };

    constructor(
        fields: Partial<CloseVectorEmbeddingsParams> & {
            key: string;
            secret: string;
        }
    ) {
        this.batchSize = fields.batchSize || this.batchSize;
        this.stripNewLines = fields.stripNewLines || this.stripNewLines;
        this.timeout = fields.timeout;

        this.config = {
            key: fields.key,
            secret: fields.secret,
        };
    }

    async embedDocuments(texts: string[]): Promise<number[][]> {
        const subPrompts = chunkArray(
            this.stripNewLines ? texts.map(t => t.replace(/\n/g, ' ')) : texts,
            this.batchSize
        );

        const embeddings: number[][] = [];

        for (let i = 0; i < subPrompts.length; i += 1) {
            const input = subPrompts[i];
            const { data } = await this.embeddingWithRetry(input);
            for (let j = 0; j < input.length; j += 1) {
                embeddings.push(data[j].embedding);
            }
        }

        return embeddings;
    }

    async embedQuery(text: string): Promise<number[]> {
        const { data } = await this.embeddingWithRetry(
            this.stripNewLines ? [text.replace(/\n/g, ' ')] : [text]
        );
        return data[0].embedding;
    }

    async embeddingWithRetry(textList: string[]): Promise<any> {
        let accessKey = this.config.key;
        let secret = this.config.secret;

        const fetchUsingToken = async function fetchUsingToken(url: string, options: RequestInit) {
            let res = await fetch(url, {
                ...options,
                headers: {
                    ...options.headers,
                    Authorization: `Bearer ${accessKey}:${await encryptToken(
                        { accessKey },
                        secret
                    )}`,
                },
            });

            if (res.status !== 200) {
                let json = await res.json();
                throw new Error(await (json as any).message);
            }

            return res;
        };

        async function createEmbeddings(textList: string[]) {
            // let resp = await fetchUsingToken(`https://vector-kv.mega-ug.uk/vectorEmbedding`, {
            let resp = await fetchUsingToken(`${END_POINT}/embeddings`, {
                method: 'POST',
                body: JSON.stringify(textList),
            });

            let data = await resp.json();
            return data;
        }

        return createEmbeddings(textList);
    }
}