import { ChatCompletionCreateParams, ChatCompletionChunk, ChatCompletion, CreateEmbeddingResponse, EmbeddingCreateParams, Embedding, ImageGenerateParams, ImagesResponse } from "openai/resources";
import { OpenAILikeInterface } from "../interfaces/OpenAILike";
import { Client } from "tencentcloud-sdk-nodejs-hunyuan/tencentcloud/services/hunyuan/v20230901/hunyuan_client";
import { ChatCompletionsRequest, ChatCompletionsResponse, Content, Message, TextToImageLiteRequest, Tool } from "tencentcloud-sdk-nodejs-hunyuan/tencentcloud/services/hunyuan/v20230901/hunyuan_models";
import { Stream } from "stream";
import { v4 } from "uuid";
import { HunYuanEmbeddingModel } from "../types/models";
import { ChatCompletionStream } from "../utils/stream";
import { Logger } from "../utils/Logger";
import { HunYuanArtStyle } from "../types/hunyuan";

export class HunYuanService implements OpenAILikeInterface {

    client: Client

    constructor(secretId: string, secretKey: string) {
        this.client = new Client({
            credential: {
                secretId,
                secretKey
            }
        })
    }
    async chatComplete(params: ChatCompletionCreateParams): Promise<ChatCompletionStream<ChatCompletionChunk> | ChatCompletion> {

        const messages: Message[] = []
        params.messages.forEach((msg) => {
            const message: Message = {
                Role: msg.role,
                Content: "",
            }
            if (msg.role === "tool") {
                message.ToolCallId = msg.tool_call_id
            } else if (msg.role === "assistant") {
                let index = -1
                message.ToolCalls = msg.tool_calls?.map((t) => {
                    index++
                    return {
                        Id: t.id,
                        Index: index,
                        Type: "function",
                        Function: {
                            Name: t.function.name,
                            Arguments: t.function.arguments
                        }
                    }
                })
            }

            if (msg.content instanceof Array) {
                if (msg.content.length === 1 && msg.content[0].type === "text") {
                    // Logger.warn("HunYuan does not support multiple content parts, only the first part will be used")
                    message.Content = msg.content[0].text
                } else {
                    const contents: Content[] = []
                    msg.content.forEach(part => {
                        if (part.type === "text") {
                            contents.push({
                                Type: part.type,
                                Text: part.text
                            })
                        } else if (part.type === "image_url") {
                            contents.push({
                                Type: part.type,
                                ImageUrl: {
                                    Url: part.image_url.url
                                }
                            })
                        }
                    })
                    message.Contents = contents.length ? contents : undefined
                }
                
            } else {
                message.Content = msg.content || ""
            }

            messages.push(message)
        })

        const tools: Tool[] = []
        params.tools?.forEach((tool) => {
            tools.push({
                Type: tool.type,
                Function: {
                    Name: tool.function.name,
                    Description: tool.function.description,
                    Parameters: JSON.stringify(tool.function.parameters)
                }
            })
        })
        const req: ChatCompletionsRequest = {
            Model: params.model,
            Messages: messages,
            Stream: params.stream || false,
            TopP: params.top_p || undefined,
            Temperature: params.temperature || undefined,
            Tools: tools.length > 0 ? tools : undefined,
            ToolChoice: params.tool_choice as string | undefined || undefined,
        }

        if (!params.stream) {
            const res = await this.client.ChatCompletions(req)
            const result: ChatCompletion = {
                id: res.Id || res.RequestId || "",
                created: res.Created || 0,
                model: params.model,
                object: "chat.completion",
                usage: {
                    completion_tokens: res.Usage?.CompletionTokens || 0,
                    prompt_tokens: res.Usage?.PromptTokens || 0,
                    total_tokens: res.Usage?.TotalTokens || 0
                },
                choices: [
                    {
                        index: 0,
                        message: {
                            role: "assistant",
                            content: res.Choices ? res.Choices[0].Message?.Content || null : null,
                            tool_calls: res.Choices ? res.Choices[0].Message?.ToolCalls?.map((t) => {
                                return {
                                    id: t.Id,
                                    index: undefined,
                                    type: "function",
                                    function: {
                                        name: t.Function.Name,
                                        arguments: t.Function.Arguments
                                    }
                                }
                            }) : undefined,
                            refusal: null
                        },
                        finish_reason: res.Choices ? res.Choices[0].FinishReason : "stop" as any,
                        logprobs: null
                    }
                ]
            }
            return result
        } else {

            const stream = new ChatCompletionStream<ChatCompletionChunk>()
            let index = 0
            let finish = null as 'stop' | 'length' | 'tool_calls' | 'content_filter' | 'function_call' | null

            const response = await this.client.ChatCompletions(req) as Stream
            const id = v4()
            response.on("message", (data: any) => {
                // Logger.debug("stream message", data)
                try {
                    const rep = JSON.parse(data.data) as ChatCompletionsResponse
                    if (rep.Choices && rep.Choices[0].FinishReason) {
                        finish = rep.Choices[0].FinishReason as 'stop' | 'length' | 'tool_calls' | 'content_filter' | 'function_call' | null
                    }

                    const content = rep.Choices ? (rep.Choices[0].Delta?.Content || "") : ""
                    const chuck: ChatCompletionChunk = {
                        id: rep.Id || rep.RequestId || id,
                        created: rep.Created || 0,
                        model: params.model,
                        object: "chat.completion.chunk",
                        choices: [
                            {
                                index,
                                delta: {
                                    role: "assistant",
                                    content: content,
                                    tool_calls: rep.Choices ? rep.Choices[0].Delta?.ToolCalls?.map(t => {
                                        return {
                                            id: t.Id,
                                            index: 0,
                                            type: t.Type,
                                            function: {
                                                name: t.Function.Name,
                                                arguments: t.Function.Arguments
                                            }
                                        } as ChatCompletionChunk.Choice.Delta.ToolCall
                                    }) : undefined
                                },
                                finish_reason: finish,
                                logprobs: null
                            }
                        ]
                    }
                    stream.write(chuck)
                    index++

                    if (finish) {
                        setTimeout(() => {
                            stream.end()
                        }, 500)
                    }
                } catch (err) {
                    Logger.error('Error:', err)
                }

            })

            return stream
        }
    }

