package db

import (
	"errors"
	"gitee.com/h79/goutils/dao/config"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlserver"
	"gorm.io/plugin/dbresolver"
	"time"

	"gorm.io/gorm"
	"runtime"
)

type DialFunc func(string) gorm.Dialector

var openFuncs = map[string]DialFunc{
	"mysql":     mysql.Open,
	"postgres":  postgres.Open,
	"sqlserver": sqlserver.Open,
}

var _ Sql = (*Adapter)(nil)

// Adapter represents the Gorm adapter for policy storage.
type Adapter struct {
	driverName     string
	dataSourceName string
	databaseName   string
	db             *gorm.DB
}

type ScopesFunc func(db *gorm.DB) *gorm.DB

// finalizer is the destructor for Adapter.
func finalizer(a *Adapter) {
	sqlDB, err := a.db.DB()
	if err != nil {
		panic(err)
	}
	err = sqlDB.Close()
	if err != nil {
		panic(err)
	}
}

// NewAdapter is the constructor for Adapter.
func NewAdapter(cfg *config.Sql) (*Adapter, error) {
	a := &Adapter{}
	a.driverName = cfg.Master.DriverType
	a.databaseName = cfg.Master.Name
	a.dataSourceName = cfg.Master.GetDataSourceName()

	// Open the DB
	db, err := openDB(a.driverName, a.dataSourceName)
	if err != nil {
		return nil, err
	}
	var sources []gorm.Dialector
	var replicas []gorm.Dialector
	for _, source := range cfg.Sources {
		dr, er := getDriver(source.DriverType, source.GetDataSourceName())
		if er != nil {
			return nil, er
		}
		sources = append(sources, dr)
	}
	for _, replica := range cfg.Replicas {
		dr, er := getDriver(replica.DriverType, replica.GetDataSourceName())
		if er != nil {
			return nil, er
		}
		replicas = append(replicas, dr)
	}
	resolver := dbresolver.Register(dbresolver.Config{
		Sources:  sources,
		Replicas: replicas,
		// sources/replicas load balancing policy
		Policy: dbresolver.RandomPolicy{},
	})
	if cfg.MaxOpenConns > 0 {
		resolver.SetMaxOpenConns(cfg.MaxOpenConns)
	}
	if cfg.MaxIdleConns > 0 {
		resolver.SetMaxIdleConns(cfg.MaxIdleConns)
	}
	if cfg.MaxLifetime > 0 {
		resolver.SetConnMaxLifetime(cfg.MaxLifetime)
	}
	if cfg.MaxIdleTime > 0 {
		resolver.SetConnMaxIdleTime(time.Minute * cfg.MaxIdleTime)
	}
	if cfg.Logger.LogLevel > 1 {
		if cfg.Logger.SlowThreshold <= 0 {
			cfg.Logger.SlowThreshold = 200
		}
		cfg.Logger.SlowThreshold = cfg.Logger.SlowThreshold * time.Millisecond
		db.Logger = &Logger{
			SqlLogger: cfg.Logger,
		}
	}
	if err = db.Use(resolver); err != nil {
		return nil, err
	}

	a.db = db

	// Call the destructor when the object is released.
	runtime.SetFinalizer(a, finalizer)

	return a, nil
}

func (a *Adapter) Db() *gorm.DB {
	return a.db
}

func (a *Adapter) Name() string {
	return a.databaseName
}

func (a *Adapter) Close() {
	a.db = nil
}

func AddDriver(driverName string, dial DialFunc) {
	openFuncs[driverName] = dial
}

func getDriver(driverName, dataSourceName string) (gorm.Dialector, error) {
	driver, ok := openFuncs[driverName]
	if !ok {
		return nil, errors.New("database dialect is not supported")
	}
	return driver(dataSourceName), nil
}

func openDB(driverName, dataSourceName string) (*gorm.DB, error) {
	dr, err := getDriver(driverName, dataSourceName)
	if err != nil {
		return nil, err
	}
	return gorm.Open(dr, &gorm.Config{})
}