veridian/app/composables/useChat.ts

import { nanoid } from 'nanoid';
import * as schema from '~~/drizzle/schema';
import { Err, Ok, type Result } from '~~/types/result';
import { buildMessageTree } from '~~/utils/message';

export type BaseMessage = {
    content: string;
    fileIds: string[];
}

export type ToolCall = typeof schema.toolCalls.$inferSelect

export type MessagePart = typeof schema.messageParts.$inferSelect & { toolCall: ToolCall | null }

export type MessageEntity = typeof schema.messages.$inferSelect & { parts: MessagePart[] | undefined } & { generation: typeof schema.generations.$inferSelect | null } & { attachments: (typeof schema.attachments.$inferSelect & { file: typeof schema.files.$inferSelect })[] }

export type Message = MessageEntity & { children: (MessageEntity | undefined)[] }

export enum ChatErrorType {
    NoModel = 0,
    NoProvider,
    NoAgent,
    NoUser,
    NoTopic,
    DatabaseOperationFailed,
    FailedToDecryptProviderApiKey,
    GenerationFailed,
    MarshallFailed,
    NoProviderApiKey,
    NoMessage,
    Unimplemented,
}

type Event = { type: string; payload: any; timestamp: number };

const TARGET_UPDATES_PER_SECOND = 24;

