package zhipu

import (
	"context"
	"errors"
	"time"

	"apigo.cc/ai"
	"github.com/yankeguo/zhipu"
)

func MakeImage(aiConf *ai.AIConfig, conf ai.ImageConfig) (ai.ImageResult, error) {
	c, err := getClient(aiConf)
	if err != nil {
		return ai.ImageResult{}, err
	}

	cc := c.ImageGeneration(conf.Model).SetPrompt(conf.SystemPrompt + conf.Prompt)
	//cc.SetSize(config.GetSize())

	t1 := time.Now().UnixMilli()
	if r, err := cc.Do(context.Background()); err == nil {
		t2 := time.Now().UnixMilli() - t1
		results := make([]string, 0)
		for _, item := range r.Data {
			results = append(results, item.URL)
		}
		if len(results) == 0 {
			results = append(results, "")
		}
		return ai.ImageResult{
			Results:  results,
			UsedTime: t2,
		}, nil
	} else {
		return ai.ImageResult{}, err
	}
}

func MakeVideo(aiConf *ai.AIConfig, conf ai.VideoConfig) (string, error) {
	c, err := getClient(aiConf)
	if err != nil {
		return "", err
	}

	cc := c.VideoGeneration(conf.Model).SetPrompt(conf.SystemPrompt + conf.Prompt)
	if len(conf.Ref) > 0 {
		cc.SetImageURL(conf.Ref[0])
	}

	if resp, err := cc.Do(context.Background()); err == nil {
		return resp.ID, nil
	} else {
		return "", err
	}
}

func GetVideoResult(aiConf *ai.AIConfig, taskId string, waitSeconds int) (ai.VideoResult, error) {
	c, err := getClient(aiConf)
	if err != nil {
		return ai.VideoResult{}, err
	}

	var r zhipu.AsyncResultResponse
	for i := 0; i < waitSeconds; i += 3 {
		r, err = c.AsyncResult(taskId).Do(context.Background())
		if err != nil {
			return ai.VideoResult{}, err
		}
		if r.TaskStatus == zhipu.VideoGenerationTaskStatusSuccess {
			results := make([]string, 0)
			previews := make([]string, 0)
			for _, item := range r.VideoResult {
				results = append(results, item.URL)
				previews = append(previews, item.CoverImageURL)
			}
			if len(results) == 0 {
				results = append(results, "")
				previews = append(previews, "")
			}
			return ai.VideoResult{
				Results:  results,
				Previews: previews,
				UsedTime: 0,
			}, nil
		}
		if r.TaskStatus == zhipu.VideoGenerationTaskStatusFail {
			return ai.VideoResult{}, errors.New("fail on video task " + taskId)
		} else if r.TaskStatus == zhipu.VideoGenerationTaskStatusProcessing {
			if waitSeconds == 0 {
				return ai.VideoResult{IsProcessing: true}, nil
			}
			time.Sleep(3 * time.Second)
		} else {
			return ai.VideoResult{}, errors.New("unknow status " + r.TaskStatus + " on video task " + taskId)
		}
	}
	return ai.VideoResult{IsProcessing: true}, errors.New("timeout on video task " + taskId)
}