import { db } from "~~/server/lib/db";
import * as z from 'zod';
import { type MessageEntity } from '~/composables/useChat';
import { promises as fs } from 'fs';
import { glob } from 'glob';
import { isLoopFinished, type ModelMessage, rerank, type RerankingModel, streamText, type StreamTextTransform, Tool, tool } from "ai";
import { generations, messageParts, messages, toolCalls, ToolCallType } from "~~/drizzle/schema";
import { topicEvents } from "~~/server/utils/events";
import { nanoid } from "nanoid";
import { Monty, MontyRuntimeError, MontySyntaxError, MontyTypingError } from '@pydantic/monty';
import { type Model } from "~/composables/useModels";
import { eq } from "drizzle-orm";
import { buildFocusedMessageTree, buildMessageTree, marshallMessages } from "~~/utils/message";
import path from "path";
import { isRerankingProvider } from "~~/server/utils/ai-provider";

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

    const topicId = getRouterParam(event, 'topicId')!;
    const userId = event.context.user!.id as string;

    const result = await readValidatedBody(event, z.object({
        parentMessageId: z.string().optional(),
        modelId: z.string(),
        rerank: z.object({
            modelId: z.string(),
            providerApiKey: z.string().optional(),
        }).optional(),
        args: z.record(z.string(), z.any()).optional(),
        providerApiKey: z.string().optional(),
    }).safeParse);
    if (!result.success) {
        throw createError({
            statusCode: 400,
            statusMessage: 'Bad Request',
            message: result.error.issues.map(issue => issue.message).join(', '),
        });
    }

    const { parentMessageId, modelId, rerank: rerankConfig, providerApiKey, args } = result.data;

    const topic = await db.query.topics.findFirst({
        where: {
            id: topicId,
            userId,
        },
        with: {
            agent: true,
            messages: {
                orderBy: {
                    createdAt: 'asc',
                },
                with: {
                    parts: {
                        orderBy: {
                            createdAt: 'asc',
                        },
                        with: {
                            toolCall: true,
                        }
                    },
                    attachments: {
                        with: {
                            file: true,
                        }
                    },
                    generation: true,
                }
            },
        },
    });
    if (!topic) {
        throw createError({
            statusCode: 404,
            statusMessage: 'Not Found',
            message: 'Topic not found',
        });
    }

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

    if (!model) {
        throw createError({
            statusCode: 404,
            statusMessage: 'Not Found',
            message: 'Model not found',
        });
    }

    const providerDetails = await getProviderDetails(model.provider, providerApiKey, model);
    if (!providerDetails.ok) {
        switch (providerDetails.error) {
            case GatewayFetchError.NoProviderApiKey: {
                throw createError({
                    statusCode: 400,
                    message: `${model.provider.type} provider requires an API key`,
                });
            }
            case GatewayFetchError.NoProviderBaseUrl: {
                throw createError({
                    statusCode: 400,
                    message: 'Invalid provider URL',
                });
            }
        }
    }

    const agentSearchConfig = topic.agent.config?.search;
    let searchParam: false | { config: SearchTheWebConfig } = false;

    if (agentSearchConfig?.enabled) {
        searchParam = { config: { rerank: false, maxResults: agentSearchConfig.maxResults ?? 10 } };

        if (agentSearchConfig.rerank && rerankConfig) {
            const rerankModel = await db.query.models.findFirst({
                where: {
                    id: rerankConfig.modelId,
                    userId,
                },
                with: {
                    provider: true,
                }
            });

            if (rerankModel === undefined) {
                throw createError({
                    statusCode: 400,
                    statusMessage: 'Invalid rerank model',
                    data: {
                        code: 'INVALID_RERANK_MODEL',
                        ok: false,
                    }
                });
            }

            const rerankProviderDetails = await getProviderDetails(rerankModel.provider, rerankConfig.providerApiKey, rerankModel);
            if (!rerankProviderDetails.ok) {
                switch (rerankProviderDetails.error) {
                    case GatewayFetchError.NoProviderApiKey: {
                        throw createError({
                            statusCode: 400,
                            statusMessage: `${rerankModel.provider.type} provider requires an API key`,
                            data: {
                                code: 'NO_RERANK_PROVIDER_API_KEY',
                                ok: false,
                            }
                        });
                    }
                    case GatewayFetchError.NoProviderBaseUrl: {
                        throw createError({
                            statusCode: 400,
                            statusMessage: 'Invalid provider URL',
                            data: {
                                code: 'BAD_RERANK_PROVIDER_URL',
                                ok: false,
                            }
                        });
                    }
                }
            }

            const { gateway: rerankGateway } = rerankProviderDetails.data;
            if (rerankGateway === null) {
                throw createError({
                    statusCode: 500,
                    statusMessage: 'Invalid gateway',
                    data: {
                        code: 'INVALID_RERANK_GATEWAY',
                        ok: false,
                    }
                });
            }

            if (isRerankingProvider(rerankGateway.gateway)) {
                searchParam.config = {
                    rerank: true,
                    maxResults: agentSearchConfig.maxResults ?? 10,
                    model: rerankGateway.gateway.reranking(rerankModel.externalId)
                };
            } else {
                throw createError({
                    statusCode: 500,
                    statusMessage: 'Invalid rerank model',
                    data: {
                        code: 'INVALID_RERANK_MODEL',
                        ok: false,
                    }
                });
            }
        }
    }

    let agentmessage = await db.transaction(async tx => {
        const generationId = nanoid();

        const [generation] = await tx.insert(generations).values({
            id: generationId,
            userId,
            topicId,
            modelId: model.externalId,
            status: 'pending',
        }).returning();

        if (!generation) {
            throw createError({
                statusCode: 500,
                statusMessage: 'Failed to insert generation',
                message: 'Failed to insert generation',
            });
        }

        const [agentmessage] = await tx.insert(messages).values({
            id: nanoid(),
            userId,
            topicId,
            content: null,
            role: 'assistant',
            generationId,
            parentMessageId,
        }).returning();

        if (!agentmessage) {
            throw createError({
                statusCode: 500,
                statusMessage: 'Failed to insert message',
                message: 'Failed to insert message',
            });
        }

        agentmessage.generationId = generation.id;
        // @ts-ignore
        agentmessage.generation = generation;

        return agentmessage;
    });
    const events = [{ type: 'MESSAGE_CREATED', payload: agentmessage }] as { type: string; payload: any }[];

    let topicMessageTree = buildMessageTree(topic.messages);

    if (parentMessageId) {
        const parentMessage = await db.query.messages.findFirst({
            where: {
                id: parentMessageId,
            }
        });
        if (!parentMessage) {
            throw createError({
                statusCode: 500,
                statusMessage: 'Failed to find parent message',
                message: 'Failed to find parent message',
            });
        }

        await db.update(messages).set({ focusedIndex: (parentMessage.focusedIndex ?? 0) + 1 }).where(eq(messages.id, parentMessageId));
        events.push({ type: 'MESSAGE_UPDATED', payload: { focusedIndex: (parentMessage.focusedIndex ?? 0) + 1 } });

        // we need to make the topic messages all the messages excluding the ones after the message we wish to regenerate
        // and if parentMessageId is undefined, then excluding the last message
        if (parentMessageId === undefined) {
            topicMessageTree = topicMessageTree.slice(0, topicMessageTree.length - 1);
        } else {
            const parentMessageIdx = topicMessageTree.findIndex(m => m.id === parentMessageId);
            if (parentMessageIdx === -1) {
                throw createError({
                    statusCode: 500,
                    statusMessage: 'Failed to find parent message',
                    message: 'Failed to find parent message',
                });
            }

            topicMessageTree = topicMessageTree.slice(0, parentMessageIdx);
        }
    }

    topicEvents.emit(topicId, { type: 'MESSAGE_CREATED', payload: agentmessage });

    const { gateway } = providerDetails.data;
    if (gateway === null) {
        throw createError({
            statusCode: 500,
            statusMessage: 'Invalid gateway',
            data: {
                code: 'INVALID_GATEWAY',
                ok: false,
            }
        });
    }

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

    if (process.env.GENERATION_DEBUG) {
        if (process.env.LOG_DIR) {
            await fs.mkdir(process.env.LOG_DIR!, { recursive: true });

            logFile = await fs.open(path.join(process.env.LOG_DIR!, `${Date.now()}-${agentmessage.generationId!}.log`), 'w');
            logMessage = (message: string) => {
                logFile!.write(message + '\n');
            };
        } else {
            console.warn('Generation debug logging is enabled but LOG_DIR is not set');
        }
    }

    const topicMessages = marshallMessages(topic.agent, buildFocusedMessageTree(topicMessageTree));
    if (topicMessages.ok === false) {
        throw createError({
            statusCode: 500,
            statusMessage: 'Failed to marshall messages',
            message: topicMessages.error,
        });
    }

    event.waitUntil(
        generateResponse(
            agentmessage as MessageEntity,
            { gateway: gateway.gateway, model, parameters: args },
            searchParam,
            agentmessage.generationId!,
            userId,
            topicId,
            topicMessages.data,
            gateway.streamTransformer,
            logMessage,
            logFile
        )
    );

    return {
        ok: true
    };
});


