package database

import (
	"context"
	"errors"
	"fmt"
	"time"

	"git.toki-labs.com/toki/ariadne/internal/domain/organization"
	"github.com/jackc/pgx/v5"
)

// WithinOrganization runs a transaction with PostgreSQL's tenant context set.
// Tenant-owned repositories must use this method instead of a raw pool.
func (database *Database) WithinOrganization(
	ctx context.Context,
	organizationID organization.ID,
	operation func(pgx.Tx) error,
) error {
	if database == nil || database.runtimePool == nil {
		return errors.New("database is not configured")
	}
	if err := organizationID.Validate(); err != nil {
		return fmt.Errorf("organization scope: %w", err)
	}
	if operation == nil {
		return errors.New("organization operation is required")
	}

	transaction, err := database.runtimePool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return fmt.Errorf("begin organization transaction: %w", err)
	}
	defer func() {
		rollbackContext, cancel := context.WithTimeout(
			context.Background(),
			5*time.Second,
		)
		defer cancel()
		_ = transaction.Rollback(rollbackContext)
	}()

	if _, err := transaction.Exec(
		ctx,
		"SELECT set_config('ariadne.organization_id', $1, true)",
		organizationID.String(),
	); err != nil {
		return fmt.Errorf("set organization scope: %w", err)
	}

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