yoyo/internal/cache/sqlite.go

package cache

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

// SQLiteCache SQLite缓存实现
type SQLiteCache struct {
	db         *sql.DB
	config     *CacheConfig
	cleanupTTL time.Duration
}

// NewSQLiteCache 创建SQLite缓存实例
func NewSQLiteCache(config *CacheConfig) (*SQLiteCache, error) {
	if config == nil {
		config = NewCacheConfig()
	}

	// 展开路径中的~符号
	dbPath, err := expandPath(config.DBPath)
	if err != nil {
		return nil, fmt.Errorf("无效的数据库路径: %w", err)
	}

	// 确保目录存在
	dir := filepath.Dir(dbPath)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, fmt.Errorf("创建缓存目录失败: %w", err)
	}

	// 打开数据库连接
	db, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_synchronous=NORMAL")
	if err != nil {
		return nil, fmt.Errorf("打开数据库失败: %w", err)
	}

	// 设置连接池参数
	db.SetMaxOpenConns(1) // SQLite只支持单个写入连接
	db.SetMaxIdleConns(1)

	cache := &SQLiteCache{
		db:         db,
		config:     config,
		cleanupTTL: time.Duration(config.ExpireDays) * 24 * time.Hour,
	}

	// 初始化数据库表
	if err := cache.initTable(); err != nil {
		db.Close()
		return nil, fmt.Errorf("初始化缓存表失败: %w", err)
	}

	// 设置清理定时器
	go cache.startCleanupTimer()

	return cache, nil
}

// initTable 初始化缓存表
func (c *SQLiteCache) initTable() error {
	query := `
	CREATE TABLE IF NOT EXISTS translation_cache (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		cache_key TEXT NOT NULL UNIQUE,
		original_text TEXT NOT NULL,
		translated_text TEXT NOT NULL,
		from_lang TEXT NOT NULL,
		to_lang TEXT NOT NULL,
		model TEXT NOT NULL,
		prompt_name TEXT,
		prompt_content TEXT,
		prompt_tokens INTEGER DEFAULT 0,
		completion_tokens INTEGER DEFAULT 0,
		total_tokens INTEGER DEFAULT 0,
		created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
		last_used_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
	);
	
	CREATE INDEX IF NOT EXISTS idx_cache_key ON translation_cache(cache_key);
	CREATE INDEX IF NOT EXISTS idx_original_text ON translation_cache(original_text);
	CREATE INDEX IF NOT EXISTS idx_created_at ON translation_cache(created_at);
	CREATE INDEX IF NOT EXISTS idx_last_used_at ON translation_cache(last_used_at);
	`

	_, err := c.db.Exec(query)
	return err
}

// Get 获取缓存
func (c *SQLiteCache) Get(ctx context.Context, key string) (*CacheEntry, error) {
	query := `
	SELECT 
		id, cache_key, original_text, translated_text, from_lang, to_lang,
		model, prompt_name, prompt_content, prompt_tokens, completion_tokens,
		total_tokens, created_at, last_used_at
	FROM translation_cache
	WHERE cache_key = ?
	`

	entry := &CacheEntry{}
	var promptName, promptContent sql.NullString
	var createdAt, lastUsedAt string

	err := c.db.QueryRowContext(ctx, query, key).Scan(
		&entry.ID, &entry.CacheKey, &entry.OriginalText, &entry.TranslatedText,
		&entry.FromLang, &entry.ToLang, &entry.Model, &promptName, &promptContent,
		&entry.PromptTokens, &entry.CompletionTokens, &entry.TotalTokens,
		&createdAt, &lastUsedAt,
	)

	if err == sql.ErrNoRows {
		return nil, nil // 缓存未命中
	}
	if err != nil {
		return nil, fmt.Errorf("查询缓存失败: %w", err)
	}

	// 处理可空字段
	if promptName.Valid {
		entry.PromptName = promptName.String
	}
	if promptContent.Valid {
		entry.PromptContent = promptContent.String
	}

	// 解析时间
	entry.CreatedAt, _ = time.Parse("2006-01-02 15:04:05", createdAt)
	entry.LastUsedAt, _ = time.Parse("2006-01-02 15:04:05", lastUsedAt)

	// 更新最后使用时间
	go c.updateLastUsed(context.Background(), key)

	return entry, nil
}

// Set 设置缓存
func (c *SQLiteCache) Set(ctx context.Context, entry *CacheEntry) error {
	// 开始事务
	tx, err := c.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("开始事务失败: %w", err)
	}
	defer tx.Rollback()

	// 插入或替换缓存
	query := `
	INSERT OR REPLACE INTO translation_cache 
	(cache_key, original_text, translated_text, from_lang, to_lang, 
	model, prompt_name, prompt_content, prompt_tokens, completion_tokens, 
	total_tokens, created_at, last_used_at)
	VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`

	now := time.Now().Format("2006-01-02 15:04:05")
	_, err = tx.ExecContext(ctx, query,
		entry.CacheKey, entry.OriginalText, entry.TranslatedText,
		entry.FromLang, entry.ToLang, entry.Model, entry.PromptName,
		entry.PromptContent, entry.PromptTokens, entry.CompletionTokens,
		entry.TotalTokens, now, now,
	)
	if err != nil {
		return fmt.Errorf("插入缓存失败: %w", err)
	}

	// 提交事务
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("提交事务失败: %w", err)
	}

	// 触发清理（异步）
	go c.Cleanup(context.Background())

	return nil
}

