import * as z from 'zod';
import { SupportedModalities } from '~/types/model';
import { Providers } from '~/types/model';
import { GatewayFetchError, getProviderDetails } from '~~/server/utils/ai-provider';
import { getModelsDevData } from '~~/server/utils/models-dev';
import { Err, Ok, type Result } from '~~/types/result';
import { db } from '~~/server/lib/db';
import { type Provider } from '~/composables/useModels';
import { models } from '~~/drizzle/schema';
import { and, eq, notInArray } from 'drizzle-orm';

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

    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');
    if (!providerId) throw createError({ statusCode: 400, message: 'Invalid provider' });

    const provider = await db.query.providers.findFirst({
        where: {
            id: providerId,
            userId,
        },
        with: { models: true }
    });

    if (!provider) throw createError({ statusCode: 404, message: 'Provider not found' });

    const [providerModelsRes, modelsDevRes] = await Promise.all([
        fetchProviderModels(provider, result.data.providerApiKey),
        getModelsDevData(),
    ]);

    if (!providerModelsRes.ok) {
        switch (providerModelsRes.error) {
            case ProviderFetchError.NoProviderApiKey: {
                throw createError({
                    statusCode: 400,
                    message: `${provider.type} provider requires an API key`,
                });
            }
            case ProviderFetchError.NoProviderBaseUrl: {
                throw createError({
                    statusCode: 400,
                    message: 'Invalid provider URL',
                });
            }
        }
    }

    const normalizedModels = await normalizeResponse(
        providerModelsRes.data.data,
        provider.type as typeof Providers[number],
        providerModelsRes.data.baseURL,
        modelsDevRes
    );

    const existingModels = provider.models || [];
    const apiModelExternalIds = normalizedModels.map((m: any) => m.id);

    console.log({ existingModels, apiModelExternalIds, normalizedModels });

    await db.transaction(async (tx) => {
        if (existingModels.length > 0) {
            await tx.delete(models)
                .where(
                    and(
                        eq(models.providerId, providerId),
                        eq(models.isCustom, false),
                        notInArray(models.externalId, apiModelExternalIds)
                    )
                );
        }

        for (const model of normalizedModels) {
            const existing = existingModels.find(m => m.externalId === model.id);

            if (existing) {
                await tx.update(models)
                    .set({
                        name: model.name || existing.name,
                        cost: model.cost,
                        inputModalities: model.attributes.inputModalities,
                        outputModalities: model.attributes.outputModalities,
                        capabilities: model.attributes.capabilities,
                        contextWindow: model.attributes.contextWindow,
                        releasedAt: model.releasedAt ? new Date(model.releasedAt) : existing.releasedAt,
                    })
                    .where(eq(models.id, existing.id));
            } else {
                await tx.insert(models).values({
                    userId,
                    providerId,
                    externalId: model.id,
                    name: model.name || model.id,
                    cost: model.cost || {},
                    inputModalities: model.attributes.inputModalities,
                    outputModalities: model.attributes.outputModalities,
                    capabilities: model.attributes.capabilities,
                    contextWindow: model.attributes.contextWindow,
                    isCustom: false,
                    enabled: false,
                    releasedAt: model.releasedAt ? new Date(model.releasedAt) : null,
                });
            }
        }
    });

    const updatedModels = await db.query.models.findMany({
        where: {
            providerId,
        },
        orderBy: {
            releasedAt: 'desc'
        }
    });

    return {
        models: updatedModels
    };
});

enum ProviderFetchError {
    NoProviderApiKey = 0,
    NoProviderBaseUrl,
}

