// *****************************************************************************
// Copyright (C) 2024 TypeFox GmbH.
//
// This program and the accompanying materials are made available under the
// terms of the Eclipse Public License v. 2.0 which is available at
// http://www.eclipse.org/legal/epl-2.0.
//
// This Source Code may also be made available under the following Secondary
// Licenses when the conditions for such availability set forth in the Eclipse
// Public License v. 2.0 are satisfied: GNU General Public License, version 2
// with the GNU Classpath Exception which is available at
// https://www.gnu.org/software/classpath/license.html.
//
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
// *****************************************************************************

import {
    LanguageModel,
    LanguageModelParsedResponse,
    LanguageModelMessage,
    LanguageModelResponse,
    LanguageModelStreamResponse,
    LanguageModelStreamResponsePart,
    ReasoningSettings,
    ReasoningSupport,
    ToolCall,
    ToolRequest,
    ToolRequestParametersProperties,
    ImageContent,
    LanguageModelRequest,
    LanguageModelStatus,
    LanguageModelTextResponse,
    UserRequest
} from '@theia/ai-core';
import { CancellationToken } from '@theia/core';
import { ChatRequest, Message, Ollama, Options, Tool, ToolCall as OllamaToolCall } from 'ollama';
import { createProxyFetch } from '@theia/ai-core/lib/node';
import { ollamaThinkParamFor } from './ollama-reasoning';

export const OllamaModelIdentifier = Symbol('OllamaModelIdentifier');

export class OllamaModel implements LanguageModel {

    protected readonly DEFAULT_REQUEST_SETTINGS: Partial<Omit<ChatRequest, 'stream' | 'model'>> = {
        keep_alive: '15m',
        // options see: https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values
        options: {}
    };

    readonly providerId = 'ollama';
    readonly vendor: string = 'Ollama';

    /**
     * @param id the unique id for this language model. It will be used to identify the model in the UI.
     * @param model the unique model name as used in the Ollama environment.
     * @param hostProvider a function to provide the host URL for the Ollama server.
     */
    constructor(
        public readonly id: string,
        protected readonly model: string,
        public status: LanguageModelStatus,
        protected host: () => string | undefined,
        public proxy?: string,
        public reasoningSupport?: ReasoningSupport
    ) { }

    async request(request: UserRequest, cancellationToken?: CancellationToken): Promise<LanguageModelResponse> {
        const settings = this.getSettings(request);
        const ollama = this.initializeOllama();
        const stream = !(request.settings?.stream === false); // true by default, false only if explicitly specified
        const ollamaRequest: ExtendedChatRequest = {
            model: this.model,
            ...this.DEFAULT_REQUEST_SETTINGS,
            ...settings,
            messages: this.mergeConsecutiveAssistantMessages(
                request.messages.map(m => this.toOllamaMessage(m)).filter((m): m is Message => m !== undefined)
            ),
            tools: request.tools?.map(t => this.toOllamaTool(t)),
            stream
        };
        const structured = request.response_format?.type === 'json_schema';
        return this.dispatchRequest(ollama, ollamaRequest, structured, cancellationToken, request.reasoning);
    }

    /**
     * Retrieves the settings for the chat request, merging the request-specific settings with the default settings.
     * @param request The language model request containing specific settings.
     * @returns A partial ChatRequest object containing the merged settings.
     */
    protected getSettings(request: LanguageModelRequest): Partial<ChatRequest> {
        const settings = request.settings ?? {};
        return {
            options: settings as Partial<Options>
        };
    }

    protected async dispatchRequest(
        ollama: Ollama,
        ollamaRequest: ExtendedChatRequest,
        structured: boolean,
        cancellation?: CancellationToken,
        reasoning?: ReasoningSettings
    ): Promise<LanguageModelResponse> {

        // Handle structured output request
        if (structured) {
            return this.handleStructuredOutputRequest(ollama, ollamaRequest);
        }

        if (isNonStreaming(ollamaRequest)) {
            // handle non-streaming request
            return this.handleNonStreamingRequest(ollama, ollamaRequest, cancellation, reasoning);
        }

        // handle streaming request
        return this.handleStreamingRequest(ollama, ollamaRequest, cancellation, reasoning);
    }

