registry/internal/database/postgres.go

package database

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strings"
	"time"

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

	apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0"
	"github.com/modelcontextprotocol/registry/pkg/model"
)

// PostgreSQL is an implementation of the Database interface using PostgreSQL
type PostgreSQL struct {
	pool *pgxpool.Pool
}

// Executor is an interface for executing queries (satisfied by both pgx.Tx and pgxpool.Pool)
type Executor interface {
	Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
	Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}

// getExecutor returns the appropriate executor (transaction or pool)
func (db *PostgreSQL) getExecutor(tx pgx.Tx) Executor {
	if tx != nil {
		return tx
	}
	return db.pool
}

// NewPostgreSQL creates a new instance of the PostgreSQL database
func NewPostgreSQL(ctx context.Context, connectionURI string) (*PostgreSQL, error) {
	// Parse connection config for pool settings
	config, err := pgxpool.ParseConfig(connectionURI)
	if err != nil {
		return nil, fmt.Errorf("failed to parse PostgreSQL config: %w", err)
	}

	// Configure pool for stability-focused defaults
	config.MaxConns = 30                      // Handle good concurrent load
	config.MinConns = 5                       // Keep connections warm for fast response
	config.MaxConnIdleTime = 30 * time.Minute // Keep connections available for bursts
	config.MaxConnLifetime = 2 * time.Hour    // Refresh connections regularly for stability

	// Create connection pool with configured settings
	pool, err := pgxpool.NewWithConfig(ctx, config)
	if err != nil {
		return nil, fmt.Errorf("failed to create PostgreSQL pool: %w", err)
	}

	// Test the connection
	if err = pool.Ping(ctx); err != nil {
		return nil, fmt.Errorf("failed to ping PostgreSQL: %w", err)
	}

	// Run migrations using a single connection from the pool
	conn, err := pool.Acquire(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to acquire connection for migrations: %w", err)
	}
	defer conn.Release()

	migrator := NewMigrator(conn.Conn())
	if err := migrator.Migrate(ctx); err != nil {
		return nil, fmt.Errorf("failed to run database migrations: %w", err)
	}

	return &PostgreSQL{
		pool: pool,
	}, nil
}

// buildFilterConditions constructs WHERE clause conditions from a ServerFilter
//
//nolint:unparam // argIndex is always 1 currently but kept for API flexibility
func buildFilterConditions(filter *ServerFilter, argIndex int) ([]string, []any, int) {
	var conditions []string
	var args []any

	if filter == nil {
		return conditions, args, argIndex
	}

	if filter.Name != nil {
		conditions = append(conditions, fmt.Sprintf("server_name = $%d", argIndex))
		args = append(args, *filter.Name)
		argIndex++
	}
	if filter.RemoteURL != nil {
		conditions = append(conditions, fmt.Sprintf("EXISTS (SELECT 1 FROM jsonb_array_elements(value->'remotes') AS remote WHERE remote->>'url' = $%d)", argIndex))
		args = append(args, *filter.RemoteURL)
		argIndex++
	}
	if filter.UpdatedSince != nil {
		conditions = append(conditions, fmt.Sprintf("updated_at > $%d", argIndex))
		args = append(args, *filter.UpdatedSince)
		argIndex++
	}
	if filter.SubstringName != nil {
		conditions = append(conditions, fmt.Sprintf("server_name ILIKE $%d", argIndex))
		args = append(args, "%"+*filter.SubstringName+"%")
		argIndex++
	}
	if filter.Version != nil {
		conditions = append(conditions, fmt.Sprintf("version = $%d", argIndex))
		args = append(args, *filter.Version)
		argIndex++
	}
	if filter.IsLatest != nil {
		conditions = append(conditions, fmt.Sprintf("is_latest = $%d", argIndex))
		args = append(args, *filter.IsLatest)
		argIndex++
	}
	if filter.IncludeDeleted == nil || !*filter.IncludeDeleted {
		conditions = append(conditions, "status != 'deleted'")
	}

	return conditions, args, argIndex
}

