import type { FilePart, ImagePart, ModelMessage } from "ai";
import { ToolCallType } from "~~/drizzle/schema";
import { Err, Ok, type Result } from "~~/types/result";
import type { Message, MessageEntity } from "~/composables/useChat";
import type { Agent } from "~/composables/useAgents";

export const buildMessageTree = (flatMessages: MessageEntity[]) => {
    const messagesMap = new Map(flatMessages.map(m => [m.id, { ...m, children: [] as MessageEntity[] }]));
    const roots = [] as Message[];

    for (const msg of messagesMap.values()) {
        if (msg.parentMessageId && messagesMap.has(msg.parentMessageId)) {
            messagesMap.get(msg.parentMessageId)!.children.push(msg);
        } else {
            roots.push(msg);
        }
    }
    return roots;
};

export const buildFocusedMessageTree = (messages: Readonly<Message[]>): MessageEntity[] => {
    let focusedMessageTree: MessageEntity[] = [];

    for (const message of messages) {
        if (message.focusedIndex) {
            if (message.focusedIndex === 0) {
                focusedMessageTree.push(message);
                continue;
            }

            const activeChild = message.children[Math.min(Math.max(message.focusedIndex! - 1, 0), message.children.length - 1)]
            if (activeChild) {
                focusedMessageTree.push(activeChild);
            } else {
                focusedMessageTree.push(message);
            }

            continue;
        }

        focusedMessageTree.push(message);
    }

    return focusedMessageTree;
}

export const marshallMessages = (agent: Agent, messages: Readonly<MessageEntity[]>): Result<ModelMessage[], string> => {
    const marshalledMessages: ModelMessage[] = [];

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

    for (const message of messages) {
        switch (message.role) {
            case 'user': {
                const attachments = message.attachments.map(attachment => {
                    if (attachment.file.mimeType.startsWith('image/')) {
                        return {
                            type: 'image',
                            image: attachment.file.url,
                        };
                    }

                    return {
                        type: 'file',
                        data: attachment.file.url,
                        filename: attachment.file.name,
                        mediaType: attachment.file.mimeType,
                    };
                }) as (FilePart | ImagePart)[];

                let messageDate = new Date(message.createdAt);
                marshalledMessages.push({
                    role: 'user',
                    content: [
                        {
                            type: 'text',
                            text: `[${messageDate.toDateString()} ${messageDate.toLocaleTimeString()}]: ${message.content!}`
                        },
                        ...attachments,
                    ],
                });
                break;
            }
            case 'assistant':
                for (const part of (message.parts || [])) {
                    if (!part) return Err('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) return Err('Tool call is null');

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

                            let inputValue: string = '';

                            switch (part.toolCall.input!.type) {
                                case ToolCallType.Text:
                                    inputValue = part.toolCall.input!.value;
                                    break;
                                case ToolCallType.Json:
                                    inputValue = JSON.stringify(part.toolCall.input!.value);
                                    break;
                            }

                            marshalledMessages.push({
                                role: 'assistant',
                                content: [
                                    {
                                        type: 'tool-call',
                                        toolCallId: part.toolCall.id,
                                        toolName: part.toolCall.toolName,
                                        input: inputValue,
                                    },
                                ],
                                providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
                            });

                            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 ToolCallType.Text:
                                            failureType = 'error-text';
                                            failureValue = part.toolCall.error!.value;
                                            break;
                                        case ToolCallType.Json:
                                            failureType = 'error-json';
                                            failureValue = JSON.stringify(part.toolCall.error!.value);
                                            break;
                                    }

                                    failureType = 'error-json';
                                    failureValue = JSON.stringify(part.toolCall.error!.value);
                                }

                                marshalledMessages.push({
                                    role: 'tool',
                                    content: [
                                        {
                                            type: 'tool-result',
                                            toolCallId: part.toolCall.id,
                                            toolName: part.toolCall.toolName,
                                            output: {
                                                type: failureType,
                                                value: failureValue,
                                            },
                                        },
                                    ],
                                    providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
                                });
                                break;
                            }

                            if (part.toolCall.status === 'completed') {
                                let outputType: 'text' | 'json';
                                let outputValue: string;

                                switch (part.toolCall.output!.type) {
                                    case ToolCallType.Text:
                                        outputType = 'text';
                                        outputValue = part.toolCall.output!.value;
                                        break;
                                    case ToolCallType.Json:
                                        outputType = 'json';
                                        outputValue = JSON.stringify(part.toolCall.output!.value);
                                        break;
                                }

                                marshalledMessages.push({
                                    role: 'tool',
                                    content: [
                                        {
                                            type: 'tool-result',
                                            toolCallId: part.toolCall.id,
                                            toolName: part.toolCall.toolName,
                                            output: {
                                                type: outputType,
                                                value: outputValue,
                                            },
                                        },
                                    ],
                                    providerOptions: part.providerOptions ? part.providerOptions as Record<string, any> : undefined,
                                });
                                break;
                            }
                        } break;
                        default:
                            return Err(`Unknown part type: ${part.type}`);
                    }
                }
                break;
            default:
                return Err(`Unknown message role: ${message.role}`);
        }
    }

    return Ok(marshalledMessages);
};