    protected async handleStreamingRequest(
        ollama: Ollama,
        chatRequest: ExtendedChatRequest,
        cancellation?: CancellationToken,
        reasoning?: ReasoningSettings
    ): Promise<LanguageModelStreamResponse> {
        const supportsThinking = await this.checkThinkingSupport(ollama, chatRequest.model);
        const thinkParam = supportsThinking ? this.getThinkingParameter(reasoning, chatRequest.model) : false;
        const responseStream = await ollama.chat({
            ...chatRequest,
            stream: true,
            think: thinkParam
        });

        cancellation?.onCancellationRequested(() => {
            responseStream.abort();
        });

        const that = this;

        const asyncIterator = {
            async *[Symbol.asyncIterator](): AsyncIterator<LanguageModelStreamResponsePart> {
                // Process the response stream and collect thinking, content messages, and tool calls.
                // Tool calls are handled when the response stream is done.
                const toolCalls: OllamaToolCall[] = [];
                let currentContent = '';
                let currentThought = '';

                // Ollama does not have ids, so we use the most recent chunk.created_at timestamp as repalcement
                let lastUpdated: Date = new Date();

                try {
                    for await (const chunk of responseStream) {
                        lastUpdated = chunk.created_at;

                        const thought = chunk.message.thinking;
                        if (thought) {
                            currentThought += thought;
                            yield { thought, signature: '' };
                        }
                        const textContent = chunk.message.content;
                        if (textContent) {
                            currentContent += textContent;
                            yield { content: textContent };
                        }

                        if (chunk.message.tool_calls && chunk.message.tool_calls.length > 0) {
                            toolCalls.push(...chunk.message.tool_calls);
                        }

                        if (chunk.done) {
                            if (chunk.prompt_eval_count !== undefined && chunk.eval_count !== undefined) {
                                yield { input_tokens: chunk.prompt_eval_count, output_tokens: chunk.eval_count };
                            }
                            if (chunk.done_reason && chunk.done_reason !== 'stop') {
                                throw new Error('Ollama stopped unexpectedly. Reason: ' + chunk.done_reason);
                            }
                        }
                    }

                    if (toolCalls && toolCalls.length > 0) {
                        chatRequest.messages.push({
                            role: 'assistant',
                            content: currentContent,
                            thinking: currentThought,
                            tool_calls: toolCalls
                        });

                        // Create tool call message parts and yield them.
                        // This is required because when calling a tool, Theia AI expects the corresponding message part to exist.
                        const toolCallsForResponse = that.createToolCalls(toolCalls, lastUpdated);
                        yield { tool_calls: toolCallsForResponse };

                        // Now handle the tool calls
                        const processedToolCallsForResponse = await that.processToolCalls(toolCallsForResponse, chatRequest);
                        yield { tool_calls: processedToolCallsForResponse };

                        // Continue the conversation with tool results
                        const continuedResponse = await that.handleStreamingRequest(
                            ollama,
                            chatRequest,
                            cancellation,
                            reasoning
                        );

                        // Stream the continued response
                        for await (const nestedEvent of continuedResponse.stream) {
                            yield nestedEvent;
                        }
                    }
                } catch (error) {
                    console.error('Error in Ollama streaming:', error.message);
                    throw error;
                }
            }
        };

        return { stream: asyncIterator };
    }

    /**
     * Check if the Ollama server supports thinking.
     *
     * Use the Ollama 'show' request to get information about the model, so we can check the capabilities for the 'thinking' capability.
     *
     * @param ollama The Ollama client instance.
     * @param model The name of the Ollama model.
     * @returns A boolean indicating whether the Ollama model supports thinking.
     */
    protected async checkThinkingSupport(ollama: Ollama, model: string): Promise<boolean> {
        const result = await ollama.show({ model });
        return result?.capabilities?.includes('thinking') || false;
    }

    protected getThinkingParameter(reasoning: ReasoningSettings | undefined, model: string): boolean | 'low' | 'medium' | 'high' {
        return ollamaThinkParamFor(reasoning?.level, this.requiresEffortLevel(model));
    }

    /** Checks if the model requires effort levels instead of a boolean for `think`. */
    protected requiresEffortLevel(model: string): boolean {
        return model.toLowerCase().includes('gpt-oss');
    }

    protected async handleStructuredOutputRequest(ollama: Ollama, chatRequest: ChatRequest): Promise<LanguageModelParsedResponse> {
        const response = await ollama.chat({
            ...chatRequest,
            format: 'json',
            stream: false,
        });
        try {
            const result: LanguageModelParsedResponse = {
                content: response.message.content,
                parsed: JSON.parse(response.message.content)
            };
            if (response.prompt_eval_count !== undefined && response.eval_count !== undefined) {
                result.usage = { input_tokens: response.prompt_eval_count, output_tokens: response.eval_count };
            }
            return result;
        } catch (error) {
            // TODO use ILogger
            console.log('Failed to parse structured response from the language model.', error);
            const result: LanguageModelParsedResponse = {
                content: response.message.content,
                parsed: {}
            };
            if (response.prompt_eval_count !== undefined && response.eval_count !== undefined) {
                result.usage = { input_tokens: response.prompt_eval_count, output_tokens: response.eval_count };
            }
            return result;
        }
    }

