import { BASE64, convert, formatDataURL } from './storage.mjs';
import { countTokens, trimText } from './alan.mjs';
import { ensureArray, ensureString, need } from './utilitas.mjs';

const _NEED = ['openai', '@google-cloud/discoveryengine'];
const embeddingClients = {};
const rerankerClients = {};

const [
    OPENAI, GOOGLE, OPENROUTER, JINA,
    GOOGLE_DEFAULT_LOCATION, GOOGLE_RERANK_CONFIG_ID,
    OPENAI_EMBED,
    GEMINI_EMBED,
    GOOGLE_MODEL_SEMANTIC_RANKER,
    JINA_MODEL_RERANKER_M0,
] = [
        'OPENAI', 'GOOGLE', 'OPENROUTER', 'JINA',
        'global', 'default_ranking_config',
        'text-embedding-3-small', // dim: 1536
        // https://blog.google/innovation-and-ai/models-and-research/gemini-models/gemini-embedding-2/
        'gemini-embedding-2-preview', // dim: from 128 to 3072, default 768.
        'semantic-ranker-default@latest',
        'jina-reranker-m0',
    ];

const PROVIDER_BASE_URL = {
    [OPENROUTER]: 'https://openrouter.ai/api/v1',
    [JINA]: 'https://api.jina.ai/v1/',
};

const DEFAULT_EMBEDDING_MODELS = {
    [OPENAI]: OPENAI_EMBED,
    [OPENROUTER]: GEMINI_EMBED,
};

const DEFAULT_RERANKER_MODELS = {
    [GOOGLE]: GOOGLE_MODEL_SEMANTIC_RANKER,
    [JINA]: JINA_MODEL_RERANKER_M0,
};

const MODEL_CONFIG = {
    [OPENAI_EMBED]: {
        source: 'openai', image: false, maxTokens: 8192,
    },
    [GEMINI_EMBED]: {
        source: 'google', image: true, maxTokens: 8192,
        options: { dimensions: 768, encoding_format: 'float' },
    },
    [GOOGLE_MODEL_SEMANTIC_RANKER]: {
        source: 'google', image: false, maxTokens: 1024, recordsLimit: 200,
        options: { ignoreRecordDetailsInResponse: true },
    },
    [JINA_MODEL_RERANKER_M0]: {
        source: 'jina', image: true, maxTokens: 1024, recordsLimit: 200,
        options: { return_documents: false },
    },
};

const ensureEmbeddingProvider = (options) => {
    options.provider = ensureString(options?.provider, { case: 'UP' });
    assert(
        DEFAULT_EMBEDDING_MODELS?.[options.provider],
        'Embedding provider is required.', 400
    );
    return options.provider;
};

const ensureRerankerProvider = (options) => {
    options.provider = ensureString(options?.provider, { case: 'UP' });
    assert(
        DEFAULT_RERANKER_MODELS?.[options.provider],
        'Reranker provider is required.', 400
    );
    return options.provider;
};

const ensureApiKey = (options) => {
    assert(options?.apiKey, 'API key is required.', 400);
    return options.apiKey;
};

const ensureCredentials = (options) => {
    assert(options?.credentials, 'Google credentials are required.', 400);
    assert(options?.projectId, 'Google project ID is required.', 400);
    return options;
};

const buildOpenAIClientOptions = (options = {}, baseURL) => {
    const clientOptions = { ...options, baseURL };
    for (const key of [
        'provider', 'model', 'modelConfig', 'source', 'debug',
    ]) { delete clientOptions[key]; }
    return clientOptions;
};

const getEmbeddingClient = (provider) => {
    provider = ensureString(provider, { case: 'UP' })
        || Object.keys(embeddingClients || {})[0];
    assert(provider, 'No embedding provider has been initialized.', 500);
    return { ...embeddingClients?.[provider], provider };
};

const getRerankerClient = (provider) => {
    provider = ensureString(provider, { case: 'UP' })
        || Object.keys(rerankerClients || {})[0];
    assert(provider, 'No reranker provider has been initialized.', 500);
    return { ...rerankerClients?.[provider], provider };
};

