package storage

import (
	"context"

	"github.com/nose7en/ToyBoomServer/common"
	"gorm.io/gorm"
)

const (
	DefaultDBName = "default"
)

type DBManager interface {
	GetDB(dbType, dbName string) *gorm.DB
	GetDefaultDB() *gorm.DB
	SetDB(dbType, dbName string, db *gorm.DB)
	RemoveDB(dbType, dbName string)
	Init(tables ...interface{})
}

type dbManagerImpl struct {
	DBs           map[string]map[string]*gorm.DB // [key: dbname, value: dbs] => [key: dbtype value: db]
	defaultDBType string
}

func (dbm *dbManagerImpl) Init(tables ...interface{}) {
	dbs := dbm.DBs[DefaultDBName]
	for _, db := range dbs {
		if err := db.AutoMigrate(tables...); err != nil {
			common.Logger(context.Background()).WithError(err).Fatalf("auto migrate error!!!")
		}
	}
}

var (
	dbm *dbManagerImpl
)

func MustInitDBManager(dbs map[string]map[string]*gorm.DB, defaultDBType string) {
	if dbm == nil {
		dbm = NewDBManager(dbs, defaultDBType)
	}
}

func NewDBManager(dbs map[string]map[string]*gorm.DB, defaultDBType string) *dbManagerImpl {
	if dbs == nil {
		dbs = make(map[string]map[string]*gorm.DB)
	}
	return &dbManagerImpl{
		DBs:           dbs,
		defaultDBType: defaultDBType,
	}
}

func GetDBManager() DBManager {
	if dbm == nil {
		dbm = NewDBManager(nil, "")
	}
	return dbm
}

func (dbm *dbManagerImpl) GetDB(dbType, dbName string) *gorm.DB {
	return dbm.DBs[dbName][dbType]
}

func (dbm *dbManagerImpl) SetDB(dbType, dbName string, db *gorm.DB) {
	if dbm.DBs[dbName] == nil {
		dbm.DBs[dbName] = make(map[string]*gorm.DB)
	}
	dbm.DBs[dbName][dbType] = db
}

func (dbm *dbManagerImpl) RemoveDB(dbType, dbName string) {
	delete(dbm.DBs[dbName], dbType)
}

func (dbm *dbManagerImpl) GetDefaultDB() *gorm.DB {
	return dbm.DBs[DefaultDBName][dbm.defaultDBType]
}