Meshray-Manager/internal/service/backup.go

package service

import (
	"archive/zip"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"gorm.io/gorm"
)

// BackupService 备份服务
type BackupService struct {
	db *gorm.DB
}

// NewBackupService 创建备份服务
func NewBackupService(db *gorm.DB) *BackupService {
	return &BackupService{
		db: db,
	}
}

// CreateBackup 创建系统备份
func (s *BackupService) CreateBackup(ctx context.Context, backupFile string) error {
	// 1. 导出数据库数据到临时文件
	tempDir := filepath.Join("data", "temp_backup")
	if err := os.MkdirAll(tempDir, 0755); err != nil {
		return fmt.Errorf("创建临时目录失败：%w", err)
	}
	defer os.RemoveAll(tempDir)

	// 导出数据库
	dbDumpFile := filepath.Join(tempDir, "meshray.sql")
	if err := s.dumpDatabase(dbDumpFile); err != nil {
		return fmt.Errorf("导出数据库失败：%w", err)
	}

	// 2. 备份配置文件
	configFiles := []string{
		"config.yaml",
	}

	for _, configFile := range configFiles {
		if _, err := os.Stat(configFile); err == nil {
			// 复制配置文件到临时目录
			src, err := os.Open(configFile)
			if err != nil {
				return fmt.Errorf("打开配置文件失败：%w", err)
			}
			defer src.Close()

			dstPath := filepath.Join(tempDir, filepath.Base(configFile))
			dst, err := os.Create(dstPath)
			if err != nil {
				return fmt.Errorf("创建配置文件副本失败：%w", err)
			}
			defer dst.Close()

			if _, err := io.Copy(dst, src); err != nil {
				return fmt.Errorf("复制配置文件失败：%w", err)
			}
		}
	}

	// 3. 打包成 zip 文件
	if err := s.createZipFile(backupFile, tempDir); err != nil {
		return fmt.Errorf("创建压缩文件失败：%w", err)
	}

	return nil
}

// RestoreBackup 恢复备份
func (s *BackupService) RestoreBackup(ctx context.Context, backupFile string) error {
	// 1. 解压备份文件
	tempDir := filepath.Join("data", "temp_restore")
	if err := os.MkdirAll(tempDir, 0755); err != nil {
		return fmt.Errorf("创建临时目录失败：%w", err)
	}
	defer os.RemoveAll(tempDir)

	if err := s.extractZipFile(backupFile, tempDir); err != nil {
		return fmt.Errorf("解压备份文件失败：%w", err)
	}

	// 2. 恢复数据库
	dbDumpFile := filepath.Join(tempDir, "meshray.sql")
	if _, err := os.Stat(dbDumpFile); err == nil {
		if err := s.restoreDatabase(dbDumpFile); err != nil {
			return fmt.Errorf("恢复数据库失败：%w", err)
		}
	}

	// 3. 恢复配置文件
	configFile := filepath.Join(tempDir, "config.yaml")
	if _, err := os.Stat(configFile); err == nil {
		// 备份当前配置
		if _, err := os.Stat("config.yaml"); err == nil {
			os.Rename("config.yaml", "config.yaml.bak."+time.Now().Format("20060102_150405"))
		}

		// 复制新配置
		src, err := os.Open(configFile)
		if err != nil {
			return fmt.Errorf("打开备份配置文件失败：%w", err)
		}
		defer src.Close()

		dst, err := os.Create("config.yaml")
		if err != nil {
			return fmt.Errorf("创建配置文件失败：%w", err)
		}
		defer dst.Close()

		if _, err := io.Copy(dst, src); err != nil {
			return fmt.Errorf("复制配置文件失败：%w", err)
		}
	}

	return nil
}

