import { createOllama } from 'ai-sdk-ollama';
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: 'ollama',
    baseUrl: '',
    modelsEndpoint: '/api/tags',

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

        if (model !== undefined) {
            const innerGateway = createOllama({
                apiKey,
                baseURL,
            });

            return Ok({
                gateway: ((modelId: string) => innerGateway(modelId, { think: model.capabilities.includes('reasoning') })) as any,
                streamTransformer: undefined,
                textTransformer: undefined,
            });
        }

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

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

    async fetchModels({ baseURL, apiKey, modelsDevData }) {
        const headers: Record<string, string> = {
            'Content-Type': 'application/json',
            'User-Agent': 'Mozilla/5.0'
        };

        if (apiKey) {
            headers['Authorization'] = `Bearer ${apiKey}`;
        }

        const res = await fetch(`${baseURL}/api/tags`, {
            method: 'GET',
            headers,
        });
        const data = await res.json();

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

        const models = new Map<string, any>();
        const infoPromises = [];

        for (const model of data.models) {
            infoPromises.push(fetch(`${baseURL}/api/show`, {
                method: 'POST',
                headers,
                body: JSON.stringify({ model: model.name }),
            }).then((res) => res.json()).then((json) => {
                const inputModalities = new Set<string>();
                const capabilities = new Set<string>();

                for (const capability of json.capabilities) {
                    switch (capability) {
                        case 'thinking':
                            capabilities.add('reasoning');
                            break;
                        case 'tools':
                            capabilities.add('tools');
                            break;
                        case 'completion':
                            inputModalities.add('text');
                            break;
                        case 'vision':
                            inputModalities.add('image');
                            break;
                    }
                }

                let contextWindow: number | undefined;
                try {
                    for (const key of Object.keys(json.model_info)) {
                        if (key.endsWith('context_length')) {
                            contextWindow = json.model_info[key];
                            break;
                        }
                    }
                } catch (e) {
                    console.error(e);
                }

                let normalizedId = model.name;
                normalizedId = normalizedId.replace(/:cloud$/, '');
                normalizedId = normalizedId.replace(/-cloud$/, '');
                normalizedId = normalizedId.replace(/:latest$/, '');
                console.log(normalizedId);
                const modelData = getModelData(normalizedId, 'ollama-cloud', modelsDevData);

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

        await Promise.all(infoPromises);
        return Array.from(models.values());
    },
} satisfies ProviderModule;