const initEmbedding = async (options = {}) => {
    if (options?.debug) {
        (await need('node:util')).inspect.defaultOptions.depth = null;
        options.logLevel = 'debug';
    }
    ensureApiKey(options);
    const provider = ensureEmbeddingProvider(options);
    const OpenAI = await need('openai');
    const baseURL = options?.baseURL || PROVIDER_BASE_URL[provider];
    const model = options?.model || DEFAULT_EMBEDDING_MODELS[provider];
    embeddingClients[provider] = {
        client: new OpenAI(buildOpenAIClientOptions(options, baseURL)),
        model, source: MODEL_CONFIG[model]?.source,
    };
    return getEmbeddingClient(provider);
};

const parseDataUrl = (value) => {
    const match = ensureString(value).match(/^data:([^;]+);base64,(.*)$/i);
    return match ? { mime: match[1], data: match[2], url: value } : null;
};

const normalizeEmbedImage = async (value, options) => {
    const dataUrl = parseDataUrl(value);
    if (dataUrl) { return dataUrl; }
    if (!options?.input && /^https?:\/\//i.test(ensureString(value))) {
        return { url: value, mime: options?.mime };
    }
    if (!options?.input && String.isString(value)) {
        assert(options?.mime, 'Image MIME type is required.', 400);
        return { data: value, mime: options.mime };
    }
    const { content, mime } = await convert(value, {
        ...options, expected: BASE64, meta: true,
    });
    assert(
        (options?.mime || mime)?.startsWith('image/'),
        `Invalid image MIME type: ${options?.mime || mime}.`, 400
    );
    return { data: content, mime: options?.mime || mime };
};

const formatOpenRouterInput = (items) => items.map((x) => {
    if (x.text !== undefined) {
        return { content: [{ type: 'text', text: x.text }] };
    }
    if (x.image) {
        return {
            content: [{
                type: 'image_url',
                image_url: {
                    url: x.image.url || formatDataURL(
                        x.image.mime, x.image.data
                    ),
                },
            }],
        };
    }
    assert(false, 'Unsupported embedding input type.', 400);
});

/**
 * Generate embeddings for text or image input.
 *
 * Usage:
 *   await initEmbedding({ provider: 'OPENROUTER', apiKey });
 *   await embed('hello world', { provider: 'OPENROUTER' });
 *   await embed({ image: '/path/to/image.jpg' }, {
 *       provider: 'OPENROUTER',
 *       input: 'FILE',
 *   });
 *   await embed([
 *       'hello world',
 *       { image: 'https://example.com/image.jpg' },
 *   ], { provider: 'OPENROUTER' });
 *
 * Input accepts a string, { text }, { image }, or an array of those values.
 * Local image files should pass { input: 'FILE' }. Raw base64 image strings
 * need { mime: 'image/...' }. OpenRouter image embeddings are sent through
 * its multimodal `image_url` embeddings format.
 */
const embed = async (input, options = {}) => {
    let [{ client, model: selectedModel, provider, source }, resp]
        = [getEmbeddingClient(options?.provider), null];
    const model = options?.model || selectedModel;
    const multiple = Array.isArray(input);
    input = await Promise.all(ensureArray(input).map(async x => {
        x = Object.isObject(x) ? x : { text: x };
        assert(
            Object.keys(x).length == 1,
            'Only one type of input is allowed at a time.', 400
        );
        if (x.text !== undefined) {
            x.text = trimText(x.text, MODEL_CONFIG[model]?.maxTokens);
        } else if (x.image !== undefined) {
            assert(
                MODEL_CONFIG[model]?.image,
                `Model ${model} does not support image embeddings.`, 400
            );
            x.image = await normalizeEmbedImage(x.image, options);
        } else {
            assert(false, 'Unsupported embedding input type.', 400);
        }
        return x;
    }));
    assert(input.length, 'Input is required.', 400);
    const body = {
        model,
        input: input.map(x => x.text),
        ...MODEL_CONFIG[model]?.options || {},
        ...options?.requestOptions || {},
    };
    switch (provider) {
        case OPENROUTER:
            source = options?.source || source
                || MODEL_CONFIG[body.model]?.source;
            body.model = `${source ? `${source}/` : ''}${body.model}`;
            input.some(x => x.image)
                && (body.input = formatOpenRouterInput(input));
        case OPENAI:
            resp = await client.embeddings.create(body);
            break;
        default:
            throw new Error(`Unsupported embedding provider: ${provider}`);
    }
    assert(resp?.data?.length, 'No embeddings returned.', 500);
    if (options?.raw) { return resp; }
    const vectors = resp.data.map(x => x.embedding);
    return multiple ? vectors : vectors[0];
};

