import { nanoid } from 'nanoid';
import { toRaw } from 'vue';
import * as schema from '~~/drizzle/schema';
import { Err, Ok, type Result } from '~~/types/result';
import { buildMessageTree } from '~~/utils/message';
import { appendTextDelta, serializeTopicSnapshot } from '~~/utils/topic-sync';

export type BaseMessage = {
    content: string;
    fileIds: string[];
    files?: {
        id: string;
        name: string;
        mimeType: string;
        url: string;
        size?: number;
    }[];
}

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; replay?: boolean };

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;
    let disposed = false;
    let connectionAttempt = 0;
    let reconnectTimeout: ReturnType<typeof setTimeout> | undefined;
    const { promise: connectionReady, resolve: markConnected } = Promise.withResolvers<void>();

    const textDeltaBuffer: Map<string, Map<string, {
        content: string;
        lastUpdatedAt: 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.content;
                    part.lastUpdatedAt = new Date(content.lastUpdatedAt);
                }
            }
        }

        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 resetTextDeltas = () => {
        if (flushTimeout !== undefined) clearTimeout(flushTimeout);
        flushTimeout = undefined;
        textDeltaBuffer.clear();
        flushCbs.clear();
    };

    const scheduleReconnect = (forceSnapshot = false) => {
        sse?.close();
        sse = undefined;
        if (disposed || reconnectTimeout !== undefined) return;
        reconnectTimeout = setTimeout(() => {
            reconnectTimeout = undefined;
            void connectSSE(forceSnapshot);
        }, 1500);
    };

    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: payload.parts ?? [], attachments: payload.attachments ?? [] });
                } else {
                    // replace the existing data with the new data
                    const index = data.value.messages.findIndex(m => m.id === payload.id);
                    if (index !== -1 && !event.replay) {
                        const existing = data.value.messages[index]!;
                        data.value.messages[index] = {
                            ...payload,
                            parts: payload.parts ?? [],
                            attachments: payload.attachments ?? existing.attachments ?? [],
                        };
                    }
                }
                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 = [];
                    if (!msg.parts.some(part => part.id === payload.part.id)) {
                        msg.parts.push(payload.part);
                    }
                }
                break;
            }

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

                if (!textDeltaBuffer.has(payload.messageId)) {
                    textDeltaBuffer.set(payload.messageId, new Map());
                }
                const messageBuffer = textDeltaBuffer.get(payload.messageId)!;
                const existing = messageBuffer.get(payload.partId)?.content ?? part.content ?? '';
                if (typeof payload.offset === 'number' && payload.offset > existing.length) {
                    scheduleReconnect(true);
                    break;
                }
                const content = appendTextDelta(existing, payload.content, payload.offset);
                if (content === existing) break;
                messageBuffer.set(payload.partId, { content, lastUpdatedAt: payload.lastUpdatedAt });

                scheduleFlush();
                break;
            }

            case 'text-delete': {
                const msg = data.value.messages.find(m => m.id === payload.messageId);
                if (msg?.parts) msg.parts = msg.parts.filter(part => part.id !== payload.partId);
                textDeltaBuffer.get(payload.messageId)?.delete(payload.partId);
                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 && part.type !== 'tool-call') {
                    nextFlush(() => {
                        part.content = payload.content;
                        if (payload.lastUpdatedAt) {
                            part.lastUpdatedAt = payload.lastUpdatedAt;
                        }
                        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 = [];
                    }
                    const part = { ...payload.part };
                    if (!part.createdAt) {
                        part.createdAt = new Date().toISOString();
                    }
                    if (!msg.parts.some(existing => existing.id === part.id)) {
                        msg.parts.push(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;
                        if (payload.status === 'completed' || payload.status === 'failed') {
                            part.lastUpdatedAt = payload.lastUpdatedAt || new Date().toISOString();
                            part.finished = true;
                        }
                    }
                }
                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 (forceSnapshot = false) => {
        if (disposed) return;
        const attempt = ++connectionAttempt;
        const id = unref(topicId);
        sse?.close();
        if (flushTimeout !== undefined) {
            clearTimeout(flushTimeout);
            flushTextDeltas();
        }
        let offeredSnapshot: string | undefined;
        const query = new URLSearchParams();
        if (!forceSnapshot && data.value?.id === id && globalThis.crypto?.subtle) {
            try {
                offeredSnapshot = serializeTopicSnapshot(toRaw(data.value));
                const digest = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(offeredSnapshot));
                query.set('snapshot', Array.from(new Uint8Array(digest), byte => byte.toString(16).padStart(2, '0')).join(''));
            } catch {
                // Fingerprinting is an optimization; a full snapshot remains safe.
                offeredSnapshot = undefined;
            }
        }
        if (disposed || attempt !== connectionAttempt) return;

        const source = new EventSource(`/api/topic/${id}?${query}`);
        sse = source;
        source.addEventListener('error', () => {
            if (sse === source) scheduleReconnect();
        });

        source.addEventListener('message', (e) => {
            if (disposed || sse !== source) return;
            const raw = JSON.parse(e.data);
            const events: Event[] = Array.isArray(raw) ? raw : [raw];

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

                if (type === 'synced') {
                    // Optimistic edits can occur while the handshake is in flight.
                    // Only retain the exact state that we offered to the server.
                    if (offeredSnapshot !== serializeTopicSnapshot(toRaw(data.value))) {
                        scheduleReconnect(true);
                        return;
                    }
                    offeredSnapshot = undefined;
                    markConnected();
                    continue;
                }

                if (type === 'initial_state') {
                    resetTextDeltas();
                    data.value = payload;
                    offeredSnapshot = undefined;
                    markConnected();
                    continue;
                }

                processEvent(event);
            }
        });

        await connectionReady;
    };


    if (connect) {
        onBeforeUnmount(() => {
            disposed = true;
            connectionAttempt++;
            if (reconnectTimeout !== undefined) clearTimeout(reconnectTimeout);
            resetTextDeltas();
            markConnected();
            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 {
            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,
        role: 'user' | 'assistant' = 'user',
    ): Promise<Result<void, ChatErrorType>> => {
        const { user } = useAuth();
        if (!user.value) return Err(ChatErrorType.NoUser);

        try {
            const message = {
                id: nanoid(),
                role,
            }

            await $fetch(`/api/topic/${unref(topicId)}/message`, {
                method: 'POST',
                body: {
                    message: {
                        ...message,
                        content: baseMessage.content,
                        fileIds: baseMessage.fileIds,
                    }
                },
                onRequest() {
                    if (data.value) {
                        switch (role) {
                            case 'user': {
                                const optimisticAttachments = baseMessage.files?.map(file => ({
                                    id: nanoid(),
                                    userId: user.value!.id,
                                    topicId: unref(topicId),
                                    messageId: message.id,
                                    fileId: file.id,
                                    createdAt: new Date(),
                                    file: {
                                        id: file.id,
                                        userId: user.value!.id,
                                        name: file.name,
                                        mimeType: file.mimeType,
                                        size: file.size ?? 0,
                                        url: file.url,
                                        createdAt: new Date(),
                                    },
                                })) ?? [];

                                data.value!.messages.push({
                                    topicId: unref(topicId),
                                    userId: user.value!.id,
                                    attachments: optimisticAttachments,
                                    parts: undefined,
                                    generation: null,
                                    parentMessageId: null,
                                    generationId: null,
                                    activeChildId: null,
                                    children: [],
                                    deleted: false,
                                    content: baseMessage.content,
                                    createdAt: new Date(),
                                    updatedAt: new Date(),
                                    ...message
                                });
                                break;
                            }
                            case 'assistant': {
                                data.value!.messages.push({
                                    topicId: unref(topicId),
                                    userId: user.value!.id,
                                    // TODO
                                    attachments: [],
                                    generation: null,
                                    parentMessageId: null,
                                    generationId: null,
                                    activeChildId: null,
                                    children: [],
                                    deleted: false,
                                    content: null,
                                    createdAt: new Date(),
                                    updatedAt: new Date(),
                                    parts: [
                                        {
                                            id: 'bogus-id',
                                            userId: user.value!.id,
                                            topicId: unref(topicId),
                                            messageId: message.id,
                                            type: 'text',
                                            content: baseMessage.content,
                                            finished: true,
                                            createdAt: new Date(),
                                            lastUpdatedAt: new Date(),
                                            providerOptions: null,
                                            toolCallId: null,
                                            toolCall: null,
                                        }
                                    ],
                                    ...message
                                });
                                break;
                            }
                        }
                    }
                    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);
        }

        await $fetch(`/api/topic/${unref(topicId)}/chat`, {
            method: 'POST',
            body: {
                modelId: model.id,
                providerApiKey,
                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
        );
    }

    const deleteMessagesFocusedDescendant = async (messageId: string): Promise<Result<void, Error>> => {
        if (!data.value) {
            return Err(new Error("No data"));
        }

        const msg = data.value.messages.find((m) => {
            return m.id === messageId;
        });
        if (!msg) {
            return Err(new Error("Message not found"));
        }

        const previousMessages = data.value.messages;

        const res = await $fetch(`/api/messages/${messageId}`, {
            method: "DELETE",
            onRequest() {
                if (!data.value) {
                    return;
                }

                if (msg.parentMessageId) {
                    const parent = data.value.messages.find((m) => {
                        return m.id === msg.parentMessageId;
                    });
                    data.value.messages = data.value.messages.filter((m) => {
                        return m.id !== msg.id;
                    });
                    if (parent) {
                        const remainingChildren = data.value.messages.filter((m) => {
                            return m.parentMessageId === parent.id;
                        });
                        if (parent.deleted && remainingChildren.length === 0) {
                            data.value.messages = data.value.messages.filter((m) => {
                                return m.id !== parent.id;
                            });
                        } else if (parent.activeChildId === msg.id) {
                            parent.activeChildId = remainingChildren.length > 0
                                ? (parent.deleted ? remainingChildren[0]!.id : remainingChildren[remainingChildren.length - 1]!.id)
                                : null;
                        }
                    }
                    return;
                }

                const children = data.value.messages.filter((m) => {
                    return m.parentMessageId === msg.id;
                });
                if (children.length === 0) {
                    data.value.messages = data.value.messages.filter((m) => {
                        return m.id !== msg.id;
                    });
                    return;
                }

                msg.deleted = true;
                msg.activeChildId = children[0]?.id ?? null;
            },
            onResponseError() {
                if (data.value) {
                    data.value.messages = previousMessages;
                }
            },
        });

        if (!res.ok) {
            return Err(new Error("Failed to delete message"));
        }

        return Ok(undefined);
    };

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