package com.sonic.cow.client.impl;

import com.alibaba.fastjson.JSON;
import com.sonic.cow.client.TextGenImageClient;
import com.byteplus.ark.runtime.model.images.generation.GenerateImagesRequest;
import com.byteplus.ark.runtime.model.images.generation.ImagesResponse;
import com.byteplus.ark.runtime.model.images.generation.ResponseFormat;
import com.byteplus.ark.runtime.service.ArkService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

@Slf4j
@Service
public class TextToImageClientImpl implements TextGenImageClient {

    @Value("${volcengine.textToImageModel}")
    private String textToImageModel;

    @Autowired
    private ArkService volcengineArkService;

    @Override
    public String textGenImage(String prompt) throws Exception{
        try {
            Long startTime = System.currentTimeMillis();
            GenerateImagesRequest generateRequest = GenerateImagesRequest.builder()
                    .model(textToImageModel)
                    .prompt(prompt)
                    .responseFormat(ResponseFormat.Base64)
                    .seed(1234567890)
                    .watermark(false)
                    //图片大小待确定
                    .size("720x1280")
                    .guidanceScale(2.5)
                    .build();
            ImagesResponse imagesResponse = volcengineArkService.generateImages(generateRequest);
            log.info("textToImage imagesResponse:" + JSON.toJSONString(imagesResponse));
            Long endTime = System.currentTimeMillis();
            log.info("textToImage 生成图片所需要的时间：{}", (endTime - startTime) / 1000);
            return imagesResponse.getData().get(0).getB64Json();
        } catch (Exception e) {
            log.error("textToImage error:", e);
            throw new Exception(e.getMessage());
        } finally {
            volcengineArkService.shutdownExecutor();
        }
    }
}
