import { createCohere } from '@ai-sdk/cohere';
import { GatewayFetchError } from '~~/server/utils/ai-provider';
import { Err, Ok } from '~~/types/result';
import type { ProviderModule } from '.';
import { getModelData, mergeSets } from './helpers';

export default {
    id: 'cohere',
    baseUrl: 'https://api.cohere.ai/v2',
    modelsEndpoint: '/models',

    createGateway({ apiKey, baseURL }) {
        if (apiKey === undefined) {
            return Err(GatewayFetchError.NoProviderApiKey);
        }

        return Ok({
            gateway: createCohere({ apiKey, baseURL }),
            streamTransformer: undefined,
            textTransformer: undefined,
        });
    },

    getAuthHeaders(apiKey) {
        return apiKey ? { Authorization: `Bearer ${apiKey}` } : {};
    },

    async fetchModels({ baseURL, apiKey, modelsDevData }) {
        const res = await fetch(`${baseURL}/models`, {
            method: 'GET',
            headers: { Authorization: `Bearer ${apiKey}` },
        });
        const data = await res.json();

        if (!res.ok) {
            throw createError({ statusCode: res.status, message: JSON.stringify(data) });
        }

        const models = [];
        for (const model of data.models) {
            const inputModalities = new Set<string>();
            const capabilities = new Set<string>();

            for (const feature of model.features || []) {
                switch (feature) {
                    case 'tools':
                        capabilities.add('tools');
                        break;
                    case 'vision':
                        inputModalities.add('image');
                        break;
                    case 'reasoning':
                        capabilities.add('reasoning');
                        break;
                }
            }

            const modelData = getModelData(model.name, 'cohere', modelsDevData);

            models.push({
                ...modelData,
                id: model.name,
                attributes: {
                    ...modelData.attributes,
                    inputModalities: mergeSets(
                        inputModalities,
                        new Set(modelData.attributes?.inputModalities || [])
                    ),
                    capabilities: mergeSets(
                        capabilities,
                        new Set(modelData.attributes?.capabilities || [])
                    ),
                    contextWindow: model.context_length || modelData.attributes?.contextWindow,
                }
            });
        }

        return models;
    },
} satisfies ProviderModule;