package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
)

// 用于第一个 API 请求的结构体
type HeadPortraitPromptRequest struct {
	Name string `json:"name"`
}

// 用于接收第一个 API 响应的结构体
type PromptResponse struct {
	Prompt string `json:"prompt"`
}

// 第二个 API 请求的结构体
type PortraitRequest struct {
	Prompt   string `json:"prompt"`
	Negative string `json:"negative"`
}

// 处理生成头像的 HTTP 请求
func GenerateAvatarHandler(c *gin.Context) {
	name := c.PostForm("name")
	if name == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "未提供名字"})
		return
	}

	savePath := "./static/Worlds/headPortrait/"
	imagePath := savePath + name + ".png"

	// 检查本地是否已有头像图片
	if _, err := os.Stat(imagePath); err == nil {
		log.Println("已存在头像，直接返回")
		c.JSON(http.StatusOK, gin.H{"image_url": imagePath})
		return
	}

	// 打印进度条开始
	printProgressBar(0, "开始生成头像...")

	portraitData, err := GenerateHeadPortrait(name)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	// 更新进度条
	printProgressBar(70, "头像生成完毕，保存中...")

	// 保存头像图片到服务器
	imagePath, err = SaveImage(savePath, name+".png", portraitData)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	// 更新并完成进度条
	printProgressBar(100, "头像保存成功！")

	// 将图片地址直接作为响应返回
	c.JSON(http.StatusOK, gin.H{"image_url": imagePath})
}

// 生成人物头像，返回图片数据和错误信息
func GenerateHeadPortrait(name string) ([]byte, error) {
	promptAPI := "http://58.34.54.45:16666/generate_character_prompt"
	headPortraitAPI := "http://58.34.54.45:19666/"

	promptReq := HeadPortraitPromptRequest{Name: name}
	reqBody, err := json.Marshal(promptReq)
	if err != nil {
		return nil, fmt.Errorf("请求体编码失败: %v", err)
	}

	resp, err := http.Post(promptAPI, "application/json", bytes.NewBuffer(reqBody))
	if err != nil {
		return nil, fmt.Errorf("调用第一个 API 失败: %v", err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取响应体失败: %v", err)
	}

	var promptResp PromptResponse
	if err := json.Unmarshal(body, &promptResp); err != nil {
		return nil, fmt.Errorf("响应解码失败: %v", err)
	}

	portraitReq := PortraitRequest{
		Prompt:   promptResp.Prompt,
		Negative: "normal",
	}
	reqBody, err = json.Marshal(portraitReq)
	if err != nil {
		return nil, fmt.Errorf("请求体编码失败: %v", err)
	}

	resp, err = http.Post(headPortraitAPI, "application/json", bytes.NewBuffer(reqBody))
	if err != nil {
		return nil, fmt.Errorf("调用第二个 API 失败: %v", err)
	}
	defer resp.Body.Close()

	portraitData, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取第二响应体失败: %v", err)
	}

	return portraitData, nil
}

// 保存图片到指定路径
func SaveImage(directory, filename string, data []byte) (string, error) {
	if err := os.MkdirAll(directory, 0755); err != nil {
		return "", fmt.Errorf("创建目录失败: %v", err)
	}

	filePath := directory + filename
	file, err := os.Create(filePath)
	if err != nil {
		return "", fmt.Errorf("创建文件失败: %v", err)
	}
	defer file.Close()

	if _, err = file.Write(data); err != nil {
		return "", fmt.Errorf("写入文件失败: %v", err)
	}

	return filePath, nil
}

// 打印绿色方块进度条
func printProgressBar(percentage int, message string) {
	barLength := 50 // 进度条长度
	progress := int(float64(percentage) / 100.0 * float64(barLength))

	// 使用绿色方块表示进度
	bar := strings.Repeat("\033[32m█\033[0m", progress) + strings.Repeat(" ", barLength-progress)

	// 使用 \r 将光标移动到行首，确保进度条在同一行上更新
	fmt.Printf("\r[%s] %d%% - %s", bar, percentage, message)
	if percentage == 100 {
		fmt.Println() // 进度条完成，换行
	}
}

func main() {
	router := gin.Default()
	router.POST("/generate_avatar", GenerateAvatarHandler)
	router.Run(":8080")
}