export const useChat = async (topicId: MaybeRef<string>, connect = true) => {
    let data = useState<Topic & { messages: Message[] } | undefined>('useChat:data', () => undefined);

    let sse: EventSource | undefined;

    const textDeltaBuffer: Map<string, Map<string, string>> = new Map();
    let flushTimeout: ReturnType<typeof setTimeout> | undefined;
    let flushCbs: Set<() => void> = new Set();

    const flushTextDeltas = () => {
        if (!data.value) return;

        for (const [messageId, parts] of textDeltaBuffer) {
            const msg = data.value.messages.find(m => m.id === messageId);
            if (!msg) continue;

            for (const [partId, content] of parts) {
                const part = msg.parts?.find(p => p.id === partId);
                if (part) {
                    part.content += content;
                }
            }
        }

        textDeltaBuffer.clear();
        flushTimeout = undefined;

        for (const cb of flushCbs) {
            cb();
        }
        flushCbs.clear();
    };

    const nextFlush = (cb: () => void) => {
        if (flushTimeout) {
            flushCbs.add(cb);
            return;
        }

        cb();
    }

    const scheduleFlush = () => {
        if (flushTimeout) return;
        flushTimeout = setTimeout(flushTextDeltas, 1000 / TARGET_UPDATES_PER_SECOND);
    };

    const processEvent = (event: Event) => {
        if (!data.value) return;

        const { type, payload } = event;

        switch (type) {
            case 'MESSAGE_CREATED': {
                // Push the new message if it doesn't exist (prevents duplicates from HTTP vs SSE)
                if (!data.value.messages.find(m => m.id === payload.id)) {
                    data.value.messages.push({ ...payload, parts: [] });
                } else {
                    // replace the existing data with the new data
                    const index = data.value.messages.findIndex(m => m.id === payload.id);
                    if (index !== -1) {
                        data.value.messages[index] = { ...payload, parts: [] };
                    }
                }
                break;
            }

            case 'MESSAGE_UPDATED': {
                const msgIndex = data.value.messages.findIndex(m => m.id === payload.id);
                if (msgIndex !== -1) {
                    data.value.messages[msgIndex] = { ...data.value.messages[msgIndex], ...payload };
                }
                break;
            }

            case 'MESSAGE_DELETED': {
                data.value.messages = data.value.messages.filter(m => m.id !== payload.id);
                break;
            }

            case 'text-start': {
                const msg = data.value.messages.find(m => m.id === payload.messageId);
                if (msg) {
                    if (!msg.parts) msg.parts = [];
                    msg.parts.push(payload.part);
                }
                break;
            }

            case 'text-delta': {
                const msg = data.value.messages.find(m => m.id === payload.messageId);
                if (!msg) break;

                if (!textDeltaBuffer.has(payload.messageId)) {
                    textDeltaBuffer.set(payload.messageId, new Map());
                }
                const messageBuffer = textDeltaBuffer.get(payload.messageId)!;
                const existing = messageBuffer.get(payload.partId) || '';
                messageBuffer.set(payload.partId, existing + payload.content);

                scheduleFlush();
                break;
            }

            case 'text-end': {
                const msg = data.value.messages.find(m => m.id === payload.messageId);
                const part = msg?.parts?.find(p => p.id === payload.partId);
                if (part) {
                    nextFlush(() => {
                        part.content = payload.content;
                        part.finished = true;
                    });
                }
                break;
            }

            case 'tool-call-start': {
                const msg = data.value.messages.find(m => m.id === payload.messageId);
                if (msg) {
                    if (!msg.parts) msg.parts = [];
                    msg.parts.push(payload.part);
                }
                break;
            }

            case 'tool-call-delta': {
                const msg = data.value.messages.find(m => m.id === payload.messageId);
                const part = msg?.parts?.find(p => p.toolCallId === payload.toolCallId);
                if (part?.toolCall) {
                    if (payload.input) part.toolCall.input = payload.input;
                    if (payload.output) part.toolCall.output = payload.output;
                    if (payload.error) part.toolCall.error = payload.error;
                    if (payload.status) part.toolCall.status = payload.status;
                }
                break;
            }

            case 'generation-complete': {
                const msg = data.value.messages.find(m => m.id === payload.messageId);
                if (msg && msg.generation) {
                    msg.generation.status = 'completed';
                    msg.generation.tokens = payload.tokens;
                }
                break;
            }

            case 'generation-failed': {
                const msg = data.value.messages.find(m => m.id === payload.messageId);
                if (msg && msg.generation) {
                    msg.generation.status = 'failed';
                    msg.generation.error = payload.error;
                }
                break;
            }

            case 'topic_updated': {
                data.value = { ...data.value, ...payload };
                break;
            }
        }
    }

    const connectSSE = async () => {
        const { promise, resolve } = Promise.withResolvers<void>();

        const id = unref(topicId);
        const lastUpdate = data.value?.messages.at(-1)?.updatedAt || '0';
        const msgCount = data.value?.messages.length || 0;

        if (sse && sse.readyState !== EventSource.CLOSED) {
            if (!(new URL(sse.url).pathname.startsWith(`/api/topic/${id}`))) {
                sse.close();
            } else {
                return;
            }
        }

        sse = new EventSource(`/api/topic/${id}?lastUpdate=${lastUpdate}&count=${msgCount}`);

        sse.addEventListener("error", (e) => {
            console.error('sse error', e);
            setTimeout(connectSSE, 1500);
        })

        sse.addEventListener("message", (e) => {
            const raw = JSON.parse(e.data);
            const events: Event[] = Array.isArray(raw) ? raw : [raw];

            for (const { type, payload, timestamp } of events) {
                if (!type || !payload) continue;

                if (type === 'initial_state') {
                    data.value = payload;
                    resolve()
                    continue;
                }

                processEvent({ type, payload, timestamp });
            }
        });

        await promise;
    }


    if (connect) {
        onBeforeUnmount(() => {
            if (flushTimeout) {
                clearTimeout(flushTimeout);
                flushTextDeltas();
            }
            if (sse) {
                console.log('[SSE] Closing connection');
                sse.close();
                sse = undefined;
            }
        });

        if (import.meta.server) {
            const id = unref(topicId);
            const { data: ssrData } = await useFetch<Topic & { messages: Message[] }>(`/api/topic/${id}`);
            data.value = ssrData.value;
        } else {
            // if (data.value === undefined) {
            //     const id = unref(topicId);
            //     const { data: ssrData } = await useFetch<Topic & { messages: Message[] }>(`/api/topic/${id}`);
            //     data.value = ssrData.value;
            // }

            await connectSSE();
        }
    }

    const topic = computed(() => ({
        ...data.value,
        messages: buildMessageTree(data.value?.messages || [])
    }));

    const getProviderAPIKey = async (provider: Provider): Promise<Result<string | undefined, ChatErrorType>> => {
        if (provider.config.apiKey === undefined) {
            return Ok(undefined);
        }

        try {
            const key = await crypto.subtle.importKey(
                "jwk",
                JSON.parse(window.localStorage.getItem("encryptionKey")!),
                "AES-GCM",
                false,
                ["encrypt", "decrypt"]
            )

            return Ok(await decrypt(
                key,
                base64ToUint8Array(provider.config.apiKey)
            ));
        } catch (error) {
            console.error('Failed to decrypt provider API key:', error);
            return Err(ChatErrorType.FailedToDecryptProviderApiKey);
        }
    }

    const sendMessage = async (
        baseMessage: BaseMessage,
        onRequest?: () => void,
    ): Promise<Result<void, ChatErrorType>> => {
        const { user } = useAuth();
        if (!user.value) return Err(ChatErrorType.NoUser);

        try {
            const message = {
                id: nanoid(),
                role: 'user',
                content: baseMessage.content,
                fileIds: baseMessage.fileIds,
            }

            await $fetch(`/api/topic/${unref(topicId)}/message`, {
                method: 'POST',
                body: {
                    message
                },
                onRequest() {
                    if (data.value) {
                        data.value!.messages.push({
                            topicId: unref(topicId),
                            userId: user.value!.id,
                            // TODO
                            attachments: [],
                            parts: undefined,
                            generation: null,
                            parentMessageId: null,
                            generationId: null,
                            focusedIndex: null,
                            deleted: null,
                            createdAt: new Date(),
                            updatedAt: new Date(),
                            children: [],
                            ...message
                        } as Message);
                    }
                    onRequest?.();
                },
                onResponseError() {
                    data.value!.messages = data.value!.messages.filter(m => m.id !== message.id);
                },
            });
            return Ok(undefined);
        } catch (error) {
            console.error('Failed to send message:', error);
            return Err(ChatErrorType.GenerationFailed);
        }
    };

    const startGeneration = async (model: ModelWithProvider, parentMessageId?: string, agentIdOverride?: string) => {
        const providerApiKeyRes = await getProviderAPIKey(model.provider);
        if (providerApiKeyRes.ok === false) {
            return providerApiKeyRes;
        }
        const providerApiKey = providerApiKeyRes.data;

        const { agents } = await useAgents();

        const agent = agents.value.find(agent => agent.id === (agentIdOverride || topic.value?.agentId));
        if (!agent) {
            return Err(ChatErrorType.NoAgent);
        }

        const { settings } = await useUserSettings();

        let rerank = undefined;

        if (agent.config?.search?.rerank && settings.value.systemAssistants.rerank.enabled) {
            const rerankModel = await useModels().then(m => m.allModels.value.find(m => m.id === settings.value.systemAssistants.rerank.modelId));
            if (rerankModel) {
                const rerankProviderApiKeyRes = await getProviderAPIKey(rerankModel.provider);
                if (rerankProviderApiKeyRes.ok === false) {
                    return rerankProviderApiKeyRes;
                }
                const rerankProviderApiKey = rerankProviderApiKeyRes.data;
                rerank = {
                    modelId: rerankModel.id,
                    providerApiKey: rerankProviderApiKey,
                };
            }
        }

        await $fetch(`/api/topic/${unref(topicId)}/chat`, {
            method: 'POST',
            body: {
                modelId: model.id,
                providerApiKey,
                rerank,
                parentMessageId,
            },
        });
    }

    /**
     * 
     * @param messageId The ID of the message we wish to regenerate *for*, this
     * can be either a user's message or an agent's message, and we will find 
     * the message that should be regenerated automatically
     * @param topic 
     * @param topicMessages messages in the current topic, if this does not 
     * include an agent message after the user's message we select, the new
     * message will not be a child of the previous message. However, if this
     * array contains an agent message that is to be regenerated, the new
     * message will have that message's id as its parent. The message
     * reference by messageId *must* be in this array
     * @param agent 
     * @param provider 
     * @param model 
     */
    const regenerateMessage = async (
        messageId: string,
        topicMessages: MessageEntity[],
        model: ModelWithProvider
    ): Promise<Result<void, ChatErrorType>> => {
        const targetMessage = topicMessages.find(message => message.id === messageId);
        if (!targetMessage) {
            return Err(ChatErrorType.NoMessage);
        }
        let targetMessageIndex = topicMessages.indexOf(targetMessage);

        let parentMessageId = undefined;
        let focusedMessages: MessageEntity[] | undefined;
        if (targetMessage.role === 'user') {
            // we need to find the next agent message
            while (targetMessageIndex < topicMessages.length) {
                const currentMessage = topicMessages[targetMessageIndex];
                if (!currentMessage) break;

                if (currentMessage.role === 'assistant') {
                    parentMessageId = currentMessage.parentMessageId || currentMessage.id;
                    focusedMessages = topicMessages.slice(0, targetMessageIndex).filter(message => message.id !== currentMessage.id);
                    break;
                }
                targetMessageIndex++;
            }
        } else {
            parentMessageId = targetMessage.parentMessageId || topicMessages[targetMessageIndex]!.id;
            focusedMessages = topicMessages.slice(0, targetMessageIndex).filter(message => message.id !== messageId);
        }

        if (focusedMessages === undefined) {
            focusedMessages = topicMessages;
        }

        startGeneration(model, parentMessageId);

        return Ok(undefined);
    }

    const patchMessageLocally = (id: string, updates: Partial<Message>) => {
        if (!data.value) return;

        data.value.messages = data.value.messages.map(m =>
            m.id === id ? { ...m, ...updates } : m
        );
    }

    return {
        topic,
        sendMessage,
        startGeneration,
        regenerateMessage,
        patchMessageLocally
    };
}