// addCursorCondition adds pagination cursor condition to WHERE clause
func addCursorCondition(cursor string, argIndex int) (string, []any, int) {
	if cursor == "" {
		return "", nil, argIndex
	}

	// Parse cursor format: "serverName:version"
	parts := strings.SplitN(cursor, ":", 2)
	if len(parts) == 2 {
		cursorServerName := parts[0]
		cursorVersion := parts[1]
		// Use compound condition: (server_name > cursor_name) OR (server_name = cursor_name AND version > cursor_version)
		condition := fmt.Sprintf("(server_name > $%d OR (server_name = $%d AND version > $%d))", argIndex, argIndex+1, argIndex+2)
		return condition, []any{cursorServerName, cursorServerName, cursorVersion}, argIndex + 3
	}

	// Fallback for malformed cursor - treat as server name only for backwards compatibility
	condition := fmt.Sprintf("server_name > $%d", argIndex)
	return condition, []any{cursor}, argIndex + 1
}

func (db *PostgreSQL) ListServers(
	ctx context.Context,
	tx pgx.Tx,
	filter *ServerFilter,
	cursor string,
	limit int,
) ([]*apiv0.ServerResponse, string, error) {
	if limit <= 0 {
		limit = 10
	}

	if ctx.Err() != nil {
		return nil, "", ctx.Err()
	}

	// Build WHERE clause conditions
	argIndex := 1
	whereConditions, args, argIndex := buildFilterConditions(filter, argIndex)

	// Add cursor pagination
	cursorCondition, cursorArgs, argIndex := addCursorCondition(cursor, argIndex)
	if cursorCondition != "" {
		whereConditions = append(whereConditions, cursorCondition)
		args = append(args, cursorArgs...)
	}
	_ = argIndex // Silence unused variable warning

	// Build the WHERE clause
	whereClause := ""
	if len(whereConditions) > 0 {
		whereClause = "WHERE " + strings.Join(whereConditions, " AND ")
	}

	// Query servers table with hybrid column/JSON data
	query := fmt.Sprintf(`
        SELECT server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value
        FROM servers
        %s
        ORDER BY server_name, version
        LIMIT $%d
    `, whereClause, argIndex)
	args = append(args, limit)

	rows, err := db.getExecutor(tx).Query(ctx, query, args...)
	if err != nil {
		return nil, "", fmt.Errorf("failed to query servers: %w", err)
	}
	defer rows.Close()

	var results []*apiv0.ServerResponse
	for rows.Next() {
		var serverName, version, status string
		var statusChangedAt, publishedAt, updatedAt time.Time
		var statusMessage *string
		var isLatest bool
		var valueJSON []byte

		err := rows.Scan(&serverName, &version, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest, &valueJSON)
		if err != nil {
			return nil, "", fmt.Errorf("failed to scan server row: %w", err)
		}

		// Parse the ServerJSON from JSONB
		var serverJSON apiv0.ServerJSON
		if err := json.Unmarshal(valueJSON, &serverJSON); err != nil {
			return nil, "", fmt.Errorf("failed to unmarshal server JSON: %w", err)
		}

		// Build ServerResponse with separated metadata
		serverResponse := &apiv0.ServerResponse{
			Server: serverJSON,
			Meta: apiv0.ResponseMeta{
				Official: &apiv0.RegistryExtensions{
					Status:          model.Status(status),
					StatusChangedAt: statusChangedAt,
					StatusMessage:   statusMessage,
					PublishedAt:     publishedAt,
					UpdatedAt:       updatedAt,
					IsLatest:        isLatest,
				},
			},
		}

		results = append(results, serverResponse)
	}

	if err := rows.Err(); err != nil {
		return nil, "", fmt.Errorf("error iterating rows: %w", err)
	}

	// Determine next cursor using compound serverName:version format
	nextCursor := ""
	if len(results) > 0 && len(results) >= limit {
		lastResult := results[len(results)-1]
		nextCursor = lastResult.Server.Name + ":" + lastResult.Server.Version
	}

	return results, nextCursor, nil
}

