// Package repository add by blacknc <blacknc@163.com> at 2023/03/01 15:48:36
package repository

import (
	"context"
	"dataservices/commons/config"
	"dataservices/commons/database"

	"gorm.io/gorm"
)

type contextKeyType int

const (
	contextDBKey contextKeyType = iota
)

func ContextWithDB(ctx context.Context, db *gorm.DB) context.Context {
	return context.WithValue(ctx, contextDBKey, db)
}

func FromContext(ctx context.Context) (gormDB *gorm.DB) {
	defer func() {
		if gormDB != nil && !config.IsProd() {
			gormDB = gormDB.Debug()
		}
	}()
	if db, ok := ctx.Value(contextDBKey).(*gorm.DB); ok {
		return db
	}
	return nil
}

func MustFromContext(ctx context.Context) *gorm.DB {
	if db := FromContext(ctx); db != nil {
		return db
	}
	if config.IsProd() {
		return database.DBConn.DB
	} else {
		return database.DBConn.DB.Debug()
	}
}

func Transaction(ctx context.Context, fc func(ctx context.Context) error) (err error) {

	if FromContext(ctx) != nil {
		return fc(ctx)
	}

	panicked := true

	tx := database.DBConn.DB.Begin()
	if tx.Error != nil {
		return tx.Error
	}

	defer func() {
		// Make sure to rollback when panic, Block error or Commit error
		if panicked || err != nil {
			tx.Rollback()
		}
	}()

	if err = fc(ContextWithDB(ctx, tx)); err == nil {
		panicked = false
		return tx.Commit().Error
	}

	panicked = false
	return
}