AutoZero/imagecache/cache.go

// Package imagecache 实现磁盘图片缓存：历史消息中的图片只保留文件名引用，
// 最新消息的图片才加载 base64 交给大模型处理。
package imagecache

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// Cache 是一个磁盘图片缓存。
type Cache struct {
	dir     string
	maxAge  time.Duration
	maxSize int64
	mu      sync.RWMutex
}

// New 创建图片缓存目录并返回 Cache 实例。
// dir: 缓存目录（如 ~/.autogo/images/）
// maxAge: 图片最大存活时间（0 = 永不过期）
// maxSize: 缓存目录最大字节数（0 = 不限制）
func New(dir string, maxAge time.Duration, maxSize int64) (*Cache, error) {
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, fmt.Errorf("imagecache: create dir: %w", err)
	}
	c := &Cache{dir: dir, maxAge: maxAge, maxSize: maxSize}
	// 启动时清理过期图片
	go c.cleanup()
	return c, nil
}

// DefaultCache 创建默认缓存（~/.autogo/images/，30天过期，1GB限制）。
func DefaultCache() (*Cache, error) {
	home, err := os.UserHomeDir()
	if err != nil {
		home = "."
	}
	dir := filepath.Join(home, ".autogo", "images")
	return New(dir, 30*24*time.Hour, 1<<30)
}

// Store 存储图片数据到缓存，返回文件名（不含路径）。
// ext: 文件扩展名（如 ".jpg", ".png"）
func (c *Cache) Store(data []byte, ext string) (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if ext == "" {
		ext = ".jpg"
	}
	if !strings.HasPrefix(ext, ".") {
		ext = "." + ext
	}

	// 生成唯一文件名
	filename, err := c.generateFilename(ext)
	if err != nil {
		return "", err
	}

	path := filepath.Join(c.dir, filename)
	if err := os.WriteFile(path, data, 0o644); err != nil {
		return "", fmt.Errorf("imagecache: write file: %w", err)
	}

	return filename, nil
}

// StoreReader 从 reader 存储图片，返回文件名。
func (c *Cache) StoreReader(r io.Reader, ext string) (string, error) {
	data, err := io.ReadAll(r)
	if err != nil {
		return "", err
	}
	return c.Store(data, ext)
}

// Load 从缓存加载图片数据。
func (c *Cache) Load(filename string) ([]byte, error) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	path := c.safePath(filename)
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return nil, fmt.Errorf("imagecache: file not found: %s", filename)
	}

	return os.ReadFile(path)
}

// LoadBase64 从缓存加载图片并返回 base64 编码字符串。
func (c *Cache) LoadBase64(filename string) (string, error) {
	data, err := c.Load(filename)
	if err != nil {
		return "", err
	}
	return encodeBase64(data), nil
}

// LoadWithMIME 从缓存加载图片，返回 base64 和 MIME 类型。
func (c *Cache) LoadWithMIME(filename string) (string, string, error) {
	data, err := c.Load(filename)
	if err != nil {
		return "", "", err
	}
	mime := detectMIME(filename)
	return encodeBase64(data), mime, nil
}

// Exists 检查文件是否存在。
func (c *Cache) Exists(filename string) bool {
	c.mu.RLock()
	defer c.mu.RUnlock()

	path := c.safePath(filename)
	_, err := os.Stat(path)
	return err == nil
}

// Delete 删除缓存中的图片。
func (c *Cache) Delete(filename string) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	path := c.safePath(filename)
	return os.Remove(path)
}

// List 列出所有缓存的图片文件名。
func (c *Cache) List() ([]string, error) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	entries, err := os.ReadDir(c.dir)
	if err != nil {
		return nil, err
	}

	var files []string
	for _, e := range entries {
		if !e.IsDir() {
			files = append(files, e.Name())
		}
	}
	return files, nil
}

// Stats 返回缓存统计信息。
func (c *Cache) Stats() (count int, totalSize int64, err error) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	entries, err := os.ReadDir(c.dir)
	if err != nil {
		return 0, 0, err
	}

	for _, e := range entries {
		if e.IsDir() {
			continue
		}
		info, err := e.Info()
		if err != nil {
			continue
		}
		count++
		totalSize += info.Size()
	}
	return count, totalSize, nil
}

// Cleanup 手动触发清理。
func (c *Cache) Cleanup() error {
	return c.cleanup()
}

// generateFilename 生成唯一文件名。
func (c *Cache) generateFilename(ext string) (string, error) {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	timestamp := time.Now().Format("20060102-150405")
	return fmt.Sprintf("%s-%s%s", timestamp, hex.EncodeToString(b), ext), nil
}

// safePath 安全地拼接路径，防止路径遍历攻击。
func (c *Cache) safePath(filename string) string {
	// 清理文件名，移除路径分隔符
	filename = filepath.Base(filename)
	return filepath.Join(c.dir, filename)
}

// cleanup 清理过期和超限的图片。
func (c *Cache) cleanup() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	entries, err := os.ReadDir(c.dir)
	if err != nil {
		return err
	}

	now := time.Now()
	var totalSize int64
	type fileEntry struct {
		name    string
		modTime time.Time
		size    int64
	}
	var files []fileEntry

	// 收集文件信息
	for _, e := range entries {
		if e.IsDir() {
			continue
		}
		info, err := e.Info()
		if err != nil {
			continue
		}
		totalSize += info.Size()
		files = append(files, fileEntry{
			name:    e.Name(),
			modTime: info.ModTime(),
			size:    info.Size(),
		})
	}

	// 删除过期文件
	if c.maxAge > 0 {
		for _, f := range files {
			if now.Sub(f.modTime) > c.maxAge {
				os.Remove(filepath.Join(c.dir, f.name))
				totalSize -= f.size
			}
		}
	}

	// 如果超过大小限制，删除最旧的文件
	if c.maxSize > 0 && totalSize > c.maxSize {
		// 按修改时间排序（最旧在前）
		for i := 0; i < len(files)-1; i++ {
			for j := i + 1; j < len(files); j++ {
				if files[i].modTime.After(files[j].modTime) {
					files[i], files[j] = files[j], files[i]
				}
			}
		}

		for _, f := range files {
			if totalSize <= c.maxSize {
				break
			}
			os.Remove(filepath.Join(c.dir, f.name))
			totalSize -= f.size
		}
	}

	return nil
}

// encodeBase64 编码 base64。
func encodeBase64(data []byte) string {
	return base64.StdEncoding.EncodeToString(data)
}

// detectMIME 根据扩展名检测 MIME 类型。
func detectMIME(filename string) string {
	ext := strings.ToLower(filepath.Ext(filename))
	switch ext {
	case ".jpg", ".jpeg":
		return "image/jpeg"
	case ".png":
		return "image/png"
	case ".gif":
		return "image/gif"
	case ".webp":
		return "image/webp"
	default:
		return "image/jpeg"
	}
}