package data

import (
	"center-api/internal/conf"
	"center-api/internal/data/ent"
	"center-api/internal/data/ent/intercept"
	"context"
	ent2 "entgo.io/ent"
	"entgo.io/ent/dialect/sql"
	"fmt"
	"github.com/go-kratos/kratos/v2/log"
	_ "github.com/go-sql-driver/mysql"
)

// ContextTxKey 用来承载事务的上下文
type ContextTxKey struct{}

// Db .
type Db struct {
	Client *ent.Client
}

// NewDb .
func NewDb(c *conf.Bootstrap, hLog *log.Helper) (*Db, func(), error) {
	client, err := buildDbClient(c.Data.Db)

	cleanup := func() {
		if client != nil {
			if err := client.Close(); err != nil {
				hLog.Error("关闭 masterClient 失败：", err)
			}
		}
	}

	return &Db{Client: client}, cleanup, err
}

func (d Db) GetDb(ctx context.Context) *ent.Client {
	tx, ok := ctx.Value(ContextTxKey{}).(*ent.Tx)
	if ok {
		return tx.Client()
	}
	return d.Client
}

// buildDbClient 构建db client
func buildDbClient(dbConf *conf.Data_Database) (*ent.Client, error) {
	db, err := sql.Open(dbConf.GetDriver(), dbConf.GetSource())
	if err != nil {
		fmt.Println("连接db 失败: ", err)
		return nil, err
	}
	db.DB().SetMaxIdleConns(int(dbConf.MaxIdle))
	db.DB().SetMaxOpenConns(int(dbConf.MaxOpen))
	db.DB().SetConnMaxLifetime(dbConf.MaxLifetime.AsDuration())
	var client *ent.Client
	client = ent.NewClient(ent.Driver(db))

	//关闭默认启用的 select distinct
	client.Intercept(
		intercept.Func(func(ctx context.Context, q intercept.Query) error {
			// Skip setting Unique if the modifier was set explicitly.
			if ent2.QueryFromContext(ctx).Unique == nil {
				q.Unique(false)
			}
			return nil
		}),
	)
	if dbConf.IsDebug {
		client = client.Debug()
	}
	return client, nil
}