const initReranker = async (options = {}) => {
    const provider = ensureRerankerProvider(options);
    const model = options?.model || DEFAULT_RERANKER_MODELS[provider];
    switch (provider) {
        case GOOGLE:
            ensureCredentials(options);
            const { RankServiceClient } = await need(
                '@google-cloud/discoveryengine', { raw: true }
            );
            const location = options?.location || GOOGLE_DEFAULT_LOCATION;
            const clientOptions = {
                ...location ? { apiEndpoint: `${location}-discoveryengine.googleapis.com` } : {},
                ...options?.apiEndpoint ? { apiEndpoint: options.apiEndpoint } : {},
                keyFilename: options.credentials,
            };
            const client = new RankServiceClient(clientOptions);
            rerankerClients[provider] = {
                client, model, rankingConfigPath: client.rankingConfigPath(
                    options.projectId, location,
                    options?.rerankerConfigId || GOOGLE_RERANK_CONFIG_ID
                ),
            };
            break;
        case JINA:
            const OpenAI = await need('openai');
            const baseURL = options?.baseURL || PROVIDER_BASE_URL[provider];
            rerankerClients[provider] = {
                client: new OpenAI(buildOpenAIClientOptions(options, baseURL)),
                model, source: MODEL_CONFIG[model]?.source,
            };
            break;
        default:
            throw new Error(`Unsupported reranker provider: ${provider}`);
    }
    return getRerankerClient(provider);
};

const rerank = async (query, records, options = {}) => {
    assert(query, 'Query is required.', 400);
    assert(records?.length, 'Records are required.', 400);
    const { provider, model, client, rankingConfigPath }
        = getRerankerClient(options?.provider);
    records = records.map((content, id) => Object.isObject(content)
        ? content : { id: String(id), content }).slice(
            0, MODEL_CONFIG[model]?.recordsLimit || records.length
        );
    const maxTokens = MODEL_CONFIG[model]?.maxTokens || Infinity;
    let result;
    for (let i in records) {
        records[i].title = await trimText(records[i]?.title || '', maxTokens);
        const titleTokens = await countTokens(records[i].title);
        const availableTokens = maxTokens - titleTokens;
        records[i].content = availableTokens > 0 ? await trimText(
            records[i].content, availableTokens
        ) : '';
        records[i].image = records[i].image ? await convert(records[i].image, {
            ...options, expected: BASE64,
        }) : undefined;
    }
    switch (provider) {
        case GOOGLE:
            var body = {
                model, query, rankingConfig: rankingConfigPath,
                records, topN: ~~options?.topN || records.length,
                ...MODEL_CONFIG[model]?.options || {},
                ...options?.requestOptions || {},
            };
            result = (await client.rank(body))?.[0]?.records;
            options?.raw || (result = result.map(x => ({
                index: ~~x.id, score: x.score,
            })));
            break;
        case JINA:
            records = records.map(x =>
                ((x.title || x.content) ? { text: [x.title, x.content].filter(x => x).join('\n') } : null)
                || (x.image ? { image: x.image } : null)
            ).filter(x => x);
            assert(records.length, 'No valid records found.', 400);
            var body = {
                model, query, documents: records,
                ...MODEL_CONFIG[model]?.options || {},
                ...options?.requestOptions || {},
            };
            result = (await client.post('/rerank', { body }))?.results;
            options?.raw || (result = result.map(x => ({
                index: ~~(x.index), score: x.relevance_score,
            })));
            break;
        default:
            throw new Error(`Unsupported reranker provider: ${provider}`);
    }
    result.sort((a, b) => b.score - a.score);
    return result || [];
};

export {
    _NEED,
    embed,
    initEmbedding,
    initReranker,
    rerank,
};