const INTERNAL_ERROR = 'An internal error occurred';

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

const formatPartId = (partType: string, existingId: string) => {
    return `veridian__part-${partType}-${existingId}-${nanoid()}`;
};

const formatToolCallId = (nativeId: string) => {
    return `veridian__tool-${nativeId}-${nanoid()}`;
};

const evalPython = async (code: string) => {
    try {
        let stdout = ''
        const printCallback = (_: string, text: string) => {
            stdout += text
        }
        const m = new Monty(code)
        await m.run({ printCallback })
        return stdout
    } catch (error) {
        if (error instanceof MontySyntaxError) {
            console.log('Syntax error:', error.message)
        } else if (error instanceof MontyRuntimeError) {
            console.log('Runtime error:', error.message)
            console.log('Traceback:', error.traceback())
        } else if (error instanceof MontyTypingError) {
            console.log('Type error:', error.displayDiagnostics())
        }
    }
};

interface SearchTheWebRerankedConfig {
    rerank: true;
    maxResults: number;
    model: RerankingModel;
}

interface SearchTheWebUnrankedConfig {
    rerank?: false;
    maxResults: number;
}

type SearchTheWebConfig = SearchTheWebRerankedConfig | SearchTheWebUnrankedConfig;

export const searchTheWeb = (config: SearchTheWebConfig) => {
    return async (query: string) => {
        const results = await $fetch<any>(`${process.env.SEARXNG_URL}/search`, {
            query: {
                q: query,
                format: 'json',
            }
        });

        const sites = results.results.map((item: any) => ({
            title: item.title,
            link: item.url,
            snippet: item.content,
        }));

        if (config.rerank) {
            const { ranking } = await rerank({
                model: config.model,
                query,
                documents: sites.map(site => site.snippet),
            });

            const ranked_sites = [];

            for (let i = 0; i < ranking.length; i++) {
                ranked_sites.push(sites[ranking[i]!.originalIndex]);
            }

            return ranked_sites.slice(0, config.maxResults || 10);
        }

        return sites.slice(0, config.maxResults || 10);
    };
};