// GetServerByName retrieves the latest version of a server by server name
func (db *PostgreSQL) GetServerByName(ctx context.Context, tx pgx.Tx, serverName string, includeDeleted bool) (*apiv0.ServerResponse, error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}

	// Build filter conditions
	isLatest := true
	filter := &ServerFilter{
		Name:           &serverName,
		IsLatest:       &isLatest,
		IncludeDeleted: &includeDeleted,
	}

	argIndex := 1
	whereConditions, args, _ := buildFilterConditions(filter, argIndex)

	whereClause := ""
	if len(whereConditions) > 0 {
		whereClause = "WHERE " + strings.Join(whereConditions, " AND ")
	}

	query := fmt.Sprintf(`
		SELECT server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value
		FROM servers
		%s
		ORDER BY published_at DESC
		LIMIT 1
	`, whereClause)

	var name, version, status string
	var statusChangedAt, publishedAt, updatedAt time.Time
	var statusMessage *string
	var valueJSON []byte

	err := db.getExecutor(tx).QueryRow(ctx, query, args...).Scan(&name, &version, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest, &valueJSON)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("failed to get server by name: %w", err)
	}

	// Parse the ServerJSON from JSONB
	var serverJSON apiv0.ServerJSON
	if err := json.Unmarshal(valueJSON, &serverJSON); err != nil {
		return nil, fmt.Errorf("failed to unmarshal server JSON: %w", err)
	}

	// Build ServerResponse with separated metadata
	serverResponse := &apiv0.ServerResponse{
		Server: serverJSON,
		Meta: apiv0.ResponseMeta{
			Official: &apiv0.RegistryExtensions{
				Status:          model.Status(status),
				StatusChangedAt: statusChangedAt,
				StatusMessage:   statusMessage,
				PublishedAt:     publishedAt,
				UpdatedAt:       updatedAt,
				IsLatest:        isLatest,
			},
		},
	}

	return serverResponse, nil
}

// GetServerByNameAndVersion retrieves a specific version of a server by server name and version
func (db *PostgreSQL) GetServerByNameAndVersion(ctx context.Context, tx pgx.Tx, serverName string, version string, includeDeleted bool) (*apiv0.ServerResponse, error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}

	// Build filter conditions
	filter := &ServerFilter{
		Name:           &serverName,
		Version:        &version,
		IncludeDeleted: &includeDeleted,
	}

	argIndex := 1
	whereConditions, args, _ := buildFilterConditions(filter, argIndex)

	whereClause := ""
	if len(whereConditions) > 0 {
		whereClause = "WHERE " + strings.Join(whereConditions, " AND ")
	}

	query := fmt.Sprintf(`
		SELECT server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value
		FROM servers
		%s
		LIMIT 1
	`, whereClause)

	var name, vers, status string
	var statusChangedAt, publishedAt, updatedAt time.Time
	var statusMessage *string
	var isLatest bool
	var valueJSON []byte

	err := db.getExecutor(tx).QueryRow(ctx, query, args...).Scan(&name, &vers, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest, &valueJSON)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("failed to get server by name and version: %w", err)
	}

	// Parse the ServerJSON from JSONB
	var serverJSON apiv0.ServerJSON
	if err := json.Unmarshal(valueJSON, &serverJSON); err != nil {
		return nil, fmt.Errorf("failed to unmarshal server JSON: %w", err)
	}

	// Build ServerResponse with separated metadata
	serverResponse := &apiv0.ServerResponse{
		Server: serverJSON,
		Meta: apiv0.ResponseMeta{
			Official: &apiv0.RegistryExtensions{
				Status:          model.Status(status),
				StatusChangedAt: statusChangedAt,
				StatusMessage:   statusMessage,
				PublishedAt:     publishedAt,
				UpdatedAt:       updatedAt,
				IsLatest:        isLatest,
			},
		},
	}

	return serverResponse, nil
}

