package zhipu

import (
	"apigo.cc/ai/llm/llm"
	"context"
	"errors"
	"github.com/yankeguo/zhipu"
	"time"
)

func (lm *LLM) FastMakeImage(prompt string, config llm.GCConfig) ([]string, llm.Usage, error) {
	config.Model = ModelCogView3Plus
	return lm.MakeImage(prompt, config)
}

func (lm *LLM) BestMakeImage(prompt string, config llm.GCConfig) ([]string, llm.Usage, error) {
	config.Model = ModelCogView3
	return lm.MakeImage(prompt, config)
}

func (lm *LLM) MakeImage(prompt string, config llm.GCConfig) ([]string, llm.Usage, error) {
	c, err := zhipu.NewClient(zhipu.WithAPIKey(lm.config.ApiKey), zhipu.WithBaseURL(lm.config.Endpoint))
	if err != nil {
		return nil, llm.Usage{}, err
	}

	config.SetDefault(&lm.config.GCConfig)
	cc := c.ImageGeneration(config.Model).SetPrompt(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)
		}
		return results, llm.Usage{
			UsedTime: t2,
		}, nil
	} else {
		return nil, llm.Usage{}, err
	}
}

func (lm *LLM) FastMakeVideo(prompt string, config llm.GCConfig) ([]string, []string, llm.Usage, error) {
	config.Model = ModelCogVideoX
	return lm.MakeVideo(prompt, config)
}

func (lm *LLM) BestMakeVideo(prompt string, config llm.GCConfig) ([]string, []string, llm.Usage, error) {
	config.Model = ModelCogVideoX
	return lm.MakeVideo(prompt, config)
}

func (lm *LLM) MakeVideo(prompt string, config llm.GCConfig) ([]string, []string, llm.Usage, error) {
	c, err := zhipu.NewClient(zhipu.WithAPIKey(lm.config.ApiKey), zhipu.WithBaseURL(lm.config.Endpoint))
	if err != nil {
		return nil, nil, llm.Usage{}, err
	}

	config.SetDefault(&lm.config.GCConfig)
	cc := c.VideoGeneration(config.Model).SetPrompt(prompt)
	cc.SetImageURL(config.GetRef())

	t1 := time.Now().UnixMilli()
	if resp, err := cc.Do(context.Background()); err == nil {
		t2 := time.Now().UnixMilli() - t1
		for i := 0; i < 1200; i++ {
			r, err := c.AsyncResult(resp.ID).Do(context.Background())
			if err != nil {
				return nil, nil, llm.Usage{}, err
			}
			if r.TaskStatus == zhipu.VideoGenerationTaskStatusSuccess {
				covers := make([]string, 0)
				results := make([]string, 0)
				for _, item := range r.VideoResult {
					results = append(results, item.URL)
					covers = append(covers, item.CoverImageURL)
				}
				return results, covers, llm.Usage{
					UsedTime: t2,
				}, nil
			}
			if r.TaskStatus == zhipu.VideoGenerationTaskStatusFail {
				return nil, nil, llm.Usage{}, errors.New("fail on task " + resp.ID)
			}
			time.Sleep(3 * time.Second)
		}
		return nil, nil, llm.Usage{}, errors.New("timeout on task " + resp.ID)
	} else {
		return nil, nil, llm.Usage{}, err
	}
}