import { VisionInterface } from "../../interfaces/Vision";
import { API_HOSTS } from "../../types/host";
import { ImageParams, ImageResult, VideoParams, VideoResult } from "../../types/vision/vision";
import { ImageGenerateParams } from "openai/resources";
import { VisionResultStream } from "../../utils/stream";
import { isImageURL } from "../../utils";
import { VisionBase } from "./base";
import { Logger } from "../../utils/Logger";

export class VisionZhiPuService extends VisionBase implements VisionInterface {

    apiKey: string;
    constructor(apiKey: string) {
        super()
        this.apiKey = apiKey
    }

    async image( params: ImageGenerateParams & ImageParams): Promise<VisionResultStream<ImageResult>> {
        const body:any = {
            model: params.model || "cogview-3",
            prompt: params.prompt
        }
        if (params?.user) {
            body.user_id = params.user
        }
        const stream = new VisionResultStream<ImageResult>()
        const timer = setInterval(() => {
            stream.write({
                status: "processing",
                created_at: Date.now()
            } as ImageResult)
        }, 1000)

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

            if (!response.ok || response.status !== 200) {
                stream.write({
                    status: "failed",
                    created_at: Date.now(),
                    message: `${response.status} ${response.statusText}`
                } as ImageResult)
                stream.end()
                return
            }

            const json = await response.json()
            stream.write({
                status: "completed",
                created_at: Date.now(),
                data: json.data.map((item: any) => {
                    return {
                        format: "url",
                        url: item.url
                    }
                })
            } as ImageResult)
            stream.end()
        })
        
        return stream

    }

    async video(params: VideoParams): Promise<VisionResultStream<VideoResult>> {
        if (params.prompt.length > 500) {
            return Promise.reject("prompt must be less than 500 characters")
        }
        const body:any = {
            model: params.model || "cogvideox",
            prompt: params.prompt
        }
        if (params.image) {
            // if (!isImageURL(params.image)) {
            //     return Promise.reject("image must be a url or base64")
            // }
            body.image_url = params.image
        }
        if (params?.user) {
            body.user_id = params.user
        }
        if (params.requestId) {
            body.request_id = params.requestId
        }
        const stream = new VisionResultStream<VideoResult>()
        Logger.debug("body:", body)
        const response = await fetch(API_HOSTS.ZHIPU + '/videos/generations', {
            method: "POST",
            headers: {
                "Content-Type": "application/json",
                "Authorization": `Bearer ${this.apiKey}`
            },
            body: JSON.stringify(body)
        })
        Logger.debug("response:", response)
        if (!response.ok || response.status !== 200) {
            this.streamFail(stream, `${response.status} ${response.statusText}`)
            return stream
        }

        const json = await response.json()
        const taskId = json.id

        const taskRequest = async (taskId: string): Promise<void> => {
            const taskRes = await fetch(API_HOSTS.ZHIPU + '/async-result/' + taskId, {
                method: "GET",
                headers: {
                    "Content-Type": "application/json",
                    "Authorization": `Bearer ${this.apiKey}`
                }
            })

            if (!taskRes.ok || taskRes.status !== 200) {
                this.streamFail(stream, taskRes.statusText)
                return
            }

            const taskData = await taskRes.json()
            Logger.debug("taskData:", taskData)
            if (taskData.task_status === "SUCCESS") {
                stream.write({
                    status: "completed",
                    created_at: Date.now(),
                    data: {
                        url: taskData.video_result[0].url,
                        cover: taskData.video_result[0].cover_image_url
                    }
                } as VideoResult)
                stream.end()
                return
            } else if (taskData.task_status === "PROCESSING") {
                stream.write({
                    status: "progressing",
                    created_at: Date.now(),
                    message: "is progressing"
                } as VideoResult)
                setTimeout(() => {
                    taskRequest(taskId)
                }, 5000)
            } else {
                this.streamFail(stream, `${taskId}: task status is ${taskData.task_status}, request id is ${taskData.request_id}`)
                return
            }
        }

        taskRequest(taskId)        
        return stream

    }
}