// GetAllVersionsByServerName retrieves all versions of a server by server name
func (db *PostgreSQL) GetAllVersionsByServerName(ctx context.Context, tx pgx.Tx, serverName string, includeDeleted bool) ([]*apiv0.ServerResponse, error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}

	// Build filter conditions
	filter := &ServerFilter{
		Name:           &serverName,
		IncludeDeleted: &includeDeleted,
	}

	argIndex := 1
	whereConditions, args, _ := buildFilterConditions(filter, argIndex)

	whereClause := ""
	if len(whereConditions) > 0 {
		whereClause = "WHERE " + strings.Join(whereConditions, " AND ")
	}

	query := fmt.Sprintf(`
		SELECT server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value
		FROM servers
		%s
		ORDER BY published_at DESC
	`, whereClause)

	rows, err := db.getExecutor(tx).Query(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("failed to query server versions: %w", err)
	}
	defer rows.Close()

	var results []*apiv0.ServerResponse
	for rows.Next() {
		var name, version, status string
		var statusChangedAt, publishedAt, updatedAt time.Time
		var statusMessage *string
		var isLatest bool
		var valueJSON []byte

		err := rows.Scan(&name, &version, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest, &valueJSON)
		if err != nil {
			return nil, fmt.Errorf("failed to scan server row: %w", err)
		}

		// Parse the ServerJSON from JSONB
		var serverJSON apiv0.ServerJSON
		if err := json.Unmarshal(valueJSON, &serverJSON); err != nil {
			return nil, fmt.Errorf("failed to unmarshal server JSON: %w", err)
		}

		// Build ServerResponse with separated metadata
		serverResponse := &apiv0.ServerResponse{
			Server: serverJSON,
			Meta: apiv0.ResponseMeta{
				Official: &apiv0.RegistryExtensions{
					Status:          model.Status(status),
					StatusChangedAt: statusChangedAt,
					StatusMessage:   statusMessage,
					PublishedAt:     publishedAt,
					UpdatedAt:       updatedAt,
					IsLatest:        isLatest,
				},
			},
		}

		results = append(results, serverResponse)
	}

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("error iterating rows: %w", err)
	}

	if len(results) == 0 {
		return nil, ErrNotFound
	}

	return results, nil
}

// CreateServer inserts a new server version with official metadata
func (db *PostgreSQL) CreateServer(ctx context.Context, tx pgx.Tx, serverJSON *apiv0.ServerJSON, officialMeta *apiv0.RegistryExtensions) (*apiv0.ServerResponse, error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}

	// Validate inputs
	if serverJSON == nil || officialMeta == nil {
		return nil, fmt.Errorf("serverJSON and officialMeta are required")
	}

	if serverJSON.Name == "" || serverJSON.Version == "" {
		return nil, fmt.Errorf("server name and version are required")
	}

	// Marshal the ServerJSON to JSONB
	valueJSON, err := json.Marshal(serverJSON)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal server JSON: %w", err)
	}

	// Insert the new server version using composite primary key
	insertQuery := `
		INSERT INTO servers (server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
	`

	_, err = db.getExecutor(tx).Exec(ctx, insertQuery,
		serverJSON.Name,
		serverJSON.Version,
		string(officialMeta.Status),
		officialMeta.StatusChangedAt,
		officialMeta.StatusMessage,
		officialMeta.PublishedAt,
		officialMeta.UpdatedAt,
		officialMeta.IsLatest,
		valueJSON,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to insert server: %w", err)
	}

	// Return the complete ServerResponse
	serverResponse := &apiv0.ServerResponse{
		Server: *serverJSON,
		Meta: apiv0.ResponseMeta{
			Official: officialMeta,
		},
	}

	return serverResponse, nil
}

