import * as z from 'zod';
import { getProvider } from '~~/server/utils/providers';
import { getModelsDevData } from '~~/server/utils/models-dev';
import { db } from '~~/server/lib/db';
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 providerModule = getProvider(provider.type);
    if (!providerModule) {
        throw createError({ statusCode: 400, message: `Unknown provider type: ${provider.type}` });
    }

    let baseURL = undefined;
    if (provider.config.apiProxyUrl && provider.config.apiProxyUrl.trim() !== '') {
        baseURL = provider.config.apiProxyUrl;
    } else {
        baseURL = providerModule.baseUrl;
    }
    baseURL = baseURL.replace(/\/$/, '');

    const normalizedModels = await providerModule.fetchModels({
        baseURL,
        apiKey: result.data.providerApiKey,
        modelsDevData: await getModelsDevData(),
    })
    const existingModels = provider.models || [];
    const apiModelExternalIds = normalizedModels.map((m: any) => m.id);

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