import * as z from 'zod';
import { providerBaseUrls, SupportedModalities } from '~/types/model';
import { Providers } from '~/types/model';
import { httpClient } from '~~/server/lib/triplit';

export default defineEventHandler(async (event) => {
    await protectRoute(event);
    const userId = event.context.user!.id;

    const result = await readValidatedBody(event, (body) =>
        z
            .object({
                providerApiKey: z.string().optional(),
            })
            .safeParse(body),
    );

    if (!result.success) {
        throw createError({
            statusCode: 400,
            message: result.error.issues[0]!.message,
        });
    }

    const providerId = getRouterParam(event, 'providerId')

    const provider = await httpClient.fetchOne(httpClient.query('providers').Where('id', '=', providerId));
    if (provider === null || provider.userId !== userId) {
        throw createError({
            statusCode: 400,
            message: 'Invalid provider',
        });
    }

    let baseUrl;
    let fetchUrl;
    let headers;
    switch (provider.type) {
        case 'cohere':
        case 'cerebras':
            if (!result.data.providerApiKey) {
                throw createError({
                    statusCode: 400,
                    message: `${provider.type} provider requires an API key`,
                });
            }
        case 'google':
        case 'openrouter':
            baseUrl = !!provider.config.apiProxyUrl ? provider.config.apiProxyUrl : providerBaseUrls[provider.type];
            baseUrl = baseUrl.replace(/\/$/, '');

            if (baseUrl === '') {
                throw createError({
                    statusCode: 400,
                    message: 'Invalid provider URL',
                });
            }

            fetchUrl = `${baseUrl}/models`;
            break;
        case 'ollama':
            if (!provider.config.apiProxyUrl) {
                throw createError({
                    statusCode: 400,
                    message: 'Ollama provider requires an API proxy URL',
                });
            }
            baseUrl = provider.config.apiProxyUrl;
            baseUrl = baseUrl.replace(/\/$/, '');

            fetchUrl = `${baseUrl}/api/tags`;
            break;
        case 'longcat':
            // longcat doesn't have a model list endpoint so we just hardcode them here,
            // sry. I talked to Meituan and this is what they said:
            // 后续如果我们新增了这样的接口会及时同步您。
            // en (approx): If we add an interface like this in the future, we will promptly update you
            return {
                models: [
                    {
                        id: "LongCat-Flash-Chat",
                        name: "LongCat Flash Chat",
                        attributes: {
                            inputModalities: ['text'],
                            outputModalities: ['text'],
                            capabilities: ['tools'],
                            contextWindow: 256_000,
                        }
                    },
                    {
                        id: "LongCat-Flash-Thinking",
                        name: "LongCat Flash Thinking",
                        attributes: {
                            inputModalities: ['text'],
                            outputModalities: ['text'],
                            capabilities: ['reasoning', 'tools'],
                            contextWindow: 256_000,
                        }
                    },
                    {
                        id: "LongCat-Flash-Thinking-2601",
                        name: "LongCat Flash Thinking (2601)",
                        attributes: {
                            inputModalities: ['text'],
                            outputModalities: ['text'],
                            capabilities: ['reasoning', 'tools'],
                            contextWindow: 256_000,
                        }
                    },
                    {
                        id: "LongCat-Flash-Lite",
                        name: "LongCat Flash Lite",
                        attributes: {
                            inputModalities: ['text'],
                            outputModalities: ['text'],
                            capabilities: ['tools'],
                            contextWindow: 320_000,
                        }
                    }
                ]
            };
    }

    if (provider.type === 'google') {
        headers = {
            'x-goog-api-key': `${result.data.providerApiKey}`
        }
    } else {
        headers = {
            'Authorization': `Bearer ${result.data.providerApiKey}`
        }
    }

    let res;
    let data;
    try {
        res = await fetch(fetchUrl, {
            method: 'GET',
            headers
        });
        data = await res.json();
    } catch (e) {
        console.error('Failed to fetch models:', e);
        throw createError({
            statusCode: 500,
            message: 'Failed to fetch models ' + e,
        });
    }

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

    return {
        models: await normalizeResponse(data, provider.type, baseUrl)
    };
});

const normalizeResponse = async (response: Record<string, any>, provider: typeof Providers[number], baseUrl: string) => {
    switch (provider) {
        case 'cerebras': {
            console.log(response);
            return response.data.map((model: any) => ({ id: model.id, releasedAt: model.created }));
        }
        case 'openrouter': {
            const models = [];

            for (const model of response.data) {
                const capabilities = new Set<string>();

                for (const capability of model.supported_parameters) {
                    switch (capability) {
                        case 'reasoning': {
                            capabilities.add('reasoning');
                        } break;
                        case 'tools': {
                            capabilities.add('tools');
                        } break;
                    }
                }

                models.push({
                    id: model.id as string,
                    name: model.name as string,
                    pricing: model.pricing,
                    attributes: {
                        inputModalities: model.architecture.input_modalities.filter((modality: string) => (SupportedModalities as Readonly<string[]>).includes(modality)),
                        outputModalities: model.architecture.output_modalities.filter((modality: string) => (SupportedModalities as Readonly<string[]>).includes(modality)),
                        capabilities: Array.from(capabilities),
                        contextWindow: model.context_length,
                        supported_parameters: model.supported_parameters,
                    },
                    created: model.created,
                });
            }

            return models;
        }
        case 'ollama': {
            const headers: Record<string, string> = {
                'Content-Type': 'application/json',
                'User-Agent': 'Mozilla/5.0'
            }

            if (response.token) {
                headers['Authorization'] = `Bearer ${response.token}`;
            }

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

            for (const model of response.models) {
                infoPromises.push(fetch(`${baseUrl}/api/show`, {
                    method: 'POST',
                    headers,
                    body: JSON.stringify({
                        model: model.name,
                    }),
                }).then((res) => {
                    return 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);
                    }

                    models.set(model.name, {
                        id: model.name,
                        name: model.name,
                        attributes: {
                            inputModalities: Array.from(inputModalities),
                            capabilities: Array.from(capabilities),
                            contextWindow,
                        }
                    });
                }));
            }

            await Promise.all(infoPromises);

            return Array.from(models.values());
        }
        case 'google': {
            console.log(response);
            return response.models.map((model: any) => ({ id: model.name.replace('models/', ''), name: model.displayName }));
        }
        case 'longcat': {
            console.log(response);
            return response.models.map((model: any) => ({ id: model.name, name: model.name }));
        }
        case 'cohere': {
            const models = [];
            for (const model of response.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;
                    }
                }

                models.push({
                    id: model.name,
                    name: model.name,
                    attributes: {
                        contextWindow: model.context_length,
                        inputModalities: Array.from(inputModalities),
                        capabilities: Array.from(capabilities),
                    }
                });
            }

            return models;
        }
    }
}