veridian/server/utils/generation.ts

import { useDrizzle } from '~~/server/utils/drizzle';
import { generations, messages as messages_drizzle } from '~~/db/schema';
import { type GenerationStreamEvent, type ChatMessage, type MessageType } from '~~/server/types/chat';
import { eq } from 'drizzle-orm';

interface ActiveGeneration {
    userId: string;
    topicId: string;
    messages: ChatMessage[];
    content: string;
    clients: Set<ReadableStreamDefaultController<Uint8Array>>;
    complete: boolean;
}

interface PendingGeneration {
    generationId: string;
    userId: string;
    topicId: string;
    messages: ChatMessage[];
    timeout: NodeJS.Timeout;
    expired: boolean;
}

const activeGenerations = new Map<string, ActiveGeneration>();
const pendingGenerations = new Map<string, PendingGeneration>();

export const getActiveGeneration = (generationId: string): ActiveGeneration | undefined => {
    return activeGenerations.get(generationId);
};

export const getPendingGeneration = (generationId: string): PendingGeneration | undefined => {
    return pendingGenerations.get(generationId);
};

export const isGenerationActive = (generationId: string): boolean => {
    return activeGenerations.has(generationId);
};

export const isGenerationPending = (generationId: string): boolean => {
    return pendingGenerations.has(generationId);
};

export const addClientToGeneration = (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): boolean => {
    const generation = activeGenerations.get(generationId);
    if (!generation) {
        return false;
    }
    generation.clients.add(controller);
    return true;
};

export const removeClientFromGeneration = (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): void => {
    const generation = activeGenerations.get(generationId);
    if (generation) {
        generation.clients.delete(controller);
        if (generation.clients.size === 0 && generation.complete) {
            activeGenerations.delete(generationId);
        }
    }
};

const sendToClients = (generation: ActiveGeneration, event: GenerationStreamEvent): void => {
    for (const client of generation.clients) {
        sendToClient(client, event);
    }
};

export const sendToClient = (client: ReadableStreamDefaultController<Uint8Array>, event: GenerationStreamEvent): void => {
    const data = JSON.stringify(event);
    const encoder = new TextEncoder();
    try {
        client.enqueue(encoder.encode(`${data}\n`));
    } catch (error) {
        console.error('Failed to send to client:', error);
    }
}

const buildPrompt = (messages: ChatMessage[]): string => {
    return messages
        .map((msg: ChatMessage) => {
            const roleMap: Record<MessageType, string> = {
                system: 'System',
                user: 'User',
                agent: 'Assistant'
            };
            return `${roleMap[msg.type]}: ${msg.message}`;
        })
        .join('\n\n');
};

const db = useDrizzle();

export const startGeneration = async (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): Promise<void> => {
    const pending = pendingGenerations.get(generationId);

    if (!pending) {
        console.error(`Generation ${generationId} not found in pending generations`);
        return;
    }

    clearTimeout(pending.timeout);
    pendingGenerations.delete(generationId);

    const { userId, topicId, messages } = pending;
    const prompt = buildPrompt(messages);

    const generation: ActiveGeneration = {
        userId,
        topicId,
        messages,
        content: '',
        clients: new Set([controller]),
        complete: false
    };

    await db.insert(generations).values({
        id: generationId,
        userId,
        topicId,
        messageId: null
    });

    activeGenerations.set(generationId, generation);

    try {
        const dummyResponse = generateDummyResponse(prompt, messages);
        const tokens = dummyResponse.split(' ');

        for (const token of tokens) {
            await new Promise(resolve => setTimeout(resolve, 50));

            generation.content += token + ' ';

            sendToClients(generation, {
                type: 'token',
                data: token + ' '
            });
        }

        const [message] = await db.insert(messages_drizzle).values({
            topicId: generation.topicId,
            userId: generation.userId,
            isUser: false,
            content: generation.content.trim(),
            model: 'dummy-model-v1',
            tokensGenerated: tokens.length,
            tokensUsedThinking: 0
        }).returning();

        generation.complete = true;

        sendToClients(generation, {
            type: 'complete',
            data: message
        });

        if (generation.clients.size === 0) {
            activeGenerations.delete(generationId);
        } else {
            generation.clients.forEach((client) => {
                try {
                    client.close();
                } catch {
                }
            });
            activeGenerations.delete(generationId);
        }

        await db.update(generations).set({
            messageId: message.id
        }).where(eq(generations.id, generationId));

    } catch (error) {
        console.error('Generation failed:', error);

        sendToClients(generation, {
            type: 'error',
            data: error instanceof Error ? error.message : 'Unknown error'
        });

        await db.delete(generations).where(eq(generations.id, generationId));

        activeGenerations.delete(generationId);
    }
};

export const registerPendingGeneration = (userId: string, generationId: string, topicId: string, messages: ChatMessage[]): void => {
    const timeout = setTimeout(() => {
        const gen = pendingGenerations.get(generationId)
        if (gen) gen.expired = true;
        console.log(`Generation ${generationId} expired - no client connected within 60 seconds`);
    }, 60000);

    pendingGenerations.set(generationId, {
        generationId,
        userId,
        topicId,
        messages,
        timeout,
        expired: false
    });

    console.log(`Registered pending generation ${generationId}, waiting for client connection...`);
};

const generateDummyResponse = (prompt: string, messages: ChatMessage[]): string => {
    const responses = [
        "This is a simulated response to your prompt. In a real implementation, this would be generated by an AI model like GPT-4 or Claude."
        + " I'm processing your message about: " + prompt.substring(0, 50) + "... "
        + "This dummy generation demonstrates the streaming and background save functionality.",

        "I understand your query. This is a placeholder response that simulates AI-generated content."
        + " The system will continue generating this response even if you close the tab, and it will"
        + " automatically save to the database when complete.",

        "Here's a simulated AI response. This demonstrates two key features:"
        + " 1) The generation continues in the background even if you disconnect,"
        + " 2) The complete response is automatically saved to the database without requiring"
        + " a separate update request from the client."
    ];

    const lastUserMessage = messages[messages.length - 1]?.message.toLowerCase() || '';

    if (lastUserMessage.includes('hello') || lastUserMessage.includes('hi')) {
        return "Hello! I'm a dummy AI assistant. This is a simulated response to your greeting."
            + " In production, this would be replaced with actual AI-generated content from an LLM provider.";
    }

    return responses[Math.floor(Math.random() * responses.length)];
};