command-server/plugins/image/image.go

package image

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"ncatbot-command-server/command"
	"ncatbot-command-server/config"
	"os"
	"path/filepath"
	"time"

	"github.com/go-resty/resty/v2"
)

// Plugin 图片生成插件
type Plugin struct{}

var imgClient *resty.Client
var uploadClient *resty.Client

func init() {
	command.RegisterPlugin(&Plugin{})
}

// Name 实现 Plugin 接口
func (p *Plugin) Name() string {
	return "生图"
}

// Description 实现 Plugin 接口
func (p *Plugin) Description() string {
	return "根据描述生成图片"
}

// Usage 实现 Plugin 接口
func (p *Plugin) Usage() string {
	return "生图 <描述>"
}

// Init 实现 PluginWithInit 接口
func (p *Plugin) Init() {
	cfg := config.Cfg
	imgClient = newHTTPClient(cfg, true)
	uploadClient = newHTTPClient(cfg, false)
	os.MkdirAll(cfg.ImageOutputDir, 0755)
}

// Run 实现 Plugin 接口
func (p *Plugin) Run(req *command.Req) command.Resp {
	if req.Content == "" {
		return command.Resp{Reply: "请输入生图描述，例如：生图 一只猫"}
	}

	url, err := genImage(req.Content)
	if err != nil {
		log.Printf("gen image error: %v", err)
		return command.Resp{Reply: "生成失败: " + err.Error()}
	}

	return command.Resp{
		Messages: []command.Message{
			{Type: command.MsgTypeImage, URL: url},
		},
	}
}

func newHTTPClient(cfg config.Config, withAuth bool) *resty.Client {
	c := resty.New().
		SetTimeout(time.Duration(cfg.RequestTimeout) * time.Second)
	if withAuth {
		c.SetHeader("Content-Type", "application/json").
			SetHeader("Authorization", "Bearer "+cfg.OpenAIAPIKey)
	}
	if cfg.Proxy != "" {
		c.SetProxy(cfg.Proxy)
	}
	return c
}

// --- API Response types ---

type OpenAIImageResp struct {
	Created int64 `json:"created"`
	Data    []struct {
		URL           string `json:"url"`
		B64JSON       string `json:"b64_json"`
		RevisedPrompt string `json:"revised_prompt"`
	} `json:"data"`
}

type UploadResp struct {
	Code int    `json:"code"`
	Msg  string `json:"msg"`
	Data struct {
		Files []string `json:"files"`
		Path  string   `json:"path"`
	} `json:"data"`
}

// --- Pipeline ---

func genImage(prompt string) (string, error) {
	cfg := config.Cfg
	if cfg.OpenAIAPIKey == "" {
		return "", errors.New("OPENAI_API_KEY not set")
	}

	b64, err := requestImage(prompt)
	if err != nil {
		return "", err
	}

	filePath, err := saveImage(b64)
	if err != nil {
		return "", err
	}

	remotePath, err := uploadImage(filePath)
	if err != nil {
		return "", fmt.Errorf("upload failed: %w", err)
	}

	os.Remove(filePath)
	return remotePath, nil
}

func requestImage(prompt string) (string, error) {
	cfg := config.Cfg

	var result OpenAIImageResp
	resp, err := imgClient.R().
		SetBody(map[string]any{
			"model":           cfg.OpenAIModel,
			"prompt":          prompt,
			"n":               cfg.ImageCount,
			"response_format": "b64_json",
		}).
		SetResult(&result).
		Post(cfg.OpenAIBaseURL + config.ImageGenerationsPath)

	if err != nil {
		return "", fmt.Errorf("request failed: %w", err)
	}

	if resp.StatusCode() != 200 {
		var errResp struct {
			Error struct {
				Message string `json:"message"`
			} `json:"error"`
		}
		if json.Unmarshal(resp.Body(), &errResp) == nil && errResp.Error.Message != "" {
			return "", fmt.Errorf("%s", errResp.Error.Message)
		}
		return "", fmt.Errorf("API error (status %d): %s", resp.StatusCode(), resp.Status())
	}

	if len(result.Data) == 0 {
		return "", errors.New("empty image result")
	}

	return result.Data[0].B64JSON, nil
}

func saveImage(b64 string) (string, error) {
	cfg := config.Cfg

	imgBytes, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		return "", fmt.Errorf("base64 decode failed: %w", err)
	}

	filename := filepath.Join(cfg.ImageOutputDir, fmt.Sprintf("%d.png", time.Now().UnixNano()))
	if err := os.WriteFile(filename, imgBytes, 0644); err != nil {
		return "", fmt.Errorf("write file failed: %w", err)
	}

	return filename, nil
}

func uploadImage(filePath string) (string, error) {
	cfg := config.Cfg
	if cfg.UploadURL == "" {
		return "", errors.New("UPLOAD_URL not set")
	}

	filename := filepath.Base(filePath)
	var result UploadResp

	resp, err := uploadClient.R().
		SetHeader("X-API-Key", cfg.UploadAPIKey).
		SetFile(filename, filePath).
		SetResult(&result).
		Post(cfg.UploadURL)

	if err != nil {
		return "", fmt.Errorf("upload request failed: %w", err)
	}

	if resp.StatusCode() != 200 {
		return "", fmt.Errorf("upload failed (status %d): %s", resp.StatusCode(), resp.Status())
	}

	if result.Code != 0 {
		return "", fmt.Errorf("upload error: %s", result.Msg)
	}

	return result.Data.Path, nil
}