// UpdateServer updates an existing server record with new server details
func (db *PostgreSQL) UpdateServer(ctx context.Context, tx pgx.Tx, serverName, version string, serverJSON *apiv0.ServerJSON) (*apiv0.ServerResponse, error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}

	// Validate inputs
	if serverJSON == nil {
		return nil, fmt.Errorf("serverJSON is required")
	}

	// Ensure the serverJSON matches the provided serverName and version
	if serverJSON.Name != serverName || serverJSON.Version != version {
		return nil, fmt.Errorf("%w: server name and version in JSON must match parameters", ErrInvalidInput)
	}

	// Marshal updated ServerJSON
	valueJSON, err := json.Marshal(serverJSON)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal updated server: %w", err)
	}

	// Update only the JSON data (keep existing metadata columns)
	query := `
		UPDATE servers
		SET value = $1, updated_at = NOW()
		WHERE server_name = $2 AND version = $3
		RETURNING server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest
	`

	var name, vers, status string
	var statusChangedAt, publishedAt, updatedAt time.Time
	var statusMessage *string
	var isLatest bool

	err = db.getExecutor(tx).QueryRow(ctx, query, valueJSON, serverName, version).Scan(&name, &vers, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("failed to update server: %w", err)
	}

	// Return the updated ServerResponse
	serverResponse := &apiv0.ServerResponse{
		Server: *serverJSON,
		Meta: apiv0.ResponseMeta{
			Official: &apiv0.RegistryExtensions{
				Status:          model.Status(status),
				StatusChangedAt: statusChangedAt,
				StatusMessage:   statusMessage,
				PublishedAt:     publishedAt,
				UpdatedAt:       updatedAt,
				IsLatest:        isLatest,
			},
		},
	}

	return serverResponse, nil
}

// SetServerStatus updates the status of a specific server version
func (db *PostgreSQL) SetServerStatus(ctx context.Context, tx pgx.Tx, serverName, version string, status model.Status, statusMessage *string) (*apiv0.ServerResponse, error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}

	// Update the status and related fields
	// Only update status_changed_at when status actually changes
	query := `
		UPDATE servers
		SET
			status = $1,
			status_changed_at = CASE WHEN status != $1::varchar THEN NOW() ELSE status_changed_at END,
			updated_at = NOW(),
			status_message = $4
		WHERE server_name = $2 AND version = $3
		RETURNING server_name, version, status, value, published_at, updated_at, is_latest, status_changed_at, status_message
	`

	var name, vers, currentStatus string
	var publishedAt, updatedAt, statusChangedAt time.Time
	var isLatest bool
	var valueJSON []byte
	var resultStatusMessage *string

	err := db.getExecutor(tx).QueryRow(ctx, query, string(status), serverName, version, statusMessage).Scan(&name, &vers, &currentStatus, &valueJSON, &publishedAt, &updatedAt, &isLatest, &statusChangedAt, &resultStatusMessage)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("failed to update server status: %w", err)
	}

	// Unmarshal the JSON data
	var serverJSON apiv0.ServerJSON
	if err := json.Unmarshal(valueJSON, &serverJSON); err != nil {
		return nil, fmt.Errorf("failed to unmarshal server JSON: %w", err)
	}

	// Return the updated ServerResponse
	serverResponse := &apiv0.ServerResponse{
		Server: serverJSON,
		Meta: apiv0.ResponseMeta{
			Official: &apiv0.RegistryExtensions{
				Status:          model.Status(currentStatus),
				StatusChangedAt: statusChangedAt,
				StatusMessage:   resultStatusMessage,
				PublishedAt:     publishedAt,
				UpdatedAt:       updatedAt,
				IsLatest:        isLatest,
			},
		},
	}

	return serverResponse, nil
}