const fetchProviderModels = async (provider: Provider, providerApiKey: string | undefined): Promise<Result<{ data: Record<string, any>, baseURL: string }, ProviderFetchError>> => {
    const providerDetails = await getProviderDetails(provider, providerApiKey);
    if (!providerDetails.ok) {
        switch (providerDetails.error) {
            case GatewayFetchError.NoProviderApiKey: {
                return Err(ProviderFetchError.NoProviderApiKey);
            }
            case GatewayFetchError.NoProviderBaseUrl: {
                return Err(ProviderFetchError.NoProviderBaseUrl);
            }
        }
    }

    const { endpoint: { baseURL, modelsEndpoint, headers } } = providerDetails.data;

    switch (provider.type) {
        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 keep you in sync
            return Ok({
                data: {
                    models: [
                        {
                            id: "LongCat-Flash-Chat",
                            name: "LongCat Flash Chat",
                            attributes: {
                                inputModalities: ['text'],
                                outputModalities: ['text'],
                                capabilities: ['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-Omni-2603",
                            name: "LongCat Flash Omni",
                            attributes: {
                                inputModalities: ['text', 'image', 'audio', 'video'],
                                outputModalities: ['text', 'audio'],
                                capabilities: ['tools'],
                                contextWindow: 256_000,
                            }
                        },
                        {
                            id: "LongCat-Flash-Lite",
                            name: "LongCat Flash Lite",
                            attributes: {
                                inputModalities: ['text'],
                                outputModalities: ['text'],
                                capabilities: ['tools'],
                                contextWindow: 320_000,
                            }
                        }
                    ],
                },
                baseURL
            });
        case 'inception':
            return Ok({
                data: {
                    models: [
                        {
                            id: "mercury-2",
                            name: "Mercury 2",
                            attributes: {
                                inputModalities: ['text'],
                                outputModalities: ['text'],
                                capabilities: ['tools', 'reasoning'],
                                contextWindow: 128_000,
                            }
                        },
                    ],
                },
                baseURL
            });
    }

    let res;
    let data;
    try {
        res = await fetch(`${baseURL}${modelsEndpoint}`, {
            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 Ok({ data, baseURL });
}

function mergeSets(setA: Set<string>, setB: Set<string>): string[] {
    return Array.from(new Set([...setA, ...setB]));
}

const getModelData = (modelId: string, providerId: string, modelsDevData: any) => {
    const modelData = modelsDevData[providerId]?.models[modelId];

    if (modelData === undefined) return {
        cost: {},
        attributes: {
            inputModalities: ['text'],
            outputModalities: ['text'],
            capabilities: [],
        }
    };

    const capabilities = new Array<string>();
    if (modelData.reasoning) {
        capabilities.push('reasoning');
    }

    if (modelData.tool_call) {
        capabilities.push('tools');
    }

    let inputModalities = modelData.modalities.input.filter((m: string) => (SupportedModalities as Readonly<string[]>).includes(m));
    let outputModalities = modelData.modalities.output.filter((m: string) => (SupportedModalities as Readonly<string[]>).includes(m));
    let supportedParameters = [];
    if (modelData.temperature) {
        supportedParameters.push('temperature');
    }

    let contextWindow = modelData.limit?.context || null;
    let cost: Record<string, string> = {};
    for (const key in modelData.cost) {
        switch (key) {
            case 'input': {
                cost.prompt = formatMoney(modelData.cost[key].toString());
            } break;
            case 'output': {
                cost.completion = formatMoney(modelData.cost[key].toString());
            } break;
        }
    }

    let releasedAt = (new Date(modelData.release_date)).getTime();

    return {
        name: modelData.name,
        attributes: {
            inputModalities,
            outputModalities,
            capabilities: capabilities || [],
            contextWindow,
            supported_parameters: modelData.supportedParameters,
        },
        cost,
        releasedAt,
    }
}

// this function multiplies a string in the format of 'D.DD' by 1_000_000
// it does this by finding the first digit that is not a zero, and then
// left shifting it in decimal by 3 places
const lshDecimal = (number: string, shift: number) => {
    if (number.length === 0) {
        return '0.00';
    }

    let value = '';
    let isNegative = number[0] === '-';
    let [integerPart, fractionalPart] = number.substring(isNegative ? 1 : 0).split('.');

    if (!fractionalPart) {
        fractionalPart = '';
    }

    if (shift <= fractionalPart.length) {
        integerPart += fractionalPart.substring(0, shift);
        fractionalPart = fractionalPart.substring(shift);
    } else if (shift > fractionalPart.length) {
        integerPart += fractionalPart;
        for (let i = 0; i < shift - fractionalPart.length; i++) {
            integerPart += '0';
        }
        fractionalPart = '';
    }

    if (isNegative) {
        value += '-';
    }

    integerPart = integerPart!.replace(/^0+/, '');
    if (integerPart.length === 0) {
        integerPart = '0';
    }

    value += integerPart;

    if (fractionalPart.length > 0) {
        if (fractionalPart.length < 2) {
            fractionalPart += '0';
        }

        value += '.' + fractionalPart;
    } else {
        value += '.00';
    }

    return value;
}

const formatMoney = (value: string) => {
    // ensure that there are two decimal places MINIMUM I WILL STRANGLE YOU SO HELP ME GOD
    const [integerPart, fractionalPart] = value.split('.');

    if (fractionalPart === undefined) {
        return `${integerPart}.00`;
    }

    if (fractionalPart.length === 1) {
        return `${integerPart}.${fractionalPart}0`;
    }

    return `${integerPart}.${fractionalPart}`;
}

const normalizeResponse = async (response: Record<string, any>, provider: typeof Providers[number], baseUrl: string, modelsDevData: any): Promise<any[]> => {
    switch (provider) {
        case 'cerebras': {
            console.log("response.data", response.data);
            for (const model of response.data) {
                console.log(model.created);
            }

            return response.data.map((model: any) => ({ ...getModelData(model.id, provider, modelsDevData), id: model.id }));
        }
        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;
                    }
                }

                const pricing = {} as Record<string, string>;
                for (const key of Object.keys(model.pricing)) {
                    // normalize to /1M tokens since models.dev is already in /1M tokens

                    switch (key) {
                        case 'prompt':
                        case 'completion':
                        case 'image':
                        case 'audio':
                        case 'discount':
                            pricing[key] = lshDecimal(model.pricing[key], 6);
                            break;
                        case 'request':
                            pricing['request'] = lshDecimal(model.pricing[key], 3);
                            break;
                        case 'image_tokens':
                            pricing['imageTokens'] = lshDecimal(model.pricing[key], 6);
                            break;
                        case 'image_output':
                            pricing['imageOutput'] = lshDecimal(model.pricing[key], 6);
                            break;
                        case 'audio_output':
                            pricing['audioOutput'] = lshDecimal(model.pricing[key], 6);
                            break;
                        case 'input_audio_cache':
                            pricing['inputAudioCache'] = lshDecimal(model.pricing[key], 6);
                            break;
                        case 'web_search':
                            pricing['webSearch'] = lshDecimal(model.pricing[key], 3);
                            break;
                        case 'internal_reasoning':
                            pricing['internalReasoning'] = lshDecimal(model.pricing[key], 6);
                            break;
                        case 'input_cache_read':
                            pricing['inputCacheRead'] = lshDecimal(model.pricing[key], 6);
                            break;
                        case 'input_cache_write':
                            pricing['inputCacheWrite'] = lshDecimal(model.pricing[key], 6);
                            break;
                    }
                }

                // TODO: do I really need to pull in data from models.dev for OR models?
                // const modelData = getModelData(model.id, provider, modelsDevData);

                models.push({
                    id: model.id as string,
                    name: model.name as string,
                    cost: 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,
                    },
                    releasedAt: model.created * 1000,
                });
            }

            return models;
        }
        case 'closedrouter': {
            return response.data.map((model: any) => {
                const capabilities = new Array<string>();

                if (model.capabilities) {
                    if (model.capabilities.reasoning) {
                        capabilities.push('reasoning');
                    }

                    if (model.capabilities.tool_call) {
                        capabilities.push('tools');
                    }
                }

                return {
                    id: model.id,
                    name: model.name,
                    attributes: {
                        inputModalities: model.modalities?.input || ['text'],
                        outputModalities: model.modalities?.output || ['text'],
                        capabilities: capabilities,
                        contextWindow: model.context_window,
                        supported_parameters: [],
                    },
                }
            });
        }
        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);
                    }

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

                    models.set(model.name, {
                        ...modelData,
                        id: model.name,
                        name: 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());
        }
        case 'vllm': {
            return response.data.map((model: any) => {
                return {
                    id: model.id,
                    attributes: {
                        inputModalities: ['text'],
                        outputModalities: ['text'],
                        contextWindow: model.max_model_len,
                    }
                }
            })
        }
        case 'google': {
            return response.models.map((model: any) => ({ ...getModelData(model.name.replace('models/', ''), provider, modelsDevData), id: model.name.replace('models/', ''), name: model.displayName }));
        }
        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;
                    }
                }

                const modelData = getModelData(model.name, provider, 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;
        }
        case 'mistral': {
            const models = [];
            for (const model of response.data) {
                const inputModalities = new Set<string>();
                const capabilities = new Set<string>();

                for (const capability of Object.keys(model.capabilities)) {
                    if (model.capabilities[capability] === false) continue;

                    switch (capability) {
                        case 'function_calling': {
                            capabilities.add('tools');
                        } break;
                        case 'completion_chat': {
                            inputModalities.add('text');
                            capabilities.add('completion');
                        } break;
                        case 'vision': {
                            inputModalities.add('image');
                        } break;
                    }
                }

                const modelData = getModelData(model.id, provider, modelsDevData);

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

            return models;
        }
        case 'inception':
        case 'longcat': {
            return response.models;
        }
    }
}