package transaction

import (
	"context"

	"gorm.io/gorm"
)

type txKeyType struct{}

var txKey = &txKeyType{}

type providerImp struct {
	db *gorm.DB

	enableNestedTransaction bool
}

func NewProvider(db *gorm.DB, opts ...ProviderOption) *providerImp {
	p := &providerImp{
		db:                      db,
		enableNestedTransaction: false,
	}

	for _, opt := range opts {
		opt(p)
	}

	return p
}

// ReadDB return a db handler for read operation
func (p *providerImp) DB(ctx context.Context) *gorm.DB {
	transDB := p.getTransDB(ctx)
	if transDB != nil {
		return transDB
	}

	return p.newDB(ctx)
}

// Transaction start a DB transaction
func (p *providerImp) Transaction(ctx context.Context, cb func(ctx context.Context) error) error {
	transDB := p.getTransDB(ctx)

	if transDB == nil {
		wdb := p.newDB(ctx)
		return wdb.Transaction(func(tx *gorm.DB) error {
			newCtx := p.setTransDB(ctx, tx)
			return cb(newCtx)
		})
	}

	// enable transaction
	if p.enableNestedTransaction {
		return transDB.Transaction(func(tx *gorm.DB) error {
			newCtx := p.setTransDB(ctx, tx)
			return cb(newCtx)
		})
	}

	// disable nested transaction
	// call callback directly.
	// We can acquire the TRANSACTIONAL db handler from context within callback
	return cb(ctx)
}

func (p *providerImp) newDB(ctx context.Context) *gorm.DB {
	return p.db.WithContext(ctx)
}

func (p *providerImp) getTransDB(ctx context.Context) *gorm.DB {
	db, ok := ctx.Value(txKey).(*gorm.DB)
	if !ok {
		return nil
	}
	return db
}

func (p *providerImp) setTransDB(ctx context.Context, db *gorm.DB) context.Context {
	newCtx := context.WithValue(ctx, txKey, db)
	return newCtx
}