import { ChatCompletionCreateParams, ChatCompletionChunk, ChatCompletion, CreateEmbeddingResponse, EmbeddingCreateParams } from "openai/resources";
import { OpenAILikeInterface } from "../interfaces/OpenAILike";
import { SenseKnowledgeConfig, SensePlugins } from "../types/sense";
import jwt from 'jsonwebtoken'
import moment from "moment";
import { ChatCompletionStream, streamingResponseHandle } from "../utils/stream";
import { Logger } from "../utils/Logger";

const SENSE_API_URL = 'https://api.sensenova.cn'
export class SenseAPIService implements OpenAILikeInterface {

    ak: string
    sk: string
    constructor(accessKeyId: string, accessKeySecret: string) {
        this.ak = accessKeyId
        this.sk = accessKeySecret
    }
    async chatComplete(params: ChatCompletionCreateParams & {
        knowledge_config?: SenseKnowledgeConfig,
        plugins?: SensePlugins
        know_ids?: string[]
    }): Promise<ChatCompletionStream<ChatCompletionChunk> | ChatCompletion> {
        
        const response = await fetch(SENSE_API_URL + '/v1/llm/chat-completions', {
            method: "POST",
            headers: {
                'Content-Type': 'application/json',
                'Authorization': 'Bearer ' + this.generateAPIToken()
            },
            body: JSON.stringify(params)
        })
       
        if (!response.ok || response.status !== 200) {
            return Promise.reject(response.status + ":" + response.statusText)
        }

        if (params.stream) {
            const stream = new ChatCompletionStream<ChatCompletionChunk>()
            streamingResponseHandle(response, (data) => {
                if (data === null) {
                    stream.end()
                } else {
                    try {
                        const json = JSON.parse(data).data
                        const chuck: ChatCompletionChunk = {
                            id: json.id,
                            model: params.model,
                            created: Date.now(),
                            object: "chat.completion.chunk",
                            choices: json.choices.map((choice: any) => {
                                return {
                                    index: choice.index,
                                    delta: {
                                        role: choice.role,
                                        content: choice.delta,
                                        tool_calls: choice.tool_calls
                                    },
                                    finish_reason: choice.finish_reason,
                                    
                                }
                            }),
                        }
                        stream.write(chuck)

                        if (json.choices[0].finish_reason !== null) {
                            setTimeout(() => {
                                stream.end()
                            }, 500)
                        }
                    } catch (err) {
                        Logger.error("sense chat complete stream parse error", err)
                    }
                }
            }, {
                prefixStr: "data:"
            })
            return stream
        } else {
            const json = await response.json()
            const data = json.data

            return {
                id: data.id,
                object: "chat.completion",
                model: params.model,
                created: Date.now(),
                choices: data.choices.map((choice: any) => {
                    return {
                        index: choice.index,
                        message: {
                            role: choice.role,
                            content: choice.message,
                            tool_calls: choice.tool_calls
                        },
                        finish_reason: choice.finish_reason
                    }
                }),
                usage: {
                    prompt_tokens: data.usage.prompt_tokens,
                    total_tokens: data.usage.total_tokens,
                    completion_tokens: data.usage.completion_tokens
                }
            } as ChatCompletion
        }

    }

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

        const req = {
            model: params.model,
            input: typeof params.input === 'string' ? [params.input] : params.input as string[]
        }
        const response = await fetch(SENSE_API_URL + '/v1/llm/embeddings', {
            method: "POST",
            headers: {
                'Content-Type': 'application/json',
                'Authorization': 'Bearer ' + this.generateAPIToken()
            },
            body: JSON.stringify(req)
        })
        const json = await response.json()
        const result: CreateEmbeddingResponse = {
            object: "list",
            data: json.embeddings.map((item: any) => {
                return {
                    object: "embedding",
                    index: item.index,
                    embedding: item.embedding
                }
            }),
            model: params.model,
            usage: {
                prompt_tokens: json.usage.prompt_tokens,
                total_tokens: json.usage.total_tokens
            }
        }
        return result

    }

    generateAPIToken(): string {
        const headers = {
            alg: "HS256",
            typ: "JWT"
        };
        const payload = {
            iss: this.ak,
            exp: moment().add(30, 'minutes').unix(), // Set the expiration time to current time + 30 minutes
            nbf: moment().subtract(5, 'seconds').unix() // Set the not before time to current time - 5 seconds
        };
        const token = jwt.sign(payload, this.sk, {
            algorithm: "HS256",
            header: headers
        });
        return token;
    }


}