sunhpc-go/cmd/init/database.go

package initcmd

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"sunhpc/internal/auth"
	"sunhpc/internal/db"
	"sunhpc/internal/log"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	appName       string = "sunhpc"
	defaultDBPath string = "/var/lib/sunhpc"
	defaultDBName string = "sunhpc.db"
)

var (
	forceDB    bool
	dbPath     string
	dbName     string
	dbFullPath string
)

func initDBPathWithViper() error {
	/*
		从 Viper 配置文件获取数据库路径
		配置文件里的键要和 Viper.GetXXX 的键对应
		配置文件格式：
		db:
			path: "/tmp/sunhpc"  # 自定义数据库路径
			name: "my_sunhpc.db" # 自定义数据库名
	*/

	log.Infof("从 Viper 配置文件获取数据库路径...")

	// ========== 第一步：设置 Viper 配置文件规则（核心） ==========
	// 1. 设置Viper基础规则
	viper.SetConfigType("yaml")   // 配置文件类型
	viper.SetConfigName("config") // 配置文件名（不带后缀）
	viper.SetEnvPrefix(appName)   // 环境变量前缀：SUNHPC_
	viper.AutomaticEnv()          // 自动读取环境变量（可选，增强兼容性）

	// 2. 添加配置文件搜索目录（Viper 会按顺序查找，找到第一个就停止）
	// 优先级：当前目录 → 用户级目录 → 系统级目录

	// ① 当前目录（开发/测试常用）
	viper.AddConfigPath(".")

	// ② Linux/macOS 用户级目录（~/.config/sunhpc/）
	if homeDir, err := os.UserHomeDir(); err == nil {
		viper.AddConfigPath(filepath.Join(homeDir, ".config", appName))
	}
	// ③ Linux/macOS 系统级目录（/etc/sunhpc/）
	viper.AddConfigPath(filepath.Join("/etc", appName))

	// ========== 第二步：设置默认值（最低优先级） ==========
	viper.SetDefault("db.path", defaultDBPath)
	viper.SetDefault("db.name", defaultDBName)

	// ========== 第三步：绑定环境变量（优先级高于默认值，低于配置文件） ==========
	viper.BindEnv("db.path", "DB_PATH") // 绑定 SUNHPC_DB_PATH → db.path
	viper.BindEnv("db.name", "DB_NAME") // 绑定 SUNHPC_DB_NAME → db.name

	// ========== 第四步：读取配置文件（优先级高于环境变量，低于默认值） ==========
	if err := viper.ReadInConfig(); err != nil {
		if _, ok := err.(viper.ConfigFileNotFoundError); ok {
			log.Info("未找到配置文件,将使用环境变量/默认值")
			return nil // 配置文件存在但格式错误，返回错误
		}
		log.Warnf("读取配置文件失败: %v", err)
		return fmt.Errorf("读取配置文件失败: %w", err)
	}

	log.Infof("成功加载配置文件: %s", viper.ConfigFileUsed())
	return nil
}

func initDBPath() error {

	// 1. 从 Viper 配置文件获取数据库路径(加载配置文件->环境变量->默认值)
	if err := initDBPathWithViper(); err != nil {
		return fmt.Errorf("Viper初始化数据库失败: %v", err)
	}

	// 2. 从Viper获取数据库路径
	dbPath = viper.GetString("db.path")
	dbName = viper.GetString("db.name")

	// 3. 拼接数据库路径
	dbFullPath = filepath.Join(dbPath, dbName)
	log.Infof("数据库完整路径: %s", dbFullPath)

	// 3. 检查数据库文件是否存在
	dir := filepath.Dir(dbFullPath)
	// 4. 检查数据库目录是否存在
	if _, err := os.Stat(dir); os.IsNotExist(err) {
		log.Infof("数据库目录不存在,创建目录: %s", dir)

		if err := os.MkdirAll(dir, 0755); err != nil {
			return fmt.Errorf("创建数据库目录失败: %v", err)
		}
	}
	return nil
}

var databaseCmd = &cobra.Command{
	Use:   "database",
	Short: "初始化数据库",
	Long: `初始化SQLite数据库,创建所有表结构和默认数据。
    
示例:
  sunhpc init database        # 初始化数据库
  sunhpc init database --force # 强制重新初始化`,

	RunE: func(cmd *cobra.Command, args []string) error {
		if err := auth.RequireRoot(); err != nil {
			return err
		}

		log.Debug("初始化数据库...")

		if err := initDBPath(); err != nil {
			return fmt.Errorf("初始化数据库路径失败: %w", err)
		}

		log.Debugf("数据库目录存在: %s", dbPath)

		// 强制模式：用户确认
		if forceDB {
			log.Warn("⚠️  警告：强制重新初始化将清空数据库中的所有数据！")
			fmt.Printf("数据库路径: %s\n", dbFullPath)
			fmt.Print("确认要重新初始化数据库吗？这将删除所有现有数据。(Y/yes): ")

			reader := bufio.NewReader(os.Stdin)
			input, err := reader.ReadString('\n')
			if err != nil {
				return fmt.Errorf("读取用户输入失败: %v", err)
			}

			input = strings.TrimSpace(strings.ToLower(input))
			if input != "y" && input != "yes" {
				log.Info("操作已取消")
				return nil
			}

			log.Info("用户确认重新初始化数据库")
		}

		// 数据库存在且不是强制模式则跳过初始化
		if _, err := os.Stat(dbFullPath); err == nil && !forceDB {
			log.Infof("数据库文件已存在: %s", dbFullPath)
			return nil
		}

		// 初始化数据库（使用配置的路径）
		database, err := db.GetInstanceWithConfig(dbPath, dbName)
		if err != nil {
			return fmt.Errorf("初始化数据库失败: %v", err)
		}
		defer database.Close()

		// 如果是强制模式，设置强制重新初始化标志
		if forceDB {
			database.SetForceInit(true)
			log.Info("强制重新初始化数据库表...")

			// 关闭现有连接以触发重新连接
			if err := database.CloseConnection(); err != nil {
				return fmt.Errorf("关闭现有数据库连接失败: %v", err)
			}

			// 重新连接并初始化
			if err := database.Connect(); err != nil {
				return fmt.Errorf("强制重新初始化数据库失败: %v", err)
			}
		}

		log.Infof("数据库初始化成功: %s", dbFullPath)
		return nil
	},
}

func init() {
	databaseCmd.Flags().BoolVarP(&forceDB, "force", "f", false, "强制重新初始化，删除现有数据库")
	Cmd.AddCommand(databaseCmd)
}