package db

import (
	"context"
	"fmt"
	db "server/internal/db/sqlc"

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

type Store struct {
	*db.Queries
	pool *pgxpool.Pool
}

func NewStore(pool *pgxpool.Pool) *Store {
	return &Store{
		pool:    pool,
		Queries: db.New(pool),
	}
}

func (store *Store) WithTx(ctx context.Context, fn func(q *db.Queries) error, opts ...pgx.TxOptions) error {
	txOpt := pgx.TxOptions{}
	if len(opts) > 0 {
		txOpt = opts[0]
	}

	tx, err := store.pool.BeginTx(ctx, txOpt)
	if err != nil {
		return fmt.Errorf("begin tx: %w", err)
	}

	defer func() {
		// 使用 context.Background() 确保即使请求 context 被取消，
		// 回滚指令也能发送给 DB，从而安全地回收连接而不是销毁连接。
		_ = tx.Rollback(context.Background())
	}()

	q := store.Queries.WithTx(tx)

	if err := fn(q); err != nil {
		return err // 业务侧的错误保持原样返回
	}

	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("commit tx: %w", err)
	}

	return nil
}

// WithTxResult 带返回值的事务
func WithTxResult[T any](ctx context.Context, store *Store, fn func(q *db.Queries) (T, error), opts ...pgx.TxOptions) (T, error) {
	var result T

	err := store.WithTx(ctx, func(q *db.Queries) error {
		r, err := fn(q)
		if err != nil {
			return err
		}
		result = r
		return nil
	}, opts...)

	return result, err
}