// SetAllVersionsStatus updates the status of all versions of a server in a single query
func (db *PostgreSQL) SetAllVersionsStatus(ctx context.Context, tx pgx.Tx, serverName string, status model.Status, statusMessage *string) ([]*apiv0.ServerResponse, error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}

	// Update the status and related fields for all versions
	// Only update rows where status or status_message actually changes
	// Only update status_changed_at when status actually changes
	query := `
		UPDATE servers
		SET
			status = $1,
			status_changed_at = CASE WHEN status != $1::varchar THEN NOW() ELSE status_changed_at END,
			updated_at = NOW(),
			status_message = $2
		WHERE server_name = $3
			AND (status != $1::varchar OR status_message IS DISTINCT FROM $2)
		RETURNING server_name, version, status, value, published_at, updated_at, is_latest, status_changed_at, status_message
	`

	rows, err := db.getExecutor(tx).Query(ctx, query, string(status), statusMessage, serverName)
	if err != nil {
		return nil, fmt.Errorf("failed to update all server versions status: %w", err)
	}
	defer rows.Close()

	var results []*apiv0.ServerResponse
	for rows.Next() {
		var name, vers, currentStatus string
		var publishedAt, updatedAt, statusChangedAt time.Time
		var isLatest bool
		var valueJSON []byte
		var resultStatusMessage *string

		if err := rows.Scan(&name, &vers, &currentStatus, &valueJSON, &publishedAt, &updatedAt, &isLatest, &statusChangedAt, &resultStatusMessage); err != nil {
			return nil, fmt.Errorf("failed to scan server row: %w", err)
		}

		// Unmarshal the JSON data
		var serverJSON apiv0.ServerJSON
		if err := json.Unmarshal(valueJSON, &serverJSON); err != nil {
			return nil, fmt.Errorf("failed to unmarshal server JSON: %w", err)
		}

		serverResponse := &apiv0.ServerResponse{
			Server: serverJSON,
			Meta: apiv0.ResponseMeta{
				Official: &apiv0.RegistryExtensions{
					Status:          model.Status(currentStatus),
					StatusChangedAt: statusChangedAt,
					StatusMessage:   resultStatusMessage,
					PublishedAt:     publishedAt,
					UpdatedAt:       updatedAt,
					IsLatest:        isLatest,
				},
			},
		}
		results = append(results, serverResponse)
	}

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("error iterating server rows: %w", err)
	}

	if len(results) == 0 {
		return nil, ErrNotFound
	}

	return results, nil
}

// InTransaction executes a function within a database transaction
func (db *PostgreSQL) InTransaction(ctx context.Context, fn func(ctx context.Context, tx pgx.Tx) error) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	tx, err := db.pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	//nolint:contextcheck // Intentionally using separate context for rollback to ensure cleanup even if request is cancelled
	defer func() {
		rollbackCtx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
		defer cancel()
		if rbErr := tx.Rollback(rollbackCtx); rbErr != nil && !errors.Is(rbErr, pgx.ErrTxClosed) {
			log.Printf("failed to rollback transaction: %v", rbErr)
		}
	}()

	if err := fn(ctx, tx); err != nil {
		return err
	}

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

	return nil
}

// AcquirePublishLock acquires an exclusive advisory lock for publishing a server
// This prevents race conditions when multiple versions are published concurrently
// Using pg_advisory_xact_lock which auto-releases on transaction end
func (db *PostgreSQL) AcquirePublishLock(ctx context.Context, tx pgx.Tx, serverName string) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	lockID := hashServerName(serverName)

	if _, err := db.getExecutor(tx).Exec(ctx, "SELECT pg_advisory_xact_lock($1)", lockID); err != nil {
		return fmt.Errorf("failed to acquire publish lock: %w", err)
	}

	return nil
}

// hashServerName creates a consistent hash of the server name for advisory locking
// We use FNV-1a hash and mask to 63 bits to fit in PostgreSQL's bigint range
func hashServerName(name string) int64 {
	const (
		offset64 = 14695981039346656037
		prime64  = 1099511628211
	)
	hash := uint64(offset64)
	for i := 0; i < len(name); i++ {
		hash ^= uint64(name[i])
		hash *= prime64
	}
	return int64(hash & 0x7FFFFFFFFFFFFFFF)
}

