import type schema from "#triplit/schema";
import type { Entity } from "@triplit/client";
import type { ModelMessage } from "ai";
import { decrypt, base64ToUint8Array } from "~/utils/crypto";

export type Message = Entity<typeof schema, 'messages'> & { parts: (Entity<typeof schema, 'message_parts'> & { toolCall: Entity<typeof schema, 'tool_calls'> | null } | undefined)[] };

export const useChat = (agentId: string) => {
    const triplit = useTriplitClient();

    const createTopic = async () => {
        const { user } = useAuth();
        if (!user.value) {
            console.error('No user');
            return;
        }

        const newTopic = await triplit.insert('topics', {
            name: 'New Topic',
            userId: user.value.id,
            agentId,
            createdAt: new Date().toISOString(),
        });

        if ('flush' in triplit) {
            await triplit.flush();
        }

        return newTopic;
    };

    const marshallMessages = (agent: Entity<typeof schema, 'agents'>, messages: Readonly<Message[]>) => {
        const marshalledMessages: ModelMessage[] = [];

        if (agent && agent.systemPrompt) {
            marshalledMessages.push({
                role: 'system',
                content: agent.systemPrompt,
            });
        }

        messages.forEach((message) => {
            switch (message.role) {
                case 'user':
                    marshalledMessages.push({
                        role: 'user',
                        // TODO: when we have images or files, this is where we need to handle them
                        content: message.content,
                    });
                    break;
                case 'assistant':
                    message.parts.forEach((part) => {
                        if (!part) throw new Error('Part is undefined');

                        switch (part.type) {
                            case 'text':
                            case 'reasoning': {
                                marshalledMessages.push({
                                    role: 'assistant',
                                    content: part.content,
                                });
                                break;
                            }
                            case 'tool-call': {
                                if (part.toolCall === null) throw new Error('Tool call is null');

                                if (part.toolCall.status === 'pending') {
                                    throw new Error(
                                        'Marshalling tool call that is still pending. This is likely a UI bug if this happens.',
                                    );
                                }

                                let inputValue: string = '';

                                switch (typeof part.toolCall.input!.value) {
                                    case 'string':
                                        inputValue = part.toolCall.input!.value;
                                        break;
                                    case 'object':
                                        inputValue = JSON.stringify(part.toolCall.input!.value, null, 2);
                                        break;
                                }

                                marshalledMessages.push({
                                    role: 'assistant',
                                    content: [
                                        {
                                            type: 'tool-call',
                                            toolCallId: part.toolCall.id,
                                            toolName: part.toolCall.toolName,
                                            input: inputValue,
                                        },
                                    ],
                                    providerOptions: part.providerOptions,
                                });

                                if (part.toolCall.status === 'failed') {
                                    let failureType: 'error-text' | 'error-json';
                                    let failureValue: string;

                                    if (part.toolCall.error === null || part.toolCall.error === undefined) {
                                        failureType = 'error-text';
                                        failureValue = 'An unknown error occurred';
                                    } else {
                                        switch (part.toolCall.error!.type) {
                                            case 'text':
                                                failureType = 'error-text';
                                                failureValue = part.toolCall.error!.value;
                                                break;
                                            case 'json':
                                                failureType = 'error-json';
                                                failureValue = JSON.stringify(part.toolCall.error!.value, null, 2);
                                                break;
                                        }
                                    }

                                    marshalledMessages.push({
                                        role: 'tool',
                                        content: [
                                            {
                                                type: 'tool-result',
                                                toolCallId: part.toolCall.id,
                                                toolName: part.toolCall.toolName,
                                                output: {
                                                    type: failureType,
                                                    value: failureValue,
                                                },
                                            },
                                        ],
                                        providerOptions: part.providerOptions,
                                    });
                                    break;
                                }

                                if (part.toolCall.status === 'completed') {
                                    marshalledMessages.push({
                                        role: 'tool',
                                        content: [
                                            {
                                                type: 'tool-result',
                                                toolCallId: part.toolCall.id,
                                                toolName: part.toolCall.toolName,
                                                output: {
                                                    type: 'json',
                                                    value: JSON.stringify(part.toolCall.output!.value, null, 2),
                                                },
                                            },
                                        ],
                                        providerOptions: part.providerOptions,
                                    });
                                    break;
                                }
                            } break;
                            default:
                                throw new Error(`Unknown part type: ${part.type}`);
                        }
                    });
                    break;
                default:
                    throw new Error(`Unknown message role: ${message.role}`);
            }
        });

        return marshalledMessages;
    };

    const sendMessage = async (
        message: string,
        topic: Entity<typeof schema, 'topics'>,
        topicMessages: Message[],
        agent: Entity<typeof schema, 'agents'>,
        provider: Entity<typeof schema, 'providers'>,
        model: Entity<typeof schema, 'models'>,
    ) => {
        const newMessage = await triplit.insert('messages', {
            topicId: topic.id,
            createdAt: new Date().toISOString(),
            content: message,
            role: 'user',
        }) as Message;

        const messages = marshallMessages(
            agent,
            topicMessages.concat(newMessage)
        );


        let providerApiKey: string | undefined = undefined;
        if (provider.config.apiKey !== undefined) {
            const key = await crypto.subtle.importKey(
                "jwk",
                JSON.parse(window.localStorage.getItem("encryptionKey")!),
                "AES-GCM",
                false,
                ["encrypt", "decrypt"]
            )

            providerApiKey = await decrypt(
                key,
                base64ToUint8Array(provider.config.apiKey)
            );
        }

        return $fetch('/api/chat/generate', {
            method: 'POST',
            body: {
                messages,
                topicId: topic.id,
                model: {
                    providerId: provider.id,
                    modelId: model.id,
                    args: {
                        temperature: 0.7,
                        max_tokens: 100,
                        top_p: 1,
                        frequency_penalty: 0,
                        presence_penalty: 0,
                    },
                },
                providerApiKey: providerApiKey,
            },
            headers: {
                'Content-Type': 'application/json',
            },
        });
    };

    return {
        sendMessage,
        createTopic,
    };
}