import { ChatCompletion, ChatCompletionChunk, ChatCompletionCreateParams, ChatCompletionMessageParam, ChatCompletionMessageToolCall, CreateEmbeddingResponse, Embedding, EmbeddingCreateParams } from "openai/resources";
import { OpenAILikeInterface } from "../interfaces/OpenAILike";
import { ChatRequest, ChatResponse, Message, Ollama, Options, Tool, ToolCall } from "ollama";
import { v4 } from "uuid";
import { Logger } from "../utils/Logger";
import { ChatCompletionStream } from "../utils/stream";

export class OllamaService implements OpenAILikeInterface {

    host: string
    ollama: Ollama
    constructor(options?: {
        host: string
    }) {
        this.host = options?.host || "http://127.0.0.1:11434"
        // Logger.debug("host", this.host)
        this.ollama = new Ollama({
            host: this.host
        })
    }

    private messagesFormat(messages: ChatCompletionMessageParam[]): Message[] {
        const result: Message[] = []
        messages.forEach(message => {
            if (message.role === "system") {
                result.push({
                    role: "system",
                    content: message.content as string
                })
            } else if (message.role === "user") {
                let msg: Message = {
                    role: "user",
                    content: ""
                }
                if (message.content instanceof Array) {
                    const images: string[] = []
                    message.content.forEach(part => {
                        if (part.type === "text") {
                            msg.content += part.text
                        } else if (part.type === "image_url") {
                            images.push(part.image_url.url)
                        }
                    })
                    if (images.length) {
                        msg.images = images
                    }

                } else {
                    msg.content = message.content
                }
                result.push(msg)
            } else if (message.role === "assistant") {
                let msg: Message = {
                    role: "assistant",
                    content: message.content as string || ""
                }
                if (message.tool_calls) {
                    const toolCalls: ToolCall[] = []
                    message.tool_calls.forEach(call => {

                        try {
                            const args = JSON.parse(call.function.arguments)
                            const toolCall: ToolCall = {
                                function: {
                                    name: call.function.name,
                                    arguments: args
                                }
                            }
                            toolCalls.push(toolCall)
                        } catch (err) {
                            Logger.error("parse tool call error", err)
                        }

                    })
                    msg.tool_calls = toolCalls
                }
                result.push(msg)
            }
        })
        return result
    }
    async chatComplete(params: ChatCompletionCreateParams, options?: any): Promise<ChatCompletionStream<ChatCompletionChunk> | ChatCompletion> {
        const messages = this.messagesFormat(params.messages)
        const id = v4()
        let chatRequest: ChatRequest = {
            model: params.model,
            messages: messages,
            options: options
        }
        if (params.response_format) {
            chatRequest.format = params.response_format.type || "text"
        }
        chatRequest.options = {
            temperature: params.temperature || undefined,
            top_p: params.top_p || undefined,
            presence_penalty: params.presence_penalty || undefined,
            frequency_penalty: params.frequency_penalty || undefined,
            seed: params.seed || undefined,
            stop: params.stop || undefined,
        } as Options

        if (params.tools) {
            // 配置tool
            params.stream = false
            const tools: Tool[] = []
            params.tools.forEach(tool => {
                const t: Tool = {
                    type: tool.type,
                    function: {
                        name: tool.function.name,
                        description: tool.function.description || "",
                        parameters: tool.function.parameters as any
                    }
                }
                tools.push(t)
            })
            chatRequest.tools = tools
        }

        if (params.stream) {
            chatRequest.stream = true
            // const response = await this.ollama.chat(chatRequest)
            const stream = new ChatCompletionStream<ChatCompletionChunk>()

            this.ollama.chat(chatRequest as ChatRequest & { stream: true }).then(async (response) => {

                for await (const chuck of response) {

                    const content = chuck.message.content
                    
                    const data: ChatCompletionChunk = {
                        id: id,
                        created: Date.now(),
                        model: params.model,
                        object: "chat.completion.chunk",
                        choices: [
                            {
                                index: 0,
                                delta: {
                                    role: "assistant",
                                    content: content,
                                    // tool_calls: chuck.message.tool_calls ? chuck.message.tool_calls.map(t => {
                                    //     return {
                                    //         id: "",
                                    //         index: 0,
                                    //         type: "function",
                                    //         function: {
                                    //             name: t.function.name,
                                    //             arguments: JSON.stringify(t.function.arguments)
                                    //         }
                                    //     }
                                    // }) : undefined
                                    // 目前还不支持
                                },
                                finish_reason: chuck.done ? "stop" : null,
                                logprobs: null
                            }
                        ]
                    }
                    stream.write(data)

                    if (chuck.done) {
                        setTimeout(() => {
                            stream.end()
                        }, 500)
                    }

                }
            })

            return stream

        } else {
            chatRequest.stream = false
            const response: ChatResponse = await this.ollama.chat(chatRequest as ChatRequest & { stream: false })
            const toolCalls: ChatCompletionMessageToolCall[] = []
            if (response.message.tool_calls) {
                response.message.tool_calls.forEach(call => {
                    const toolCall: ChatCompletionMessageToolCall = {
                        function: {
                            name: call.function.name,
                            arguments: (typeof call.function.arguments === 'string') ? call.function.arguments : JSON.stringify(call.function.arguments)
                        },
                        id: "",
                        type: 'function'
                    }
                    toolCalls.push(toolCall)
                })
            }

            const result: ChatCompletion = {
                id: id,
                created: Date.now(),
                model: params.model,
                object: "chat.completion",
                usage: {
                    completion_tokens: response.eval_duration,
                    prompt_tokens: response.load_duration,
                    total_tokens: response.total_duration
                },
                choices: [
                    {
                        index: 0,
                        message: {
                            role: response.message.role as "assistant",
                            content: response.message.content,
                            tool_calls: this.responseToolCallsFormat(response.message.tool_calls || []),
                            refusal: null
                        },
                        finish_reason: "stop",
                        logprobs: null
                    }
                ]
            }
            return result

        }


    }

    private responseToolCallsFormat(toolCalls: ToolCall[]): ChatCompletionMessageToolCall[] | undefined {
        const calls: ChatCompletionMessageToolCall[] = []
        if (toolCalls) {
            toolCalls.forEach(call => {
                const toolCall: ChatCompletionMessageToolCall = {
                    function: {
                        name: call.function.name,
                        arguments: (typeof call.function.arguments === 'string') ? call.function.arguments : JSON.stringify(call.function.arguments)
                    },
                    id: "",
                    type: 'function'
                }
                calls.push(toolCall)
            })
        }

        return calls.length ? calls : undefined
    }

    async embeddings(params: EmbeddingCreateParams): Promise<CreateEmbeddingResponse> {

        const inputs = (typeof params.input === "string") ? [params.input] : params.input
        const embeddings: Embedding[] = []

        for (let i = 0; i < inputs.length; i++) {
            const result = await this.ollama.embeddings({
                model: params.model,
                prompt: inputs[i] as string
            })
            Logger.debug("ollama embeddings", result)
            embeddings.push({
                embedding: result.embedding,
                index: i,
                object: 'embedding'
            })
        }

        return {
            object: "list",
            data: embeddings,
            model: params.model,
            usage: {
                prompt_tokens: 500 * inputs.length,
                total_tokens: 500 * inputs.length
            }
        } as CreateEmbeddingResponse
    }
}