vision/canvas.go

package vision

import (
	"bytes"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/jpeg"
	_ "image/png"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"apigo.cc/go/file"
	"github.com/fogleman/gg"
	"golang.org/x/image/font"
)

// Canvas 代表一个绘图画布，封装了图像处理与绘制能力
type Canvas struct {
	dc        *gg.Context
	bgColor   string
	lastColor string
	lastFont  font.Face
}

// New 创建一个新的画布
func New(width, height int, backgroundColor ...string) *Canvas {
	img := image.NewRGBA(image.Rect(0, 0, width, height))
	dc := gg.NewContextForImage(img)
	c := &Canvas{
		dc: dc,
	}
	if len(backgroundColor) > 0 && backgroundColor[0] != "" {
		c.bgColor = backgroundColor[0]
		c.SetColor(c.bgColor)
		c.dc.Clear()
	}
	return c
}

// Load 从文件加载图像并创建画布
func Load(path string) (*Canvas, error) {
	if path == "" {
		return nil, fmt.Errorf("path is empty")
	}
	if !file.Exists(path) {
		return nil, fmt.Errorf("file not found: %s", path)
	}

	data, err := file.ReadBytes(path)
	if err != nil {
		return nil, err
	}

	img, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		// 尝试使用 FFmpeg 作为 fallback (用于 HEIC 等格式)
		ext := strings.ToLower(filepath.Ext(path))
		if ext == ".heic" || ext == ".heif" || ext == ".webp" || ext == ".avif" {
			return loadWithFFmpeg(path)
		}
		return nil, fmt.Errorf("decode image failed: %v", err)
	}

	return &Canvas{
		dc: gg.NewContextForImage(img),
	}, nil
}

func loadWithFFmpeg(path string) (*Canvas, error) {
	tmpFile := filepath.Join(os.TempDir(), fmt.Sprintf("vision_load_%d.png", os.Getpid()))
	defer os.Remove(tmpFile)

	// 如果是 HEIC/HEIF，优先使用专门的转换工具
	ext := strings.ToLower(filepath.Ext(path))
	if ext == ".heic" || ext == ".heif" {
		if err := ConvertHEIC(path, tmpFile); err == nil {
			return loadPNG(tmpFile)
		}
	}

	// 否则或失败后，回退到 FFmpeg
	ffmpeg, err := EnsureFFmpeg()
	if err != nil {
		return nil, fmt.Errorf("ffmpeg not found for fallback: %w", err)
	}

	// 将输入文件转换为 PNG (FFmpeg 对 HEIC 的网格重构支持较弱)
	cmd := exec.Command(ffmpeg, "-i", path, "-frames:v", "1", "-y", tmpFile)
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("ffmpeg decode fallback failed: %w", err)
	}

	return loadPNG(tmpFile)
}

func loadPNG(path string) (*Canvas, error) {
	data, err := file.ReadBytes(path)
	if err != nil {
		return nil, err
	}

	img, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, err
	}

	return &Canvas{
		dc: gg.NewContextForImage(img),
	}, nil
}

// Save 将画布保存到文件
func Save(c *Canvas, path string, quality ...int) error {
	var err error

	if strings.HasSuffix(strings.ToLower(path), ".jpg") || strings.HasSuffix(strings.ToLower(path), ".jpeg") {
		q := 85
		if len(quality) > 0 {
			q = quality[0]
		}
		// gg 没有内置 SaveJPG 到 context，我们需要手动编码
		f, createErr := file.Create(path)
		if createErr != nil {
			return createErr
		}
		defer f.Close()
		err = jpeg.Encode(f, c.dc.Image(), &jpeg.Options{Quality: q})
	} else {
		err = c.dc.SavePNG(path)
	}

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

// Image 返回底层图像
func (c *Canvas) Image() image.Image {
	return c.dc.Image()
}

// Width 返回画布宽度
func (c *Canvas) Width() int {
	return c.dc.Width()
}

// Height 返回画布高度
func (c *Canvas) Height() int {
	return c.dc.Height()
}

// SetColor 设置当前绘图颜色 (支持 hex 格式)
func (c *Canvas) SetColor(hex string) {
	c.lastColor = hex
	c.dc.SetColor(ParseColor(hex))
}

// Clear 清除指定区域，如果设置了背景色则填充背景色，否则填充透明
func (c *Canvas) Clear(x, y, w, h int) {
	if c.bgColor != "" {
		c.dc.Push()
		c.dc.DrawRectangle(float64(x), float64(y), float64(w), float64(h))
		c.dc.SetColor(ParseColor(c.bgColor))
		c.dc.Fill()
		c.dc.Pop()
		return
	}
	if img, ok := c.dc.Image().(*image.RGBA); ok {
		transparent := image.NewUniform(color.RGBA{0, 0, 0, 0})
		draw.Draw(img, image.Rect(x, y, x+w, y+h), transparent, image.Point{}, draw.Src)
	}
}

// Sub 提取子区域并返回新画布
func (c *Canvas) Sub(x, y, w, h int) *Canvas {
	newImg := image.NewRGBA(image.Rect(0, 0, w, h))
	draw.Draw(newImg, newImg.Bounds(), c.dc.Image(), image.Pt(x, y), draw.Src)
	newDC := gg.NewContextForImage(newImg)
	return &Canvas{
		dc:        newDC,
		bgColor:   c.bgColor,
		lastColor: c.lastColor,
		lastFont:  c.lastFont,
	}
}

// Clone 克隆当前画布
func (c *Canvas) Clone() *Canvas {
	bounds := c.dc.Image().Bounds()
	newImg := image.NewRGBA(bounds)
	draw.Draw(newImg, bounds, c.dc.Image(), bounds.Min, draw.Src)
	return &Canvas{
		dc:        gg.NewContextForImage(newImg),
		bgColor:   c.bgColor,
		lastColor: c.lastColor,
		lastFont:  c.lastFont,
	}
}