import { createOpenRouter, type OpenRouterProvider } from '@openrouter/ai-sdk-provider';
import type { Entity } from '@triplit/client';
import { type ModelMessage, modelMessageSchema, streamText, tool } from 'ai';
import { promises as fs } from 'fs';
import { glob } from 'glob';
import { nanoid } from 'nanoid';
import path from 'path';
import * as z from 'zod';
import { httpClient } from '~~/server/lib/triplit';
import { addPendingGeneration, completeGeneration } from '~~/server/utils/generations';
import type { schema } from '~~/triplit/schema';

export const messagesSchema = z.array(modelMessageSchema);

// quick access
// const ACTIVE_MODEL_ID = 'openrouter/free';
// const ACTIVE_MODEL_ID = 'x-ai/grok-4.1-fast';
const ACTIVE_MODEL_ID = 'google/gemini-3-flash-preview';
// const ACTIVE_MODEL_ID = 'arcee-ai/trinity-mini:free';

type ModelGateway = OpenRouterProvider;

export default defineEventHandler(async (event) => {
    await protectRoute(event);

    const result = await readValidatedBody(event, (body) =>
        z
            .object({
                messages: messagesSchema.min(1),
                topicId: z.string(),
                model: z.object({
                    providerId: z.string(),
                    modelId: z.string(),
                    args: z.any(),
                }),
                providerApiKey: z.string().optional(),
            })
            .safeParse(body),
    );

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

    const userId = event.context.user!.id;

    const { messages, topicId, model: { modelId, providerId }, providerApiKey } = result.data;

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

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

    const existingPendingGenerations = await httpClient.fetchOne(
        httpClient.query('generations').Where('topicId', '=', topicId).Where('status', '=', 'pending'),
    );
    if (existingPendingGenerations !== null) {
        throw createError({
            statusCode: 400,
            message: 'There cannot be more than one active generation per topic',
        });
    }

    let gateway: ModelGateway;

    switch (provider.type) {
        case 'openrouter': {
            if (providerApiKey === undefined) {
                throw createError({
                    statusCode: 400,
                    message: 'OpenRouter provider requires an API key',
                });
            }

            gateway = createOpenRouter({
                apiKey: providerApiKey,
                headers: {
                    'HTTP-Referer': 'https://localhost:3000',
                    'X-Title': 'Veridian',
                },
            });
            break;
        }
        default:
            throw new Error(`Unknown provider type: ${provider.type}`);
    }

    const generationId = nanoid();
    const message = await httpClient.insert('messages', {
        topicId,
        generationId,
        content: '',
        role: 'assistant',
    });
    await httpClient.insert('generations', {
        id: generationId,
        topicId,
        modelId: model.externalId,
        status: 'pending',
        messageId: message.id,
    });

    let logFile: fs.FileHandle | undefined;
    let logMessage: ((message: string) => void) | undefined;

    if (process.env.GENERATION_DEBUG) {
        logFile = await fs.open(path.join(process.env.LOG_DIR!, `${Date.now()}-${generationId}.log`), 'w');
        logMessage = (message: string) => {
            logFile!.write(message + '\n');
        };
    }

    event.waitUntil(
        generateResponse(message, { gateway, model: model.externalId }, generationId, userId, messages, logMessage, logFile),
    );

    return {
        generationId,
        messageId: message.id,
    };
});

const INTERNAL_ERROR = 'An internal error occurred';

// todo message takes in variadics like console.log
const todo = (...args: any[]) => {
    console.error('TODO', ...args);
    throw new Error('TODO');
};