    async embeddings(params: EmbeddingCreateParams): Promise<CreateEmbeddingResponse> {
        if (params.model !== HunYuanEmbeddingModel.HUNYUAN_EMBEDDING) {
            throw new Error("hunyuan embedding model only support hunyuan-embedding")
        }

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

        let promptTokens = 0
        let totalTokens = 0
        for (let index = 0; index < inputs.length; index++) {
            const input = inputs[index];
            const result = await this.client.GetEmbedding({
                Input: input as string
            })

            if (result.Data && result.Data[0].Embedding) {
                embeddings.push({
                    embedding: result.Data[0].Embedding,
                    index: index,
                    object: "embedding"
                })
                promptTokens += result.Usage?.PromptTokens || 0
                totalTokens += result.Usage?.TotalTokens || 0
            }
        }

        return {
            object: "list",
            data: embeddings,
            model: params.model,
            usage: {
                prompt_tokens: promptTokens,
                total_tokens: totalTokens
            }
        } as CreateEmbeddingResponse
    }

    async imageGenerate(params: ImageGenerateParams & {
        negative?: string,
        styleH?: HunYuanArtStyle,
        sizeH?: "768:768" | "768:1024" | "1024:768" | "1024:1024" | "720:1280" | "1280:720" | "768:1280" | "1280:768" | "1080:1920" | "1920:1080",
        logoAdd?: boolean,
        logoParam?: string
    }): Promise<ImagesResponse> {
        const type = params.response_format || "b64_json"
        const req: TextToImageLiteRequest = {
            Prompt: params.prompt,
            NegativePrompt: params.negative,
            Style: (typeof params.styleH === "string") ? params.styleH : (params.styleH as any).toString(),
            Resolution: params.sizeH || "768:768",
            RspImgType: params.response_format === "url" ? "url" : "base64",
            LogoAdd: params.logoAdd ? 1 : 0
        }
        Logger.debug("req", req)
        this.client.region = "ap-guangzhou"

        return new Promise((resolve, reject) => {
            this.client.TextToImageLite(req, (err, res) => {
                if (err) {
                    Logger.error(err)
                    reject(err)
                    return
                }
                const rep = {
                    created: Date.now(),
                    data: [
                        (type === "b64_json") ? {
                            b64_json: res.ResultImage
                        } : {
                            url: res.ResultImage
                        }
                    ]
                } as ImagesResponse
                resolve(rep)
            })
        })


        
    }

}