AntiCaptcha/internal/database/database.go

package database

import (
	"fmt"
	"log"

	"anticaptcha/internal/config"
	"anticaptcha/internal/model"

	"gorm.io/driver/mysql"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

var DB *gorm.DB

func Init(cfg *config.DatabaseConfig) error {
	var err error
	var gormConfig *gorm.Config

	// 生产环境禁用日志
	gormConfig = &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
	}

	switch cfg.Type {
	case "sqlite":
		DB, err = gorm.Open(sqlite.Open(cfg.SQLite.Path), gormConfig)
	case "mysql":
		dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Local",
			cfg.User, cfg.Password, cfg.Host, cfg.Port, cfg.Database)
		DB, err = gorm.Open(mysql.Open(dsn), gormConfig)
	default:
		return fmt.Errorf("不支持的数据库类型: %s", cfg.Type)
	}

	if err != nil {
		return fmt.Errorf("数据库连接失败: %v", err)
	}

	// 自动迁移
	if err = DB.AutoMigrate(
		&model.User{},
		&model.RegistrationCode{},
		&model.EndpointCost{},
		&model.Config{},
	); err != nil {
		return fmt.Errorf("数据库迁移失败: %v", err)
	}

	return nil
}

func Close() error {
	if DB != nil {
		sqlDB, err := DB.DB()
		if err != nil {
			return err
		}
		return sqlDB.Close()
	}
	return nil
}

// 确保 admin 用户存在
func EnsureAdminUser() error {
	var count int64
	DB.Model(&model.User{}).Where("role = ?", model.RoleAdmin).Count(&count)
	if count > 0 {
		return nil
	}

	admin := &model.User{
		Username:       "admin",
		HashedPassword: HashPassword("admin"),
		Role:           model.RoleAdmin,
		Balance:        1000000,
	}

	return DB.Create(admin).Error
}

func HashPassword(password string) string {
	// 简化处理，实际应使用 bcrypt
	// 后续补充完整实现
	return password
}