// dumpDatabase 导出数据库到 SQL 文件
func (s *BackupService) dumpDatabase(outputFile string) error {
	// 使用 SQLite 的 dump 功能
	// 通过 gorm 执行 PRAGMA 和查询来导出所有表结构和数据

	file, err := os.Create(outputFile)
	if err != nil {
		return err
	}
	defer file.Close()

	// 写入注释头
	file.WriteString("-- MeshRay Database Backup\n")
	file.WriteString(fmt.Sprintf("-- Generated at: %s\n\n", time.Now().Format(time.RFC3339)))

	// 获取所有表名
	var tables []string
	if err := s.db.Raw("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'").Scan(&tables).Error; err != nil {
		return err
	}

	// 导出每个表
	for _, table := range tables {
		// 导出表结构
		var createSQL string
		if err := s.db.Raw(fmt.Sprintf("SELECT sql FROM sqlite_master WHERE type='table' AND name='%s'", table)).Scan(&createSQL).Error; err != nil {
			continue
		}

		file.WriteString(fmt.Sprintf("-- Table structure for table `%s`\n", table))
		file.WriteString("DROP TABLE IF EXISTS `" + table + "`;\n")
		file.WriteString(createSQL + ";\n\n")

		// 导出表数据
		var rows []map[string]interface{}
		if err := s.db.Table(table).Find(&rows).Error; err != nil {
			continue
		}

		if len(rows) > 0 {
			file.WriteString(fmt.Sprintf("-- Data for table `%s`\n", table))
			file.WriteString("INSERT INTO `" + table + "` VALUES\n")

			for i, row := range rows {
				values := make([]string, 0)
				for _, v := range row {
					if v == nil {
						values = append(values, "NULL")
					} else {
						values = append(values, fmt.Sprintf("'%v'", v))
					}
				}

				if i < len(rows)-1 {
					file.WriteString("(" + strings.Join(values, ",") + "),\n")
				} else {
					file.WriteString("(" + strings.Join(values, ",") + ");\n\n")
				}
			}
		}
	}

	return nil
}

// restoreDatabase 从 SQL 文件恢复数据库
func (s *BackupService) restoreDatabase(inputFile string) error {
	// 读取 SQL 文件
	content, err := os.ReadFile(inputFile)
	if err != nil {
		return err
	}

	// 简单实现：执行 SQL 语句
	// 生产环境应该使用 SQLite 命令行工具或更完善的 SQL 解析器
	queries := strings.Split(string(content), ";")

	for _, query := range queries {
		query = strings.TrimSpace(query)
		if query == "" || strings.HasPrefix(query, "--") {
			continue
		}

		// 执行 SQL 语句
		if err := s.db.Exec(query).Error; err != nil {
			// 忽略错误（因为可能遇到 DROP TABLE 时表不存在）
			continue
		}
	}

	return nil
}

// createZipFile 创建 ZIP 压缩文件
func (s *BackupService) createZipFile(zipFile, sourceDir string) error {
	file, err := os.Create(zipFile)
	if err != nil {
		return err
	}
	defer file.Close()

	writer := zip.NewWriter(file)
	defer writer.Close()

	return filepath.Walk(sourceDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// 跳过目录本身
		if info.IsDir() {
			return nil
		}

		// 创建 ZIP 中的文件头
		header, err := zip.FileInfoHeader(info)
		if err != nil {
			return err
		}
		header.Name, _ = filepath.Rel(sourceDir, path)
		header.Method = zip.Deflate

		f, err := writer.CreateHeader(header)
		if err != nil {
			return err
		}

		// 读取源文件并写入 ZIP
		srcFile, err := os.Open(path)
		if err != nil {
			return err
		}
		defer srcFile.Close()

		_, err = io.Copy(f, srcFile)
		return err
	})
}

// extractZipFile 解压 ZIP 文件
func (s *BackupService) extractZipFile(zipFile, destDir string) error {
	reader, err := zip.OpenReader(zipFile)
	if err != nil {
		return err
	}
	defer reader.Close()

	for _, file := range reader.File {
		path := filepath.Join(destDir, file.Name)

		// 如果是目录，创建目录
		if file.FileInfo().IsDir() {
			os.MkdirAll(path, 0755)
			continue
		}

		// 创建父目录
		if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
			return err
		}

		// 解压文件
		srcFile, err := file.Open()
		if err != nil {
			return err
		}
		defer srcFile.Close()

		dstFile, err := os.Create(path)
		if err != nil {
			return err
		}
		defer dstFile.Close()

		_, err = io.Copy(dstFile, srcFile)
		if err != nil {
			return err
		}
	}

	return nil
}