NetworkAuth/database/database.go

package database

import (
	"NetworkAuth/utils"
	"context"
	"fmt"
	"log"
	"os"
	"sync"
	"time"

	"github.com/glebarez/sqlite"
	"github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	gLogger "gorm.io/gorm/logger"
)

// ============================================================================
// 全局变量
// ============================================================================

var (
	// dbInstance 全局 *gorm.DB 实例，使用单例确保全局复用
	dbInstance *gorm.DB
	// healthCheckCancel 数据库健康检查取消函数
	healthCheckCancel context.CancelFunc
	// once 确保初始化只执行一次
	once sync.Once
)

// ============================================================================
// 公共函数
// ============================================================================

// Init 初始化数据库连接（根据配置自动选择驱动）
// - 默认使用 SQLite（github.com/glebarez/sqlite）
// - 生产环境支持 MySQL（gorm.io/driver/mysql）
func Init() (*gorm.DB, error) {
	var initErr error
	once.Do(func() {
		initErr = performInit()
	})
	return dbInstance, initErr
}

// GetDB 获取全局 *gorm.DB 实例
// 如果未初始化，会尝试初始化一次
func GetDB() (*gorm.DB, error) {
	if dbInstance != nil {
		return dbInstance, nil
	}
	return Init()
}

// ReInit 重新初始化数据库连接
// 用于在修改配置后重新连接数据库
func ReInit() (*gorm.DB, error) {
	// 如果已有连接，尝试关闭它
	if dbInstance != nil {
		if healthCheckCancel != nil {
			healthCheckCancel()
			healthCheckCancel = nil
		}
		if sqlDB, err := dbInstance.DB(); err == nil {
			sqlDB.Close()
		}
	}
	dbInstance = nil

	// 重新执行初始化逻辑（不经过 once.Do）
	return dbInstance, performInit()
}

func performInit() error {
	// 检查是否已经有配置文件（通过检查文件是否存在）
	configFile := viper.ConfigFileUsed()
	// 如果 viper 没有使用配置文件（可能是因为没找到文件而使用了默认配置），
	// 或者配置文件路径为空，我们应该假设处于未安装状态。
	// 但 viper.ConfigFileUsed() 在 ReadInConfig 成功后会返回文件名。
	// 如果 ReadInConfig 失败（因为文件不存在），viper 可能会返回空或者我们在 config.go 中设置的路径。

	// 在 config.go 中，如果文件不存在，我们加载了默认配置但没有写文件。
	// 此时 viper.ConfigFileUsed() 可能是空的或者我们设置的路径。
	// 让我们检查该路径对应的文件是否存在。

	if configFile == "" {
		configFile = "config.json"
	}

	_, err := os.Stat(configFile)
	isConfigExists := !os.IsNotExist(err)

	// 如果配置文件不存在，说明还没有经过安装初始化，暂时不连接数据库
	if !isConfigExists {
		logrus.Info("尚未初始化配置，跳过数据库连接")
		return nil
	}

	var initErr error
	dbType := viper.GetString("database.type")
	switch dbType {
	case "mysql":
		initErr = initMySQL()
	default:
		initErr = initSQLite()
	}

	// 如果数据库初始化成功，配置连接池和启动健康检查
	if initErr == nil && dbInstance != nil {
		// 加载数据库配置
		var configPrefix string
		if dbType == "mysql" {
			configPrefix = "database.mysql"
		} else {
			configPrefix = "database.sqlite"
		}

		dbConfig := utils.LoadDatabaseConfig(configPrefix)

		// 验证配置
		if err := utils.ValidateDatabaseConfig(dbConfig); err != nil {
			logrus.WithError(err).Warn("数据库配置验证失败，使用默认配置")
			dbConfig = utils.GetDefaultDatabaseConfig()
		}

		// 配置连接池
		if err := utils.ConfigureConnectionPool(dbInstance, dbConfig); err != nil {
			logrus.WithError(err).Error("配置数据库连接池失败")
		}

		// 启动健康检查
		healthCheckCancel = utils.StartHealthCheck(dbInstance, dbConfig)
	}
	return initErr
}

// SetDB 设置全局 *gorm.DB 实例（用于测试）
func SetDB(db *gorm.DB) {
	dbInstance = db
}

// ============================================================================
// 私有函数
// ============================================================================

// initSQLite 初始化 SQLite 数据库
// 使用 viper 中的 database.sqlite.path 作为数据库文件路径
func initSQLite() error {
	path := viper.GetString("database.sqlite.path")
	if path == "" {
		path = "./recharge.db"
	}
	dsn := fmt.Sprintf("file:%s?cache=shared&_busy_timeout=5000&_fk=1", path)
	var logLevel gLogger.LogLevel
	switch viper.GetString("logger.level") {
	case "debug":
		logLevel = gLogger.Info
	case "error":
		logLevel = gLogger.Error
	default:
		logLevel = gLogger.Warn
	}
	gl := gLogger.New(log.New(os.Stdout, "\r\n", log.LstdFlags), gLogger.Config{SlowThreshold: 2 * time.Second, LogLevel: logLevel, IgnoreRecordNotFoundError: true, Colorful: false})
	db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{Logger: gl})
	if err != nil {
		logrus.WithError(err).Error("SQLite 初始化失败")
		return err
	}

	// SQLite 连接池配置（SQLite 对连接池支持有限，但仍可设置基本参数）
	if sqlDB, err := db.DB(); err == nil {
		// SQLite 通常使用单连接，但可以设置一些基本参数
		sqlDB.SetMaxOpenConns(1) // SQLite 建议使用单连接
		sqlDB.SetMaxIdleConns(1)
	}

	dbInstance = db
	logrus.WithField("path", path).Info("SQLite 连接已建立")
	return nil
}

// initMySQL 初始化 MySQL 数据库
// 从 viper 读取 database.mysql.* 配置构建 DSN
func initMySQL() error {
	host := viper.GetString("database.mysql.host")
	port := viper.GetInt("database.mysql.port")
	user := viper.GetString("database.mysql.username")
	pass := viper.GetString("database.mysql.password")
	dbname := viper.GetString("database.mysql.database")
	charset := viper.GetString("database.mysql.charset")
	if charset == "" {
		charset = "utf8mb4"
	}

	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=Local", user, pass, host, port, dbname, charset)
	var logLevel gLogger.LogLevel
	switch viper.GetString("logger.level") {
	case "debug":
		logLevel = gLogger.Info
	case "error":
		logLevel = gLogger.Error
	default:
		logLevel = gLogger.Warn
	}
	gl := gLogger.New(log.New(os.Stdout, "\r\n", log.LstdFlags), gLogger.Config{SlowThreshold: 2 * time.Second, LogLevel: logLevel, IgnoreRecordNotFoundError: true, Colorful: false})
	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{Logger: gl})
	if err != nil {
		logrus.WithError(err).Error("MySQL 初始化失败")
		return err
	}

	dbInstance = db
	logrus.WithField("host", host).WithField("database", dbname).Info("MySQL 连接已建立")
	return nil
}