import { protectRoute } from '~~/server/utils/auth';
import { getPendingGeneration, startGeneration, addClientToGeneration, removeClientFromGeneration, sendToClient } from '~~/server/utils/generation';
import { eventHandler, setHeader, setResponseStatus } from 'h3';
import { generations, messages } from '~~/db/schema';
import { eq } from 'drizzle-orm';

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

    const generationId = getRouterParam(event, 'id');

    if (!generationId) {
        throw createError({
            statusCode: 400,
            statusMessage: 'Missing generation ID'
        });
    }

    const pendingGeneration = getPendingGeneration(generationId);

    if (pendingGeneration && pendingGeneration.expired) {
        throw createError({
            statusCode: 410,
            statusMessage: 'Generation expired - no client connected within 60 seconds'
        });
    }

    setHeader(event, 'Content-Type', 'text/event-stream');
    setHeader(event, 'Cache-Control', 'no-cache');
    setHeader(event, 'Connection', 'keep-alive');
    setHeader(event, 'X-Accel-Buffering', 'no');

    setResponseStatus(event, 200);

    const shouldStartGeneration = pendingGeneration;

    try {
        const stream = new ReadableStream({
            async start(controller) {
                addClientToGeneration(generationId, controller);

                if (shouldStartGeneration) {
                    startGeneration(generationId, controller);
                } else {
                    const generation = await useDrizzle().select().from(generations).where(eq(generations.id, getRouterParam(event, 'id')!))
                    if (!generation) throw createError({ statusCode: 404, statusMessage: 'Generation not found' });
                    const message = await useDrizzle().select().from(messages).where(eq(messages.id, generation[0].messageId!))
                    if (!message) throw createError({ statusCode: 404, statusMessage: 'Message not found' });

                    sendToClient(controller, {
                        type: 'complete',
                        data: message[0].content
                    });
                    controller.close();
                }

                event.node.req.on('close', () => {
                    removeClientFromGeneration(generationId, controller);
                });
            }
        });

        return sendStream(event, stream);
    } catch (error) {
        // assume it failed because the generation is complete so try to send the message
    }
});