package store

import (
	"context"
	"errors"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
)

const transactionRollbackTimeout = 5 * time.Second

type Tx interface {
	Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}

func (s *Store) InTx(ctx context.Context, fn func(Tx) error) error {
	return s.beginTransaction(ctx, func(tx pgx.Tx) error {
		return fn(tx)
	})
}

func (s *Store) beginTransaction(ctx context.Context, fn func(pgx.Tx) error) (err error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return err
	}
	// A Worker shutdown cancels ctx before in-flight callbacks return. Use an
	// independent bounded context so rollback still reaches PostgreSQL.
	defer func() {
		rollbackCtx, cancel := context.WithTimeout(context.Background(), transactionRollbackTimeout)
		defer cancel()
		rollbackErr := tx.Rollback(rollbackCtx)
		if err == nil && rollbackErr != nil && !errors.Is(rollbackErr, pgx.ErrTxClosed) {
			err = rollbackErr
		}
	}()
	if err = fn(tx); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func rollbackTransaction(tx pgx.Tx) {
	if tx == nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), transactionRollbackTimeout)
	defer cancel()
	_ = tx.Rollback(ctx)
}