import { type Entity } from '@triplit/client';
import * as z from 'zod';
import { SupportedModalities } from '~/types/model';
import { Providers } from '~/types/model';
import { httpClient } from '~~/server/lib/triplit';
import { GatewayFetchError, getProviderDetails } from '~~/server/utils/ai-provider';
import { getModelsDevData } from '~~/server/utils/models-dev';
import { schema } from '~~/triplit/schema';
import { Err, Ok, type Result } from '~~/types/result';

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',
        });
    }

    console.log("apiKey", result.data.providerApiKey);

    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',
                });
            }
        }
    }

    console.log("providerModelsRes.data", providerModelsRes.data);

    const normalizedModels = await normalizeResponse(
        providerModelsRes.data.data,
        provider.type,
        providerModelsRes.data.baseURL,
        modelsDevRes
    );

    return {
        models: normalizedModels,
    };
});

enum ProviderFetchError {
    NoProviderApiKey = 0,
    NoProviderBaseUrl,
}

const fetchProviderModels = async (provider: Entity<typeof schema, 'providers'>, 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: {
                // throw createError({
                //     statusCode: 400,
                //     message: `${provider.type} provider requires an API key`,
                // });
                return Err(ProviderFetchError.NoProviderApiKey);
            }
            case GatewayFetchError.NoProviderBaseUrl: {
                // throw createError({
                //     statusCode: 400,
                //     message: 'Invalid provider URL',
                // });
                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 promptly update you
            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",
                            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,
                            }
                        }
                    ],
                },
                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];
    console.log("modelData", modelData);

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

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

    if (modelData.tool_call) {
        capabilities.add('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: modelData.capabilities || [],
            contextWindow,
            supported_parameters: modelData.supportedParameters,
        },
        cost,
        releasedAt,
    }
}

// const formatBig = (bigValue: Big) => {
//     let str = bigValue.toString();

//     if (!str.includes('.')) return str + '.00';
//     if (str.split('.')[1]!.length === 1) return str + '0';

//     return str;
// };

// 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) => {
    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: mergeSets(
                        //     new Set(modelData.inputModalities || []),
                        //     new Set(model.architecture.input_modalities.filter((modality: string) => (SupportedModalities as Readonly<string[]>).includes(modality))),
                        // ),
                        // outputModalities: mergeSets(
                        //     new Set(modelData.outputModalities || []),
                        //     new Set(model.architecture.output_modalities.filter((modality: string) => (SupportedModalities as Readonly<string[]>).includes(modality))),
                        // ),
                        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 '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 'google': {
            return response.models.map((model: any) => ({ ...getModelData(model.name.replace('models/', ''), provider, modelsDevData), id: model.name.replace('models/', ''), name: model.displayName }));
        }
        case 'longcat': {
            return response.models;
        }
        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;
        }
    }
}