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

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

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

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

    // Fetch generation from database
    const generation = await getGenerationStatus(generationId);
    if (!generation) {
        throw createError({
            statusCode: 404,
            statusMessage: 'Generation not found'
        });
    }

    // Verify ownership
    if (generation.userId !== event.context.user.id) {
        throw createError({
            statusCode: 403,
            statusMessage: 'Unauthorized'
        });
    }

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

    try {
        const stream = new ReadableStream({
            async start(controller) {
                try {
                    // If generation is already completed, send the completed message
                    if (generation.status === 'completed' && generation.messageId) {
                        const db = useDrizzle();
                        const [message] = await db
                            .select()
                            .from(messages)
                            .where(eq(messages.id, generation.messageId));

                        if (message) {
                            sendToClient(controller, {
                                type: 'complete',
                                data: message
                            });
                        }
                        controller.close();
                        return;
                    }

                    // If generation failed, send the error
                    if (generation.status === 'failed') {
                        sendToClient(controller, {
                            type: 'error',
                            data: { error: generation.error || 'Generation failed' }
                        });
                        controller.close();
                        return;
                    }

                    // If already streaming, just add this client
                    if (isGenerationStreaming(generationId)) {
                        addClientToGeneration(generationId, controller);
                    } else {
                        // Start generation if in pending status
                        if (generation.status === 'pending') {
                            // Fetch the original messages context (stored in topic messages)
                            const db = useDrizzle();
                            const topicMessages = await db
                                .select()
                                .from(messages)
                                .where(eq(messages.topicId, generation.topicId));

                            const chatMessages = topicMessages.map(m => ({
                                type: m.isUser ? 'user' as const : ('agent' as const),
                                message: m.content
                            }));

                            addClientToGeneration(generationId, controller);
                            await startGeneration(generationId, generation.userId, generation.topicId, chatMessages, controller);
                        }
                    }

                    event.node.req.on('close', () => {
                        removeClientFromGeneration(generationId, controller);
                    });
                } catch (error) {
                    console.error('Stream start error:', error);
                    sendToClient(controller, {
                        type: 'error',
                        data: { error: 'Stream initialization failed' }
                    });
                    controller.close();
                }
            }
        });

        return sendStream(event, stream);
    } catch (error) {
        console.error('Stream error:', error);
        throw createError({
            statusCode: 500,
            statusMessage: 'Stream error'
        });
    }
});