import { ChatCompletionCreateParams, ChatCompletionChunk, ChatCompletion, CreateEmbeddingResponse, EmbeddingCreateParams } from "openai/resources";
import { OpenAILikeInterface } from "../interfaces/OpenAILike";
import { ErnieChatCompletion, ErnieRequestParams } from "../types/ernie";
import { Logger } from "../utils/Logger";
import { ErnieAPI, ErnieMessage } from "../utils/erine";
import { ChatCompletionStream, StreamUtils } from "../utils/stream";


const ERINE_API_HOST = process.env.ERINE_API_HOST || "aip.baidubce.com"

export class ErnieAPIService implements OpenAILikeInterface {

    private clientId: string;
    private clientSecret: string;

    constructor(clientId: string, clientSecret: string) {
        this.clientId = clientId || "";
        this.clientSecret = clientSecret || "";

        if (!this.clientId || !this.clientSecret) {
            throw new Error("clientId and clientSecret is required");
        }
    }

    async chatComplete(params: ChatCompletionCreateParams): Promise<ChatCompletionStream<ChatCompletionChunk> | ChatCompletion> {
        const erineAPI = await ErnieAPI.getInstance(this.clientId, this.clientSecret)
        const url = "/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/" + params.model
        const messages: ErnieMessage[] = []
        params.messages.forEach(message => {
            messages.push({
                content: message.content as string,
                role: message.role as 'user' | 'assistant'
            })
        })
        const data: ErnieRequestParams = {
            stream: params.stream || false,
            temperature: params.temperature || undefined,
            top_p: params.top_p || undefined,
            frequency_penalty: params.frequency_penalty || undefined,
            presence_penalty: params.presence_penalty || undefined,
            max_output_tokens: params.max_tokens || undefined,
            messages: messages,
        }
        if (params.response_format?.type) {
            data.response_format = params.response_format.type || "text"
        }

        const response = await erineAPI.request({
            method: "POST",
            url: url,
            data: data,
        })
        if (!response.ok || response.status !== 200) {
            return Promise.reject(response.statusText)
        }

        const contentType = response.headers.get('content-type') || "";
        if (contentType.includes('application/json')) {
            const json = await response.json()
            console.error(json)
            if (json.error_code) {
                return Promise.reject(json.error_msg)
            }

            const result: ChatCompletion = {
                id: json.id,
                object: "chat.completion",
                created: json.created,
                model: params.model,
                choices: [{
                    index: 0,
                    message: {
                        role: "assistant",
                        content: json.result,
                        refusal: null
                    },
                    finish_reason: "stop",
                    logprobs: null,
                }],
                usage: {
                    prompt_tokens: json.usage.prompt_tokens,
                    completion_tokens: json.usage.completion_tokens,
                    total_tokens: json.usage.total_tokens
                }
            }
            return result
        } else if (contentType.includes('text/event-stream')) {
            if (!response.body) {
                return Promise.reject("Response body is empty")
            }

            const stream = new ChatCompletionStream<ChatCompletionChunk>()
            StreamUtils.handle(response, (chunk) => {
                if (chunk) {
                    try {
                        const json = JSON.parse(chunk) as ErnieChatCompletion
                        const result: ChatCompletionChunk = {
                            id: json.id,
                            object: "chat.completion.chunk",
                            created: json.created,
                            model: params.model,
                            choices: [{
                                index: 0,
                                delta: {
                                    role: "assistant",
                                    content: json.result
                                },
                                finish_reason: json.is_end ? "stop" : null
                            }]
                        }
                        stream.write(result)
                    } catch (err) {
                        Logger.error("")
                    }
                } else {
                    stream.end()
                }
            }, {
                spliteStr: "\n\n",
                prefixStr: "data: "
            })

            return stream
        } else {
            return Promise.reject("Invalid content type: " + contentType)
        }

    }

    embeddings(params: EmbeddingCreateParams): Promise<CreateEmbeddingResponse> {
        throw new Error("Method not implemented.");
    }

}