    protected async handleNonStreamingRequest(
        ollama: Ollama,
        chatRequest: ExtendedNonStreamingChatRequest,
        cancellation?: CancellationToken,
        reasoning?: ReasoningSettings
    ): Promise<LanguageModelResponse> {
        try {
            // even though we have a non-streaming request, we still use the streaming version for two reasons:
            // 1. we can abort the stream if the request is cancelled instead of having to wait for the entire response
            // 2. we can use think: true so the Ollama API separates thinking from content and we can filter out the thoughts in the response
            const supportsThinking = await this.checkThinkingSupport(ollama, chatRequest.model);
            const thinkParam = supportsThinking ? this.getThinkingParameter(reasoning, chatRequest.model) : false;
            const responseStream = await ollama.chat({ ...chatRequest, stream: true, think: thinkParam });
            cancellation?.onCancellationRequested(() => {
                responseStream.abort();
            });

            const toolCalls: OllamaToolCall[] = [];
            let content = '';
            let lastUpdated: Date = new Date();
            let inputTokenCount: number | undefined;
            let outputTokenCount: number | undefined;

            // process the response stream
            for await (const chunk of responseStream) {
                // if the response contains content, append it to the result
                const textContent = chunk.message.content;
                if (textContent) {
                    content += textContent;
                }

                // record requested tool calls so we can process them later
                if (chunk.message.tool_calls && chunk.message.tool_calls.length > 0) {
                    toolCalls.push(...chunk.message.tool_calls);
                }

                // if the response is done, capture token usage and check the done reason
                if (chunk.done) {
                    lastUpdated = chunk.created_at;
                    inputTokenCount = chunk.prompt_eval_count;
                    outputTokenCount = chunk.eval_count;
                    if (chunk.done_reason && chunk.done_reason !== 'stop') {
                        throw new Error('Ollama stopped unexpectedly. Reason: ' + chunk.done_reason);
                    }
                }
            }

            // process any tool calls by adding all of them to the messages of the conversation
            if (toolCalls && toolCalls.length > 0) {
                chatRequest.messages.push({
                    role: 'assistant',
                    content: content,
                    tool_calls: toolCalls
                });

                const preparedToolCalls = this.createToolCalls(toolCalls, lastUpdated);
                await this.processToolCalls(preparedToolCalls, chatRequest);
                if (cancellation?.isCancellationRequested) {
                    return { text: '' };
                }

                // recurse to get the final response content (the intermediate content remains hidden, it is only part of the conversation)
                return this.handleNonStreamingRequest(ollama, chatRequest, cancellation, reasoning);
            }

            // if no tool calls are necessary, return the final response content
            const result: LanguageModelTextResponse = { text: content };
            if (inputTokenCount !== undefined && outputTokenCount !== undefined) {
                result.usage = { input_tokens: inputTokenCount, output_tokens: outputTokenCount };
            }
            return result;
        } catch (error) {
            console.error('Error in ollama call:', error.message);
            throw error;
        }
    }

    private createToolCalls(toolCalls: OllamaToolCall[], lastUpdated: Date): ToolCall[] {
        const toolCallsForResponse: ToolCall[] = [];
        for (const [idx, toolCall] of toolCalls.entries()) {
            const args = JSON.stringify(toolCall.function?.arguments);
            toolCallsForResponse.push({
                id: `ollama_${lastUpdated}_${idx}`,
                function: {
                    name: toolCall.function.name,
                    arguments: args
                },
                finished: false
            });
        }
        return toolCallsForResponse;
    }

    private async processToolCalls(toolCalls: ToolCall[], chatRequest: ExtendedChatRequest): Promise<ToolCall[]> {
        const tools: ToolWithHandler[] = chatRequest.tools ?? [];
        const toolCallsForResponse: ToolCall[] = [];

        for (const call of toolCalls) {
            const functionToCall = tools.find(tool => tool.function.name === call.function!.name);
            let funcResult: string;

            if (functionToCall) {
                const rawResult = await functionToCall.handler(call.function!.arguments!);
                funcResult = typeof rawResult === 'string' ? rawResult : JSON.stringify(rawResult);
            } else {
                funcResult = 'error: Tool not found';
            }

            chatRequest.messages.push({
                role: 'tool',
                content: `Tool call ${call.function!.name} returned: ${String(funcResult)}`,
            });

            // update tool call message
            toolCallsForResponse.push({
                ...call,
                result: String(funcResult),
                finished: true
            });
        }
        return toolCallsForResponse;
    }

    protected initializeOllama(): Ollama {
        const host = this.host();
        if (!host) {
            throw new Error('Please provide OLLAMA_HOST in preferences or via environment variable');
        }
        return new Ollama({ host: host, fetch: createProxyFetch(this.proxy) });
    }