// GetCurrentLatestVersion retrieves the current latest version of a server by server name
func (db *PostgreSQL) GetCurrentLatestVersion(ctx context.Context, tx pgx.Tx, serverName string) (*apiv0.ServerResponse, error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}

	executor := db.getExecutor(tx)

	query := `
		SELECT server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value
		FROM servers
		WHERE server_name = $1 AND is_latest = true
	`

	row := executor.QueryRow(ctx, query, serverName)

	var name, version, status string
	var statusChangedAt, publishedAt, updatedAt time.Time
	var statusMessage *string
	var isLatest bool
	var jsonValue []byte

	err := row.Scan(&name, &version, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest, &jsonValue)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("failed to scan server row: %w", err)
	}

	// Parse the JSON value to get the server details
	var serverJSON apiv0.ServerJSON
	if err := json.Unmarshal(jsonValue, &serverJSON); err != nil {
		return nil, fmt.Errorf("failed to unmarshal server JSON: %w", err)
	}

	// Build ServerResponse with separated metadata
	serverResponse := &apiv0.ServerResponse{
		Server: serverJSON,
		Meta: apiv0.ResponseMeta{
			Official: &apiv0.RegistryExtensions{
				Status:          model.Status(status),
				StatusChangedAt: statusChangedAt,
				StatusMessage:   statusMessage,
				PublishedAt:     publishedAt,
				UpdatedAt:       updatedAt,
				IsLatest:        isLatest,
			},
		},
	}

	return serverResponse, nil
}

// CountServerVersions counts the number of versions for a server
func (db *PostgreSQL) CountServerVersions(ctx context.Context, tx pgx.Tx, serverName string) (int, error) {
	if ctx.Err() != nil {
		return 0, ctx.Err()
	}

	executor := db.getExecutor(tx)

	query := `SELECT COUNT(*) FROM servers WHERE server_name = $1`

	var count int
	err := executor.QueryRow(ctx, query, serverName).Scan(&count)
	if err != nil {
		return 0, fmt.Errorf("failed to count server versions: %w", err)
	}

	return count, nil
}

// CheckVersionExists checks if a specific version exists for a server
func (db *PostgreSQL) CheckVersionExists(ctx context.Context, tx pgx.Tx, serverName, version string) (bool, error) {
	if ctx.Err() != nil {
		return false, ctx.Err()
	}

	executor := db.getExecutor(tx)

	query := `SELECT EXISTS(SELECT 1 FROM servers WHERE server_name = $1 AND version = $2)`

	var exists bool
	err := executor.QueryRow(ctx, query, serverName, version).Scan(&exists)
	if err != nil {
		return false, fmt.Errorf("failed to check version existence: %w", err)
	}

	return exists, nil
}

// UnmarkAsLatest marks the current latest version of a server as no longer latest
func (db *PostgreSQL) UnmarkAsLatest(ctx context.Context, tx pgx.Tx, serverName string) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	executor := db.getExecutor(tx)

	query := `UPDATE servers SET is_latest = false WHERE server_name = $1 AND is_latest = true`

	_, err := executor.Exec(ctx, query, serverName)
	if err != nil {
		return fmt.Errorf("failed to unmark latest version: %w", err)
	}

	return nil
}

// SetLatestVersion sets is_latest=true on the given version and false on all other versions
// of the same server. Passing an empty version clears is_latest for all rows.
//
// The clear and set are issued as separate statements because the unique partial index
// idx_unique_latest_per_server is non-deferrable and Postgres checks it row-by-row within
// a single UPDATE, which would trip when flipping one row's flag off and another's on.
func (db *PostgreSQL) SetLatestVersion(ctx context.Context, tx pgx.Tx, serverName, version string) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	executor := db.getExecutor(tx)

	if _, err := executor.Exec(ctx,
		`UPDATE servers SET is_latest = false WHERE server_name = $1 AND is_latest = true AND version <> $2`,
		serverName, version,
	); err != nil {
		return fmt.Errorf("failed to clear previous latest version: %w", err)
	}

	if version == "" {
		return nil
	}

	if _, err := executor.Exec(ctx,
		`UPDATE servers SET is_latest = true WHERE server_name = $1 AND version = $2 AND is_latest = false`,
		serverName, version,
	); err != nil {
		return fmt.Errorf("failed to set latest version: %w", err)
	}

	return nil
}

// Close closes the database connection
func (db *PostgreSQL) Close() error {
	db.pool.Close()
	return nil
}