async function generateResponse(
    message: Entity<typeof schema, 'messages'>,
    model: {
        gateway: ModelGateway,
        model: string,
    },
    generationId: string,
    userId: string,
    messages: ModelMessage[],
    log?: (message: string) => void,
    logFile?: fs.FileHandle,
) {
    const controller = new AbortController();
    addPendingGeneration(generationId, controller);

    const activeParts = new Map<string, { id: string; accumulatedContent: string; providerOptions?: any }>();
    const activeToolCalls = new Map<string, { id: string }>();

    const response = streamText({
        model: model.gateway(model.model),
        messages,
        // a little trick that makes it so that the stream doesnt stop because of tool calls, and will continue an unbounded amount of time and steps
        stopWhen: [],
        // tools: {
        //     // writeFile: tool({
        //     //     inputSchema: z.object({
        //     //         path: z.string(),
        //     //         content: z.string(),
        //     //     }),
        //     //     outputSchema: z.object({
        //     //         success: z.boolean(),
        //     //     }),
        //     //     execute: async ({ path, content }) => {
        //     //         await fs.writeFile(path, content);
        //     //         return {
        //     //             success: true,
        //     //         };
        //     //     }
        //     // }),
        //     listDirectory: tool({
        //         inputSchema: z.object({
        //             path: z.string(),
        //         }),
        //         outputSchema: z.object({
        //             files: z.array(z.object({ name: z.string(), type: z.string() })),
        //         }),
        //         execute: async ({ path }) => {
        //             const rawFiles = await fs.readdir(path, { withFileTypes: true });
        //             const files = rawFiles.map((file) => ({
        //                 name: file.name,
        //                 type: file.isFile() ? 'file' : 'directory',
        //             }));

        //             return {
        //                 files,
        //             };
        //         },
        //     }),
        //     glob: tool({
        //         inputSchema: z.object({
        //             pattern: z.string(),
        //         }),
        //         outputSchema: z.object({
        //             files: z.array(z.object({ name: z.string(), type: z.string() })),
        //         }),
        //         execute: async ({ pattern }) => {
        //             const rawFiles = await glob(pattern, { withFileTypes: true });
        //             const files = rawFiles.map((file) => ({
        //                 name: file.name,
        //                 type: file.isFile() ? 'file' : 'directory',
        //             }));

        //             return {
        //                 files,
        //             };
        //         },
        //     }),
        //     readFile: tool({
        //         inputSchema: z.object({
        //             path: z.string(),
        //         }),
        //         outputSchema: z.object({
        //             path: z.string(),
        //             content: z.string(),
        //         }),
        //         execute: async ({ path }) => {
        //             const file = await fs.readFile(path);
        //             return {
        //                 path,
        //                 content: file.toString(),
        //             };
        //         },
        //     }),
        //     readFiles: tool({
        //         inputSchema: z.object({
        //             paths: z.array(z.string()).describe('The file paths to read'),
        //         }),
        //         outputSchema: z.object({
        //             files: z.array(
        //                 z.object({
        //                     path: z.string(),
        //                     content: z.string(),
        //                 }),
        //             ),
        //         }),
        //         execute: async ({ paths }) => {
        //             const files = await Promise.all(
        //                 paths.map(async (path) => {
        //                     const file = await fs.readFile(path);
        //                     return {
        //                         path: path,
        //                         content: file.toString(),
        //                     };
        //                 }),
        //             );

        //             return {
        //                 files,
        //             };
        //         },
        //     }),
        //     fetchUrl: tool({
        //         inputSchema: z.object({
        //             url: z.string(),
        //         }),
        //         outputSchema: z.object({
        //             content: z.string(),
        //         }),
        //         execute: async ({ url }) => {
        //             const response = await fetch(url);
        //             const content = await response.text();
        //             return {
        //                 content,
        //             };
        //         },
        //     }),
        // },
        onStepFinish: async (result) => {
            if (result.toolResults.length > 0) {
                for (const toolResult of result.toolResults) {
                    let outputType: 'text' | 'json' = 'text';
                    let outputValue: string = '';

                    switch (typeof toolResult.output) {
                        case 'string':
                            outputType = 'text';
                            outputValue = toolResult.output;
                            break;
                        case 'object':
                            outputType = 'json';
                            outputValue = JSON.stringify(toolResult.output, null, 2);
                            break;
                        default:
                            console.error('Unknown output type', toolResult.output);
                            await httpClient.update('tool_calls', toolResult.toolCallId, {
                                status: 'failed',
                                error: {
                                    type: 'text',
                                    value: 'Tool returned invalid output',
                                },
                            });
                            break;
                    }

                    await httpClient.update('tool_calls', toolResult.toolCallId, {
                        status: 'completed',
                        output: {
                            type: outputType,
                            value: outputValue,
                        },
                    });

                    activeToolCalls.delete(toolResult.toolCallId);
                }
            }
        },
        onFinish: async (result) => {
            await httpClient.update('generations', generationId, {
                status: 'completed',
                tokens: {
                    input: result.totalUsage.inputTokens,
                    cache: {
                        read: result.totalUsage.inputTokenDetails.cacheReadTokens,
                        write: result.totalUsage.inputTokenDetails.cacheWriteTokens,
                    },
                    output: result.totalUsage.outputTokens,
                    thinking: result.totalUsage.outputTokenDetails.reasoningTokens,
                },
            });
        },
        onError: async (error: any) => {
            console.error('generation error', error);
            log?.(error);
            // TODO: the docs say "The stream processing will pause until the callback promise is resolved." Suggesting that this error might not be fatal?
            for (const activePart of activeParts.values()) {
                await httpClient.update('message_parts', activePart.id, {
                    finished: true,
                    lastUpdatedAt: new Date(),
                });
            }

            for (const activeToolCall of activeToolCalls.values()) {
                await httpClient.update('tool_calls', activeToolCall.id, {
                    status: 'failed',
                    error: {
                        type: 'text',
                        value: 'An unknown error occurred',
                    },
                });
            }

            await httpClient.update('generations', generationId, {
                status: 'failed',
                error: error.message,
            });
        },
        abortSignal: controller.signal,
    });

    let curStepIdx = -1;
    let key, part, type;

    const pendingUpdates = new Map<string, NodeJS.Timeout>();

    const TARGET_UPDATES_PER_SECOND = 24;

    const scheduleUpdate = (key: string) => {
        const part = activeParts.get(key);
        if (!part || pendingUpdates.has(part.id)) return;

        pendingUpdates.set(part.id, setTimeout(async () => {
            const currentPart = activeParts.get(key);
            // Only update if the part is still active and we haven't deleted it at 'text-end'
            if (currentPart) {
                try {
                    await httpClient.update('message_parts', currentPart.id, {
                        content: currentPart.accumulatedContent,
                        providerOptions: currentPart.providerOptions,
                        lastUpdatedAt: new Date(),
                    });
                } catch (error) {
                    // the update failed, but it doesnt matter because the full message will be updated on step finish
                    console.warn('Failed to update message part', error);
                }
            }
            pendingUpdates.delete(part.id);
        }, 1000 / TARGET_UPDATES_PER_SECOND));
    };

    try {
        for await (const token of response.fullStream) {
            log?.(JSON.stringify(token, null, 2));

            switch (token.type) {
                case 'start-step': {
                    curStepIdx++;
                } break;
                case 'tool-input-start': {
                    key = `tool-call-${curStepIdx}`;

                    const toolCallId = token.id;

                    part = await httpClient.insert('message_parts', {
                        messageId: message.id,
                        toolCallId,
                        type: 'tool-call',
                        content: '',
                        finished: false,
                        createdAt: new Date(),
                        lastUpdatedAt: new Date(),
                    });

                    await httpClient.insert('tool_calls', {
                        id: toolCallId,
                        userId: userId,
                        toolName: token.toolName,
                        status: 'pending',
                        input: null,
                        output: null,
                        error: null,
                        createdAt: new Date(),
                    });

                    activeToolCalls.set(key, { id: toolCallId });

                    activeParts.set(key, { id: part.id, accumulatedContent: '' });
                } break;
                case 'text-start':
                case 'reasoning-start': {
                    type = token.type.split('-')[0];
                    key = `${type}-${curStepIdx}`;

                    part = await httpClient.insert('message_parts', {
                        messageId: message.id,
                        type: type as 'text' | 'reasoning',
                        content: '',
                        finished: false,
                        createdAt: new Date(),
                        lastUpdatedAt: new Date(),
                    });

                    activeParts.set(key, { id: part.id, accumulatedContent: '' });
                } break;
                case 'text-delta':
                case 'reasoning-delta': {
                    type = token.type.split('-')[0] as 'text' | 'reasoning';
                    key = `${type}-${curStepIdx}`;
                    part = activeParts.get(key);
                    if (part === undefined) {
                        console.error('Received delta without a start');
                        break;
                    }

                    let shouldUpdate = false;

                    // TODO: we should potentially merge providerOptions, but for now, just overwrite them
                    if (token.providerMetadata !== undefined) {
                        shouldUpdate = true;
                        part.providerOptions = token.providerMetadata;
                    }

                    // OpenRouter sometimes puts [REDACTED] in thinking if reasoning is encrypted, so we need to remove it and hide it;
                    // do not trim or else we lose intentional whitespace and newlines potentially breaking the UI and having words comebined e.g. "the" "\n\n" "assistant" would become "theassistant"
                    const text = token.text.replaceAll('[REDACTED]', '');
                    if (text !== '') {
                        shouldUpdate = true;
                        part.accumulatedContent += token.text;
                    }

                    if (shouldUpdate) {
                        scheduleUpdate(key);
                    }
                }
                    break;
                case 'text-end':
                case 'reasoning-end': {
                    type = token.type.split('-')[0];
                    key = `${type}-${curStepIdx}`;
                    part = activeParts.get(key);
                    if (part === undefined) {
                        console.error('Received end without a start');
                        break;
                    }

                    activeParts.delete(key);

                    if (part.accumulatedContent === '' && part.providerOptions === undefined) {
                        // completely empty, delete it
                        await httpClient.delete('message_parts', part.id);
                        break;
                    }

                    await httpClient.update('message_parts', part.id, {
                        content: part.accumulatedContent,
                        providerOptions: part.providerOptions,
                        finished: true,
                        lastUpdatedAt: new Date(),
                    });
                } break;
                case 'tool-call': {
                    let inputType: 'text' | 'json' = 'text';
                    let inputValue: string = '';

                    switch (typeof token.input) {
                        case 'string':
                            inputType = 'text';
                            inputValue = token.input;
                            break;
                        case 'object':
                            inputType = 'json';
                            inputValue = JSON.stringify(token.input, null, 2);
                            break;
                        default:
                            console.error('Unknown input type', token.input);
                            break;
                    }

                    await httpClient.update('tool_calls', token.toolCallId, {
                        status: 'pending',
                        input: {
                            type: inputType,
                            value: inputValue,
                        },
                    });
                } break;
                case 'tool-error': {
                    const toolCall = activeToolCalls.get(token.toolCallId);
                    if (toolCall === undefined) {
                        console.error('Received tool-error without a start');
                        break;
                    }

                    let outputType: 'text' | 'json';
                    let outputValue: string;

                    switch (typeof token.error) {
                        case 'string':
                            outputType = 'text';
                            outputValue = token.error;
                            break;
                        case 'object':
                            outputType = 'json';
                            outputValue = JSON.stringify(token.error, null, 2);
                            break;
                        default:
                            console.error('Unknown error type', token.error);
                            outputType = 'text';
                            outputValue = 'Tool returned invalid output';
                            break;
                    }

                    await httpClient.update('tool_calls', toolCall.id, {
                        status: 'failed',
                        error: {
                            type: outputType,
                            value: outputValue,
                        },
                    });
                } break;
                case 'error': {
                    let error = INTERNAL_ERROR;
                    if (typeof token.error === 'string') {
                        error = token.error;
                    } else if (typeof token.error === 'object') {
                        error = JSON.stringify(token.error, null, 2);
                    }

                    await httpClient.update('generations', generationId, {
                        status: 'failed',
                        error,
                    });
                } break;
                case 'finish': {
                    switch (token.finishReason) {
                        case 'error':
                            await httpClient.update('generations', generationId, {
                                status: 'failed',
                                error: INTERNAL_ERROR,
                            });
                            break;
                        case 'content-filter':
                            await httpClient.update('generations', generationId, {
                                status: 'failed',
                                error: 'Content was filtered',
                            });
                            break;
                    }

                    // I hate you switch fallthroughs
                } break;
                case 'abort': {
                    for (const activeToolCall of activeToolCalls.values()) {
                        await httpClient.update('tool_calls', activeToolCall.id, {
                            status: 'cancelled',
                        });
                    }

                    await httpClient.update('generations', generationId, {
                        status: 'cancelled',
                    });
                } break;
                case 'file':
                    todo('file token type', token);
                    break;
                case 'raw':
                    todo('raw token type', token);
                    break;
                case 'source':
                    todo('source token type', token);
                    break;
                case 'tool-approval-request':
                    todo('tool-approval-request token type', token);
                    break;
                case 'tool-output-denied':
                    todo('tool-output-denied token type', token);
                    break;
                case 'start':
                case 'finish-step':
                case 'tool-input-delta':
                case 'tool-input-end':
                case 'tool-result':
                    // handled or irrelevant
                    break;
            }
        }
    } catch (error: any) {
        console.error(error);

        for (const activePart of activeParts.values()) {
            await httpClient.update('message_parts', activePart.id, {
                finished: true,
                lastUpdatedAt: new Date(),
            });
        }

        for (const activeToolCall of activeToolCalls.values()) {
            await httpClient.update('tool_calls', activeToolCall.id, {
                status: 'failed',
                error: {
                    type: 'text',
                    value: 'An unknown error occurred',
                },
            });
        }

        await httpClient.update('generations', generationId, {
            status: 'failed',
            error: error.message,
        });
    } finally {
        completeGeneration(generationId);
        if (logFile !== undefined) logFile.close();
    }
}