    protected toOllamaTool(tool: ToolRequest): ToolWithHandler {
        const transform = (props: ToolRequestParametersProperties | undefined) => {
            if (!props) {
                return undefined;
            }

            const result: Record<string, { type: string, description: string, enum?: string[] }> = {};
            for (const [key, prop] of Object.entries(props)) {
                const type = prop.type;
                if (type) {
                    const description = typeof prop.description == 'string' ? prop.description : '';
                    result[key] = {
                        type: type,
                        description: description
                    };
                } else {
                    // TODO: Should handle anyOf, but this is not supported by the Ollama type yet
                }
            }
            return result;
        };
        return {
            type: 'function',
            function: {
                name: tool.name,
                description: tool.description ?? 'Tool named ' + tool.name,
                parameters: {
                    type: tool.parameters?.type ?? 'object',
                    required: tool.parameters?.required ?? [],
                    properties: transform(tool.parameters?.properties) ?? {}
                },
            },
            handler: tool.handler
        };
    }

    protected toOllamaMessage(message: LanguageModelMessage): Message | undefined {
        const result: Message = {
            role: this.toOllamaMessageRole(message),
            content: ''
        };

        if (LanguageModelMessage.isTextMessage(message) && message.text.length > 0) {
            result.content = message.text;
        } else if (LanguageModelMessage.isToolUseMessage(message)) {
            result.tool_calls = [{ function: { name: message.name, arguments: message.input as Record<string, unknown> } }];
        } else if (LanguageModelMessage.isToolResultMessage(message)) {
            result.content = `Tool call ${message.name} returned: ${message.content}`;
        } else if (LanguageModelMessage.isThinkingMessage(message)) {
            result.thinking = message.thinking;
        } else if (LanguageModelMessage.isImageMessage(message) && ImageContent.isBase64(message.image)) {
            result.images = [message.image.base64data];
        } else if (LanguageModelMessage.isCompactionMessage(message)) {
            // Ollama has no server-side compaction; the opaque marker carries no representable content and is dropped.
            return undefined;
        } else {
            console.log(`Unknown message type encountered when converting message to Ollama format: ${JSON.stringify(message)}. Ignoring message.`);
            return undefined;
        }

        return result;
    }

    protected mergeConsecutiveAssistantMessages(messages: Message[]): Message[] {
        const result: Message[] = [];
        for (const message of messages) {
            const previous = result[result.length - 1];
            if (previous?.role === 'assistant' && message.role === 'assistant') {
                const merged: Message = { ...previous, role: 'assistant' };

                const previousContent = previous.content;
                const nextContent = message.content;
                if (previousContent && nextContent) {
                    merged.content = `${previousContent}\n${nextContent}`;
                } else if (nextContent) {
                    merged.content = nextContent;
                } else if (previousContent) {
                    merged.content = previousContent;
                }

                const previousThinking = previous.thinking;
                const nextThinking = message.thinking;
                if (previousThinking && nextThinking) {
                    merged.thinking = `${previousThinking}\n${nextThinking}`;
                } else if (nextThinking) {
                    merged.thinking = nextThinking;
                } else if (previousThinking) {
                    merged.thinking = previousThinking;
                }

                const toolCalls = [...(previous.tool_calls ?? []), ...(message.tool_calls ?? [])];
                if (toolCalls.length > 0) {
                    merged.tool_calls = toolCalls;
                }

                result[result.length - 1] = merged;
            } else {
                result.push(message);
            }
        }
        return result;
    }

    protected toOllamaMessageRole(message: LanguageModelMessage): string {
        if (LanguageModelMessage.isToolResultMessage(message)) {
            return 'tool';
        }
        const actor = message.actor;
        if (actor === 'ai') {
            return 'assistant';
        }
        if (actor === 'user') {
            return 'user';
        }
        if (actor === 'system') {
            return 'system';
        }
        console.log(`Unknown actor encountered when converting message to Ollama format: ${actor}. Falling back to 'user'.`);
        return 'user'; // default fallback
    }
}

/**
 * Extended Tool containing a handler
 * @see Tool
 */
type ToolWithHandler = Tool & { handler: (arg_string: string) => Promise<unknown> };

/**
 * Extended chat request with mandatory messages and ToolWithHandler tools
 *
 * @see ChatRequest
 * @see ToolWithHandler
 */
type ExtendedChatRequest = ChatRequest & {
    messages: Message[]
    tools?: ToolWithHandler[]
};

type ExtendedNonStreamingChatRequest = ExtendedChatRequest & { stream: false };

function isNonStreaming(request: ExtendedChatRequest): request is ExtendedNonStreamingChatRequest {
    return !request.stream;
}
