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

export default {
    id: 'huggingface',
    baseUrl: 'https://router.huggingface.co/v1',
    modelsEndpoint: null,

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

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

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

    async fetchModels({ baseURL, apiKey, modelsDevData }) {
        const res = await fetch(
            'https://huggingface.co/api/models?pipeline_tag=text-generation&sort=downloads',
            {
                method: 'GET',
                headers: apiKey ? { Authorization: `Bearer ${apiKey}` } : {},
            },
        );
        const data = await res.json();

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

        return data
            .filter((model: any) => model.tags?.includes('text-generation-inference'))
            .map((model: any) => ({
                ...getModelData(model.id, 'huggingface', modelsDevData),
                id: model.id,
                name: model.id,
                releasedAt: model.createdAt,
            }));
    },
} satisfies ProviderModule;