qq_bot/util/text_to_pic.go

package util

import (
	"bytes"
	"image"
	"image/color"
	"image/jpeg"
	"image/png"
	"os"
	"path/filepath"
	"strings"

	"github.com/fogleman/gg"
	"github.com/yuin/goldmark"
	"golang.org/x/net/html"
)

const (
	defaultWidth   = 800
	defaultMargin  = 20
	lineHeight     = 1.5
	fontSize       = 14
	maxImageWidth  = 600
	maxImageHeight = 400
)

type ContentBlock struct {
	Type      string      // "text" or "image"
	Content   string      // text content or image path
	ImgWidth  float64     // image width after resize
	ImgHeight float64     // image height after resize
	Image     image.Image // loaded image
}

// TextToPicture 将文本转换为图片
func TextToPicture(text string, isMarkdown bool) (image.Image, error) {
	var htmlContent string
	if isMarkdown {
		// 将Markdown转换为HTML
		md := goldmark.New()
		var buf bytes.Buffer
		if err := md.Convert([]byte(text), &buf); err != nil {
			return nil, err
		}
		htmlContent = buf.String()
	} else {
		htmlContent = text
	}

	// 解析HTML并提取内容块
	doc, err := html.Parse(strings.NewReader(htmlContent))
	if err != nil {
		return nil, err
	}

	var blocks []ContentBlock
	var extractContent func(*html.Node)
	extractContent = func(n *html.Node) {
		if n.Type == html.TextNode {
			text := strings.TrimSpace(n.Data)
			if text != "" {
				blocks = append(blocks, ContentBlock{
					Type:    "text",
					Content: text,
				})
			}
		} else if n.Type == html.ElementNode && n.Data == "img" {
			// 处理图片节点
			var src string
			for _, attr := range n.Attr {
				if attr.Key == "src" {
					src = attr.Val
					break
				}
			}
			if src != "" && !strings.HasPrefix(src, "http") {
				// 处理本地图片
				img, w, h, err := loadAndResizeImage(src)
				if err == nil {
					blocks = append(blocks, ContentBlock{
						Type:      "image",
						Content:   src,
						Image:     img,
						ImgWidth:  w,
						ImgHeight: h,
					})
				}
			}
		}
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			extractContent(c)
		}
	}
	extractContent(doc)

	// 计算总高度
	totalHeight := calculateTotalHeight(blocks)

	// 创建图片上下文
	dc := gg.NewContext(defaultWidth, totalHeight)
	dc.SetRGB(1, 1, 1) // 设置白色背景
	dc.Clear()

	// 设置字体和颜色
	if err := dc.LoadFontFace("./resource/fonts/SourceHanSansCN-Regular.otf", fontSize); err != nil {
		return nil, err
	}
	dc.SetColor(color.Black)

	// 绘制内容
	y := float64(defaultMargin)
	for _, block := range blocks {
		if block.Type == "text" {
			wrapped := dc.WordWrap(block.Content, float64(defaultWidth-2*defaultMargin))
			for _, wrappedLine := range wrapped {
				dc.DrawString(wrappedLine, float64(defaultMargin), y+fontSize)
				y += fontSize * lineHeight
			}
			y += fontSize // 段落间额外间距
		} else if block.Type == "image" {
			// 居中绘制图片
			x := (float64(defaultWidth) - block.ImgWidth) / 2
			dc.DrawImage(block.Image, int(x), int(y))
			y += block.ImgHeight + float64(defaultMargin) // 图片后添加间距
		}
	}

	return dc.Image(), nil
}

// loadAndResizeImage 加载并调整图片大小
func loadAndResizeImage(path string) (image.Image, float64, float64, error) {
	// 确保路径是绝对路径
	if !filepath.IsAbs(path) {
		path = filepath.Join(".", path)
	}

	file, err := os.Open(path)
	if err != nil {
		return nil, 0, 0, err
	}
	defer file.Close()

	// 解码图片
	var img image.Image
	if strings.HasSuffix(strings.ToLower(path), ".png") {
		img, err = png.Decode(file)
	} else {
		img, err = jpeg.Decode(file)
	}
	if err != nil {
		return nil, 0, 0, err
	}

	// 计算调整后的尺寸
	bounds := img.Bounds()
	origWidth := float64(bounds.Dx())
	origHeight := float64(bounds.Dy())

	var newWidth, newHeight float64
	if origWidth > maxImageWidth {
		ratio := maxImageWidth / origWidth
		newWidth = maxImageWidth
		newHeight = origHeight * ratio
	} else {
		newWidth = origWidth
		newHeight = origHeight
	}

	if newHeight > maxImageHeight {
		ratio := maxImageHeight / newHeight
		newHeight = maxImageHeight
		newWidth = newWidth * ratio
	}

	// 创建新的图片上下文并调整大小
	dc := gg.NewContext(int(newWidth), int(newHeight))
	dc.DrawImage(img, 0, 0)

	return dc.Image(), newWidth, newHeight, nil
}

// calculateTotalHeight 计算总高度
func calculateTotalHeight(blocks []ContentBlock) int {
	var totalHeight float64 = float64(2 * defaultMargin) // 上下边距

	for _, block := range blocks {
		if block.Type == "text" {
			// 估算文本高度
			estimatedLines := (len(block.Content)*fontSize)/(defaultWidth-2*defaultMargin) + 1
			totalHeight += float64(estimatedLines) * fontSize * lineHeight
			totalHeight += fontSize // 段落间距
		} else if block.Type == "image" {
			totalHeight += block.ImgHeight + float64(defaultMargin) // 图片高度加间距
		}
	}

	return int(totalHeight)
}