// Delete 删除缓存
func (c *SQLiteCache) Delete(ctx context.Context, key string) error {
	query := `DELETE FROM translation_cache WHERE cache_key = ?`
	_, err := c.db.ExecContext(ctx, query, key)
	if err != nil {
		return fmt.Errorf("删除缓存失败: %w", err)
	}
	return nil
}

// Clear 清空缓存
func (c *SQLiteCache) Clear(ctx context.Context) error {
	// 先删除所有记录
	_, err := c.db.ExecContext(ctx, `DELETE FROM translation_cache`)
	if err != nil {
		return fmt.Errorf("清空缓存失败: %w", err)
	}

	// 然后执行VACUUM（不能在事务中执行）
	_, err = c.db.ExecContext(ctx, `VACUUM`)
	if err != nil {
		return fmt.Errorf("清理数据库失败: %w", err)
	}

	return nil
}

// Stats 获取缓存统计信息
func (c *SQLiteCache) Stats(ctx context.Context) (*CacheStats, error) {
	stats := &CacheStats{}

	// 获取总记录数
	err := c.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM translation_cache`).Scan(&stats.TotalRecords)
	if err != nil {
		return nil, fmt.Errorf("查询缓存统计失败: %w", err)
	}

	// 如果没有记录，直接返回
	if stats.TotalRecords == 0 {
		return stats, nil
	}

	// 获取时间范围和平均tokens
	var oldestStr, newestStr sql.NullString
	var avgTokens sql.NullFloat64
	err = c.db.QueryRowContext(ctx, `
	SELECT 
		MIN(created_at), 
		MAX(created_at),
		AVG(total_tokens)
	FROM translation_cache
	`).Scan(&oldestStr, &newestStr, &avgTokens)
	if err != nil {
		return nil, fmt.Errorf("查询缓存时间范围失败: %w", err)
	}

	// 解析时间字符串
	if oldestStr.Valid {
		stats.OldestRecord, _ = time.Parse("2006-01-02 15:04:05", oldestStr.String)
	}
	if newestStr.Valid {
		stats.NewestRecord, _ = time.Parse("2006-01-02 15:04:05", newestStr.String)
	}
	if avgTokens.Valid {
		stats.AvgTokensPerRecord = avgTokens.Float64
	}

	// 计算数据库文件大小
	dbPath, _ := expandPath(c.config.DBPath)
	if info, err := os.Stat(dbPath); err == nil {
		stats.TotalSizeBytes = info.Size()
	}

	return stats, nil
}

// Cleanup 清理过期缓存
func (c *SQLiteCache) Cleanup(ctx context.Context) error {
	tx, err := c.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("开始事务失败: %w", err)
	}
	defer tx.Rollback()

	// 清理过期缓存
	if c.cleanupTTL > 0 {
		expiredTime := time.Now().Add(-c.cleanupTTL).Format("2006-01-02 15:04:05")
		_, err = tx.ExecContext(ctx, `DELETE FROM translation_cache WHERE last_used_at < ?`, expiredTime)
		if err != nil {
			return fmt.Errorf("清理过期缓存失败: %w", err)
		}
	} else if c.cleanupTTL == 0 {
		// 如果过期时间为0，清理所有记录
		_, err = tx.ExecContext(ctx, `DELETE FROM translation_cache`)
		if err != nil {
			return fmt.Errorf("清理所有缓存失败: %w", err)
		}
	}

	// 清理超出数量限制的缓存
	if c.config.MaxRecords > 0 {
		_, err = tx.ExecContext(ctx, `
		DELETE FROM translation_cache 
		WHERE id NOT IN (
			SELECT id FROM translation_cache 
			ORDER BY last_used_at DESC 
			LIMIT ?
		)
		`, c.config.MaxRecords)
		if err != nil {
			return fmt.Errorf("清理超出数量限制的缓存失败: %w", err)
		}
	}

	return tx.Commit()
}

// Close 关闭缓存
func (c *SQLiteCache) Close() error {
	if c.db != nil {
		return c.db.Close()
	}
	return nil
}

// updateLastUsed 更新最后使用时间
func (c *SQLiteCache) updateLastUsed(ctx context.Context, key string) {
	query := `UPDATE translation_cache SET last_used_at = ? WHERE cache_key = ?`
	now := time.Now().Format("2006-01-02 15:04:05")
	c.db.ExecContext(ctx, query, now, key)
}

// startCleanupTimer 启动清理定时器
func (c *SQLiteCache) startCleanupTimer() {
	ticker := time.NewTicker(1 * time.Hour) // 每小时清理一次
	defer ticker.Stop()

	for range ticker.C {
		c.Cleanup(context.Background())
	}
}

// expandPath 展开路径中的~符号
func expandPath(path string) (string, error) {
	if strings.HasPrefix(path, "~") {
		home, err := os.UserHomeDir()
		if err != nil {
			return "", err
		}
		path = filepath.Join(home, path[1:])
	}
	return path, nil
}