import { ChatCompletion, ChatCompletionChunk, ChatCompletionCreateParams, CreateEmbeddingResponse, EmbeddingCreateParams, ImageGenerateParams, ImagesResponse } from "openai/resources";
import { OpenAILikeInterface } from "../interfaces/OpenAILike";
import { OpenAIService } from "./openai";
import { ChatCompletionStream } from "../utils/stream";
import { API_HOSTS } from "../types/host";
import { Logger } from "../utils/Logger";

export class ZhiPuAPIService extends OpenAIService implements OpenAILikeInterface {

    constructor(apiKey: string) {
        super({
            host: API_HOSTS.ZHIPU,
            apiKey: `${apiKey}`
        })
    }

    chatComplete(params: ChatCompletionCreateParams): Promise<ChatCompletionStream<ChatCompletionChunk> | ChatCompletion> {
        // zhipu支持web_search，在返回中有web_search字段
        return super.chatComplete(params)
    }

    embeddings(params: EmbeddingCreateParams): Promise<CreateEmbeddingResponse> {
        return super.embeddings(params)
    }

    async imageGenerate(params: ImageGenerateParams): Promise<ImagesResponse> {
        const body:any = {
            model: params.model || "cogview-3",
            prompt: params.prompt
        }
        if (params.user) {
            body.user_id = params.user
        }

        const response = await fetch(API_HOSTS.ZHIPU + '/images/generations', {
            method: "POST",
            headers: {
                "Content-Type": "application/json",
                "Authorization": `Bearer ${this.apiKey}`
            },
            body: JSON.stringify(body)
        }) 

        const json = await response.json()
        const result: ImagesResponse = {
            created: json.created,
            data: json.data.map((item: any) => {
                return {
                    url: item.url
                }
            })
        }
        return result
    }
}