// TODO: obviously come up with a better way for the user to define their own tools
const { listDirectoryTool, globTool, readFileTool, readFilesTool, fetchUrlTool, pythonTool, bashTool } = {
    listDirectoryTool: tool({
        description: 'Lists the contents of a directory',
        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,
            };
        },
    }),
    globTool: tool({
        description: 'Lists files matching a glob pattern',
        inputSchema: z.object({
            pattern: z.string(),
        }),
        outputSchema: z.object({
            files: z.array(z.object({ path: z.string(), type: z.string() })),
        }),
        execute: async ({ pattern }) => {
            const rawFiles = await glob(pattern, { withFileTypes: true });
            const files = rawFiles.map((file) => ({
                path: file.parentPath + '/' + file.name,
                type: file.isFile() ? 'file' : 'directory',
            }));

            return {
                files,
            };
        },
    }),
    readFileTool: tool({
        description: 'Reads the contents of a file',
        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(),
            };
        },
    }),
    readFilesTool: tool({
        description: 'Reads the contents of multiple files',
        inputSchema: z.object({
            paths: z.array(z.string()).describe('a list of 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,
            };
        },
    }),
    fetchUrlTool: tool({
        description: 'Fetches the content of a URL',
        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,
            };
        },
    }),
    pythonTool: tool({
        description: 'Executes a Python code snippet',
        inputSchema: z.object({
            code: z.string(),
        }),
        outputSchema: z.object({
            output: z.string(),
        }),
        execute: async ({ code }) => {
            const output = await evalPython(code);
            return {
                output,
            };
        },
    }),
    bashTool: tool({
        description: 'Executes a Bash command',
        inputSchema: z.object({
            code: z.string(),
        }),
        outputSchema: z.object({
            output: z.string(),
        }),
        execute: async ({ code }) => {
            const { exec } = await import('child_process');
            const { promisify } = await import('util');
            const execAsync = promisify(exec);

            async function runCommand(command: string) {
                const { stdout, stderr } = await execAsync(command);
                if (stderr) {
                    console.error(`Error: ${stderr}`);
                    return stderr;
                } else {
                    return stdout;
                }
            }

            const output = await runCommand(code);
            return {
                output,
            };
        },
    }),
}

