package dbx

import (
	"context"
	"database/sql"
)

// WithTx begins a transaction on db, runs fn, and commits on success.
// The transaction is rolled back if fn returns an error or commit fails.
func WithTx(ctx context.Context, db QueryExecTx, fn func(tx *sql.Tx) error) error {
	if err := ctx.Err(); err != nil {
		return err
	}

	tx, err := db.Begin()
	if err != nil {
		return err
	}

	if err := fn(tx); err != nil {
		_ = tx.Rollback()
		return err
	}

	return tx.Commit()
}

// WithTxValue begins a transaction on db, runs fn, and commits on success.
// Returns the value from fn or zero and an error if begin, fn, or commit fails.
// The transaction is rolled back if fn returns an error.
func WithTxValue[T any](ctx context.Context, db QueryExecTx, fn func(tx QueryExec) (T, error)) (T, error) {
	var zero T
	if err := ctx.Err(); err != nil {
		return zero, err
	}

	tx, err := db.Begin()
	if err != nil {
		return zero, err
	}

	result, err := fn(tx)
	if err != nil {
		_ = tx.Rollback()
		return zero, err
	}

	if err := tx.Commit(); err != nil {
		return zero, err
	}

	return result, nil
}