async function generateResponse(
    message: MessageEntity,
    model: {
        gateway: ModelGateway,
        model: Model,
        parameters?: Record<string, any>,
    },
    search: false | {
        config: SearchTheWebConfig,
    },
    generationId: string,
    userId: string,
    topicId: string,
    messages: ModelMessage[],
    streamTransoforms: StreamTextTransform<{}> | StreamTextTransform<{}>[] | undefined,
    log?: (message: string) => void,
    logFile?: fs.FileHandle,
) {
    const controller = new AbortController();
    addPendingGeneration(generationId, controller);

    let requestStart = undefined;
    let ttft = undefined;
    const activeParts = new Map<string, { id: string; accumulatedContent: string; providerOptions?: any }>();
    const activeToolCalls = new Set<string>();
    const nativeToDbToolCallId = new Map<string, string>();

    // TODO: somehow let the user turn on and off tools
    const tools: Record<string, Tool> = {
        // 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: listDirectoryTool,
        glob: globTool,
        readFile: readFileTool,
        readFiles: readFilesTool,
        // fetchUrl: fetchUrlTool,
        python: pythonTool,
        bash: bashTool,
    };

    if (search) {
        tools.search = tool({
            description: 'Searches the web',
            inputSchema: z.object({
                query: z.string(),
            }),
            outputSchema: z.array(
                z.object({
                    title: z.string(),
                    link: z.string(),
                    snippet: z.string(),
                    engine: z.string(),
                })
            ),
            execute: async ({ query }) => {
                return await searchTheWeb(search.config)(query);
            },
        });
    }

    const response = streamText({
        model: model.gateway(model.model.externalId),
        messages,
        providerOptions: {
            openrouter: {
                debug: {
                    echo_upstream_body: true,
                },
                user: userId,
            }
        },
        experimental_transform: streamTransoforms,
        stopWhen: isLoopFinished(),
        tools: model.model.capabilities.includes('tools') ? tools : undefined,
        onError: async (error: any) => {
            // TODO: the docs say "The stream processing will pause until the callback promise is resolved." Suggesting that this error might not be fatal?
            console.error('generation error', error);
            log?.(error);

            // Mark all active parts as finished
            await Promise.all([...activeParts.values()].map(async part => {
                await db.update(messageParts).set({ finished: true, lastUpdatedAt: new Date() }).where(eq(messageParts.id, part.id))
                await topicEvents.emit(topicId, {
                    type: 'text-end',
                    payload: {
                        messageId: message.id,
                        partId: part.id,
                        content: part.accumulatedContent,
                    }
                })
            }));

            // Mark failed tool calls
            await Promise.all([...activeToolCalls].map(async id => {
                await db.update(toolCalls).set({
                    status: 'failed',
                    error: { type: ToolCallType.Text, value: 'Generation failed' }
                }).where(eq(toolCalls.id, id))

                await topicEvents.emit(topicId, {
                    type: 'tool-call-delta',
                    payload: {
                        messageId: message.id,
                        toolCallId: id,
                        toolName: null,
                        error: {
                            type: ToolCallType.Text,
                            value: 'Generation failed'
                        }
                    }
                })
            }));

            await db.update(generations).set({
                status: 'failed',
                error: typeof error === 'string' ? error : JSON.stringify(error)
            }).where(eq(generations.id, generationId));

            await topicEvents.emit(topicId, {
                type: 'generation-failed',
                payload: {
                    messageId: message.id,
                    generationId,
                    error: typeof error === 'string' ? error : JSON.stringify(error),
                }
            })
        },
        // onStepFinish: async (step) => {
        //     step.content.forEach(async (part) => {
        //         switch (part.type) {
        //             case 'text': {
        //                 await db.insert(messageParts).values({
        //                     userId,
        //                     topicId,
        //                     messageId: message.id,
        //                     type: 'text',
        //                     content: part.text,
        //                     providerOptions: part.providerMetadata,
        //                     finished: true,
        //                     createdAt: new Date(),
        //                     lastUpdatedAt: new Date(),
        //                 });
        //             } break;
        //         }
        //     })
        // },
        abortSignal: controller.signal,
    });

    let curStepIdx = -1;
    let key, type;

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

    // since we arent streaming straight from the database, we can update less often
    const TARGET_UPDATES_PER_SECOND = 10;

    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);
            if (!currentPart) return;

            try {
                await db.update(messageParts)
                    .set({
                        content: currentPart.accumulatedContent,
                        providerOptions: currentPart.providerOptions,
                        lastUpdatedAt: new Date(),
                    })
                    .where(eq(messageParts.id, currentPart.id));
            } catch (e) {
                console.warn('Failed to update message part', e);
            }
            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': {
                    requestStart = performance.now();
                    break;
                }

                case 'start-step': {
                    curStepIdx++;
                    break;
                }

                case 'text-start':
                case 'reasoning-start': {
                    const type = token.type.split('-')[0] as 'text' | 'reasoning';
                    const key = `${type}-${curStepIdx}`;
                    const nativeId = nanoid();

                    const [part] = await db.insert(messageParts).values({
                        id: formatPartId(type, nativeId),
                        userId,
                        topicId,
                        messageId: message.id,
                        type,
                        content: '',
                        finished: false,
                        createdAt: new Date(),
                        lastUpdatedAt: new Date(),
                    }).returning();

                    if (!part) {
                        throw new Error('Failed to insert message part');
                    }

                    await topicEvents.emit(topicId, {
                        type: 'text-start',
                        payload: {
                            messageId: message.id,
                            part,
                        }
                    })

                    activeParts.set(key, { id: part.id, accumulatedContent: '' });
                    break;
                }

                case 'text-delta':
                case 'reasoning-delta': {
                    if (ttft === undefined) {
                        ttft = performance.now() - requestStart!;
                    }

                    type = token.type.split('-')[0] as 'text' | 'reasoning';
                    key = `${type}-${curStepIdx}`;
                    const 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;
                        await topicEvents.emit(topicId, {
                            type: 'text-delta',
                            payload: {
                                messageId: message.id,
                                partId: part.id,
                                content: token.text,
                            }
                        })
                    }

                    if (shouldUpdate) {
                        scheduleUpdate(key);
                    }

                    break;
                }

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

                    activeParts.delete(key);

                    if (part.accumulatedContent === '' && !part.providerOptions) {
                        await db.delete(messageParts).where(eq(messageParts.id, part.id));

                        await topicEvents.emit(topicId, {
                            type: 'text-delete',
                            payload: {
                                messageId: message.id,
                                partId: part.id,
                            }
                        })
                    } else {
                        const [dbPart] = await db.update(messageParts)
                            .set({
                                content: part.accumulatedContent,
                                providerOptions: part.providerOptions,
                                finished: true,
                                lastUpdatedAt: new Date(),
                            })
                            .where(eq(messageParts.id, part.id)).returning();

                        if (!dbPart) {
                            throw new Error('Failed to update message part');
                        }

                        await topicEvents.emit(topicId, {
                            type: 'text-end',
                            payload: {
                                messageId: message.id,
                                partId: part.id,
                                content: part.accumulatedContent,
                            }
                        })
                    }

                    break;
                }

                case 'tool-input-start': {
                    key = `tool-call-${curStepIdx}`;

                    const toolCallId = token.id;
                    const dbToolCallId = formatToolCallId(toolCallId);
                    nativeToDbToolCallId.set(toolCallId, dbToolCallId);

                    const [toolCall] = await db.insert(toolCalls).values({
                        id: dbToolCallId,
                        userId: userId,
                        toolName: token.toolName,
                        status: 'pending',
                        input: null,
                        output: null,
                        error: null,
                        createdAt: new Date(),
                    }).returning();

                    if (!toolCall) {
                        throw new Error('Failed to insert tool call');
                    }

                    const [part] = await db.insert(messageParts).values({
                        id: formatPartId('tool-call', dbToolCallId),
                        userId,
                        topicId,
                        messageId: message.id,
                        toolCallId: dbToolCallId,
                        type: 'tool-call',
                        content: null,
                        finished: false,
                        createdAt: new Date(),
                        lastUpdatedAt: new Date(),
                    }).returning();

                    if (!part) {
                        throw new Error('Failed to insert message part');
                    }

                    // @ts-ignore
                    part.toolCall = toolCall;

                    await topicEvents.emit(topicId, {
                        type: 'tool-call-start',
                        payload: {
                            messageId: message.id,
                            partId: part.id,
                            part,
                        }
                    })

                    activeToolCalls.add(dbToolCallId);

                    activeParts.set(key, { id: part.id, accumulatedContent: '' });
                } break;
                case 'tool-call': {
                    let inputType: ToolCallType = ToolCallType.Text;
                    let inputValue: string = '';

                    switch (typeof token.input) {
                        case 'string':
                            inputType = ToolCallType.Text;
                            inputValue = token.input;
                            break;
                        case 'object':
                            inputType = ToolCallType.Json;
                            inputValue = JSON.stringify(token.input);
                            break;
                        default:
                            console.error('Unknown input type', token.input);
                            break;
                    }

                    const dbToolCallIdFromMap = nativeToDbToolCallId.get(token.toolCallId);
                    if (dbToolCallIdFromMap && activeToolCalls.has(dbToolCallIdFromMap)) {
                        const dbToolCallId = dbToolCallIdFromMap;
                        await db.update(toolCalls)
                            .set({
                                status: 'pending',
                                input: {
                                    type: inputType,
                                    value: inputValue,
                                },
                            })
                            .where(eq(toolCalls.id, dbToolCallId));

                        await topicEvents.emit(topicId, {
                            type: 'tool-call-delta',
                            payload: {
                                messageId: message.id,
                                toolCallId: dbToolCallId,
                                toolName: token.toolName,
                                input: {
                                    type: inputType,
                                    value: inputValue,
                                },
                            }
                        })
                    } else {
                        key = `tool-call-${curStepIdx}`;

                        const toolCallId = token.toolCallId;
                        const dbToolCallId = formatToolCallId(toolCallId);
                        nativeToDbToolCallId.set(toolCallId, dbToolCallId);

                        const [toolCall] = await db.insert(toolCalls).values({
                            id: dbToolCallId,
                            userId: userId,
                            toolName: token.toolName,
                            status: 'pending',
                            input: {
                                type: inputType,
                                value: inputValue,
                            },
                            output: null,
                            error: null,
                            createdAt: new Date(),
                        }).returning();

                        if (!toolCall) {
                            throw new Error('Failed to insert tool call');
                        }

                        const [part] = await db.insert(messageParts).values({
                            id: formatPartId('tool-call', dbToolCallId),
                            userId,
                            topicId,
                            messageId: message.id,
                            toolCallId: dbToolCallId,
                            type: 'tool-call',
                            content: null,
                            finished: false,
                            createdAt: new Date(),
                            lastUpdatedAt: new Date(),
                        }).returning();

                        if (!part) {
                            throw new Error('Failed to insert message part');
                        }

                        // @ts-ignore
                        part.toolCall = toolCall;

                        await topicEvents.emit(topicId, {
                            type: 'tool-call-start',
                            payload: {
                                messageId: message.id,
                                part,
                            }
                        })

                        activeToolCalls.add(dbToolCallId);

                        activeParts.set(key, { id: part.id, accumulatedContent: '' });
                    }
                } break;
                case 'tool-result': {
                    let outputType: ToolCallType = ToolCallType.Text;
                    let outputValue: string = '';
                    const dbToolCallId = nativeToDbToolCallId.get(token.toolCallId);

                    switch (typeof token.output) {
                        case 'string':
                            outputType = ToolCallType.Text;
                            outputValue = token.output;
                            break;
                        case 'object':
                            outputType = ToolCallType.Json;
                            outputValue = JSON.stringify(token.output);
                            break;
                        default:
                            console.error('Unknown output type', token.output);
                            if (dbToolCallId) {
                                await db.update(toolCalls).set({
                                    status: 'failed',
                                    error: { type: ToolCallType.Text, value: 'Tool returned invalid output' }
                                }).where(eq(toolCalls.id, dbToolCallId));

                                await topicEvents.emit(topicId, {
                                    type: 'tool-call-delta',
                                    payload: {
                                        messageId: message.id,
                                        toolCallId: dbToolCallId,
                                        toolName: token.toolName,
                                        output: {
                                            type: outputType,
                                            value: outputValue,
                                        },
                                        status: 'failed',
                                    }
                                })

                                activeToolCalls.delete(dbToolCallId);
                            }
                            break;
                    }

                    if (dbToolCallId) {
                        await db.update(toolCalls).set({
                            status: 'completed',
                            output: {
                                type: outputType,
                                value: outputValue,
                            },
                        }).where(eq(toolCalls.id, dbToolCallId));

                        await topicEvents.emit(topicId, {
                            type: 'tool-call-delta',
                            payload: {
                                messageId: message.id,
                                toolCallId: dbToolCallId,
                                toolName: token.toolName,
                                output: {
                                    type: outputType,
                                    value: outputValue,
                                },
                                status: 'completed',
                            }
                        })

                        activeToolCalls.delete(dbToolCallId);
                    }
                    break;
                }

                case 'tool-error': {
                    console.error('Tool error:', token);

                    let outputType: ToolCallType;
                    let outputValue: string;

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

                    const existingDbToolCallId = nativeToDbToolCallId.get(token.toolCallId);
                    if (existingDbToolCallId && activeToolCalls.has(existingDbToolCallId)) {
                        await db.update(toolCalls).set({
                            status: 'failed',
                            error: {
                                type: outputType as ToolCallType,
                                value: outputValue as string,
                            }
                        }).where(eq(toolCalls.id, existingDbToolCallId));

                        await topicEvents.emit(topicId, {
                            type: 'tool-call-delta',
                            payload: {
                                messageId: message.id,
                                toolCallId: existingDbToolCallId,
                                toolName: token.toolName,
                                error: {
                                    type: outputType as ToolCallType,
                                    value: outputValue as string,
                                }
                            }
                        })
                    } else {
                        key = `tool-call-${curStepIdx}`;

                        const toolCallId = token.toolCallId;
                        const dbToolCallId = formatToolCallId(toolCallId);
                        nativeToDbToolCallId.set(toolCallId, dbToolCallId);

                        const [toolCall] = await db.insert(toolCalls).values({
                            id: dbToolCallId,
                            userId: userId,
                            toolName: token.toolName,
                            status: 'failed',
                            input: null,
                            output: null,
                            error: {
                                type: outputType,
                                value: outputValue,
                            },
                            createdAt: new Date(),
                        }).returning();

                        if (!toolCall) {
                            throw new Error('Failed to insert tool call');
                        }

                        const [part] = await db.insert(messageParts).values({
                            id: formatPartId('tool-call', dbToolCallId),
                            userId,
                            topicId,
                            messageId: message.id,
                            toolCallId: dbToolCallId,
                            type: 'tool-call',
                            content: null,
                            finished: false,
                            createdAt: new Date(),
                            lastUpdatedAt: new Date(),
                        }).returning();

                        if (!part) {
                            throw new Error('Failed to insert message part');
                        }

                        await topicEvents.emit(topicId, {
                            type: 'tool-call-start',
                            payload: {
                                messageId: message.id,
                                part,
                            }
                        })

                        activeParts.set(key, { id: part.id, accumulatedContent: '' });
                    }

                    if (existingDbToolCallId) {
                        activeToolCalls.delete(existingDbToolCallId);
                    }
                    break;
                }

                case 'finish': {
                    let tps;
                    if (ttft !== undefined && token.totalUsage.outputTokens !== undefined) {
                        const tokenStreamStart = requestStart! + ttft;
                        // this is the *real* request duration, excluding the 
                        // TTFT
                        const requestDuration = performance.now() - tokenStreamStart;

                        tps = token.totalUsage.outputTokens / (requestDuration / 1000);
                    }

                    await Promise.all([...activeParts.values()].map(async part => {
                        await db.update(messageParts)
                            .set({ finished: true, lastUpdatedAt: new Date() })
                            .where(eq(messageParts.id, part.id))

                        await topicEvents.emit(topicId, {
                            type: 'text-end',
                            payload: {
                                messageId: message.id,
                                partId: part.id,
                                content: part.accumulatedContent,
                            }
                        })
                    }));


                    switch (token.finishReason) {
                        case 'error':
                            await db.update(generations).set({
                                status: 'failed',
                                error: INTERNAL_ERROR,
                            }).where(eq(generations.id, generationId));

                            await topicEvents.emit(topicId, {
                                type: 'generation-failed',
                                payload: {
                                    messageId: message.id,
                                    generationId,
                                    error: INTERNAL_ERROR,
                                }
                            })
                            break;
                        case 'content-filter':
                            await db.update(generations).set({
                                status: 'failed',
                                error: 'Content was filtered',
                            }).where(eq(generations.id, generationId));

                            await topicEvents.emit(topicId, {
                                type: 'generation-failed',
                                payload: {
                                    messageId: message.id,
                                    generationId,
                                    error: 'Content was filtered',
                                }
                            })
                            break;
                    }

                    await db.update(generations).set({
                        status: 'completed',
                        tokens: {
                            input: token.totalUsage.inputTokens,
                            cache: {
                                read: token.totalUsage.inputTokenDetails.cacheReadTokens,
                                write: token.totalUsage.inputTokenDetails.cacheWriteTokens
                            },
                            output: token.totalUsage.outputTokens,
                            thinking: token.totalUsage.outputTokenDetails.reasoningTokens,
                            ttft,
                            tps,
                        },
                    }).where(eq(generations.id, generationId));

                    await topicEvents.emit(topicId, {
                        type: 'generation-complete',
                        payload: {
                            messageId: message.id,
                            generationId,
                            tokens: {
                                input: token.totalUsage.inputTokens,
                                cache: {
                                    read: token.totalUsage.inputTokenDetails.cacheReadTokens,
                                    write: token.totalUsage.inputTokenDetails.cacheWriteTokens
                                },
                                output: token.totalUsage.outputTokens,
                                thinking: token.totalUsage.outputTokenDetails.reasoningTokens,
                                ttft,
                                tps,
                            },
                        }
                    })

                    // I hate you switch fallthroughs
                    break;
                }

                case 'abort': {
                    await Promise.all([
                        ...[...activeParts.values()].map(async part => {
                            await db.update(messageParts)
                                .set({ finished: true, lastUpdatedAt: new Date() })
                                .where(eq(messageParts.id, part.id))

                            await topicEvents.emit(topicId, {
                                type: 'text-end',
                                payload: {
                                    messageId: message.id,
                                    partId: part.id,
                                    content: part.accumulatedContent,
                                }
                            })
                        }),
                        ...[...activeToolCalls].map(async id => {
                            await db.update(toolCalls)
                                .set({ status: 'cancelled' })
                                .where(eq(toolCalls.id, id))

                            await topicEvents.emit(topicId, {
                                type: 'tool-call-cancel',
                                payload: {
                                    messageId: message.id,
                                    toolCallId: id,
                                }
                            })
                        }),
                        db.update(generations)
                            .set({ status: 'cancelled' })
                            .where(eq(generations.id, generationId)),
                        topicEvents.emit(topicId, {
                            type: 'generation-cancelled',
                            payload: {
                                messageId: message.id,
                                generationId,
                            }
                        })
                    ]);
                } 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;
                // typescript thinks this is not a real token type?
                // case 'tool-output-denied':
                //     todo('tool-output-denied token type', token);
                //     break;
                case 'error':
                case 'finish-step':
                case 'tool-input-delta':
                case 'tool-input-end':
                    // handled or irrelevant
                    break;
            }
        }
    } catch (error: any) {
        console.error(error);

        await db.update(generations)
            .set({ status: 'failed', error: String(error) })
            .where(eq(generations.id, generationId));

        await topicEvents.emit(topicId, {
            type: 'generation-failed',
            payload: {
                messageId: message.id,
                generationId,
                error: String(error),
            }
        })
    } finally {
        completeGeneration(generationId);
        if (logFile !== undefined) logFile.close();
    }
}