registry/internal/database/postgres.go

package database

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

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/modelcontextprotocol/registry/internal/model"
)

// PostgreSQL is an implementation of the Database interface using PostgreSQL
type PostgreSQL struct {
	conn *pgx.Conn
}

// NewPostgreSQL creates a new instance of the PostgreSQL database
func NewPostgreSQL(ctx context.Context, connectionURI string) (*PostgreSQL, error) {
	conn, err := pgx.Connect(ctx, connectionURI)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to PostgreSQL: %w", err)
	}

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

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

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

// List retrieves MCPRegistry entries with optional filtering and pagination
func (db *PostgreSQL) List(
	ctx context.Context,
	filter map[string]any,
	cursor string,
	limit int,
) ([]*model.Server, string, error) {
	if limit <= 0 {
		limit = 10
	}

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

	// Build WHERE clause
	whereClause := "WHERE is_latest = true"
	args := []any{}
	argIndex := 1

	// Add filters
	for k, v := range filter {
		switch k {
		case "name":
			whereClause += fmt.Sprintf(" AND name = $%d", argIndex)
			args = append(args, v)
			argIndex++
		case "version":
			whereClause += fmt.Sprintf(" AND version = $%d", argIndex)
			args = append(args, v)
			argIndex++
		case "status":
			whereClause += fmt.Sprintf(" AND status = $%d", argIndex)
			args = append(args, v)
			argIndex++
		}
	}

	// Add cursor pagination
	if cursor != "" {
		if _, err := uuid.Parse(cursor); err != nil {
			return nil, "", fmt.Errorf("invalid cursor format: %w", err)
		}
		whereClause += fmt.Sprintf(" AND id > $%d", argIndex)
		args = append(args, cursor)
		argIndex++
	}

	// Build query
	query := fmt.Sprintf(`
		SELECT id, name, description, status, repository, version, release_date, is_latest
		FROM servers
		%s
		ORDER BY id
		LIMIT $%d
	`, whereClause, argIndex)
	args = append(args, limit)

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

	var results []*model.Server
	for rows.Next() {
		var server model.Server
		var repositoryJSON []byte
		var releaseDate time.Time

		err := rows.Scan(
			&server.ID,
			&server.Name,
			&server.Description,
			&server.Status,
			&repositoryJSON,
			&server.VersionDetail.Version,
			&releaseDate,
			&server.VersionDetail.IsLatest,
		)
		if err != nil {
			return nil, "", fmt.Errorf("failed to scan server row: %w", err)
		}

		// Parse repository JSON
		if len(repositoryJSON) > 0 {
			if err := json.Unmarshal(repositoryJSON, &server.Repository); err != nil {
				return nil, "", fmt.Errorf("failed to unmarshal repository: %w", err)
			}
		}

		server.VersionDetail.ReleaseDate = releaseDate.Format(time.RFC3339)
		results = append(results, &server)
	}

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

	// Determine next cursor
	nextCursor := ""
	if len(results) > 0 && len(results) >= limit {
		nextCursor = results[len(results)-1].ID
	}

	return results, nextCursor, nil
}

// GetByID retrieves a single ServerDetail by its ID
func (db *PostgreSQL) GetByID(ctx context.Context, id string) (*model.ServerDetail, error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}

	query := `
		SELECT id, name, description, status, repository, version, release_date, is_latest, packages, remotes
		FROM servers
		WHERE id = $1
	`

	var serverDetail model.ServerDetail
	var repositoryJSON, packagesJSON, remotesJSON []byte
	var releaseDate time.Time

	err := db.conn.QueryRow(ctx, query, id).Scan(
		&serverDetail.ID,
		&serverDetail.Name,
		&serverDetail.Description,
		&serverDetail.Status,
		&repositoryJSON,
		&serverDetail.VersionDetail.Version,
		&releaseDate,
		&serverDetail.VersionDetail.IsLatest,
		&packagesJSON,
		&remotesJSON,
	)

	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("failed to get server by ID: %w", err)
	}

	// Parse JSON fields
	if len(repositoryJSON) > 0 {
		if err := json.Unmarshal(repositoryJSON, &serverDetail.Repository); err != nil {
			return nil, fmt.Errorf("failed to unmarshal repository: %w", err)
		}
	}

	if len(packagesJSON) > 0 {
		if err := json.Unmarshal(packagesJSON, &serverDetail.Packages); err != nil {
			return nil, fmt.Errorf("failed to unmarshal packages: %w", err)
		}
	}

	if len(remotesJSON) > 0 {
		if err := json.Unmarshal(remotesJSON, &serverDetail.Remotes); err != nil {
			return nil, fmt.Errorf("failed to unmarshal remotes: %w", err)
		}
	}

	serverDetail.VersionDetail.ReleaseDate = releaseDate.Format(time.RFC3339)

	return &serverDetail, nil
}

// Publish adds a new ServerDetail to the database
func (db *PostgreSQL) Publish(ctx context.Context, serverDetail *model.ServerDetail) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	tx, err := db.conn.Begin(ctx)
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer func() {
		if err := tx.Rollback(ctx); err != nil {
			log.Printf("Failed to rollback transaction: %v", err)
		}
	}()

	// Check if there's an existing latest version for this server
	var existingVersion string
	checkQuery := `
		SELECT version 
		FROM servers 
		WHERE name = $1 AND is_latest = true
	`
	err = tx.QueryRow(ctx, checkQuery, serverDetail.Name).Scan(&existingVersion)
	if err != nil && !errors.Is(err, pgx.ErrNoRows) {
		return fmt.Errorf("failed to check existing version: %w", err)
	}

	// Validate version ordering
	if existingVersion != "" && serverDetail.VersionDetail.Version <= existingVersion {
		return ErrInvalidVersion
	}

	// Prepare JSON data
	repositoryJSON, err := json.Marshal(serverDetail.Repository)
	if err != nil {
		return fmt.Errorf("failed to marshal repository: %w", err)
	}

	packagesJSON, err := json.Marshal(serverDetail.Packages)
	if err != nil {
		return fmt.Errorf("failed to marshal packages: %w", err)
	}

	remotesJSON, err := json.Marshal(serverDetail.Remotes)
	if err != nil {
		return fmt.Errorf("failed to marshal remotes: %w", err)
	}

	// Generate ID and set metadata
	if serverDetail.ID == "" {
		serverDetail.ID = uuid.New().String()
	}
	serverDetail.VersionDetail.IsLatest = true
	if serverDetail.VersionDetail.ReleaseDate == "" {
		serverDetail.VersionDetail.ReleaseDate = time.Now().Format(time.RFC3339)
	}

	releaseDate, err := time.Parse(time.RFC3339, serverDetail.VersionDetail.ReleaseDate)
	if err != nil {
		return fmt.Errorf("failed to parse release date: %w", err)
	}

	// Update existing latest version to not be latest
	if existingVersion != "" {
		updateQuery := `
			UPDATE servers 
			SET is_latest = false 
			WHERE name = $1 AND is_latest = true
		`
		_, err = tx.Exec(ctx, updateQuery, serverDetail.Name)
		if err != nil {
			return fmt.Errorf("failed to update existing latest version: %w", err)
		}
	}

	// Insert new server version
	insertQuery := `
		INSERT INTO servers (id, name, description, status, repository, version, release_date, is_latest, packages, remotes)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
	`
	_, err = tx.Exec(ctx, insertQuery,
		serverDetail.ID,
		serverDetail.Name,
		serverDetail.Description,
		serverDetail.Status,
		repositoryJSON,
		serverDetail.VersionDetail.Version,
		releaseDate,
		serverDetail.VersionDetail.IsLatest,
		packagesJSON,
		remotesJSON,
	)

	if err != nil {
		return fmt.Errorf("failed to insert server: %w", err)
	}

	return tx.Commit(ctx)
}

// ImportSeed imports initial data from a seed file into PostgreSQL
func (db *PostgreSQL) ImportSeed(ctx context.Context, seedFilePath string) error {
	servers, err := ReadSeedFile(ctx, seedFilePath)
	if err != nil {
		return fmt.Errorf("failed to read seed file: %w", err)
	}

	log.Printf("Importing %d servers into PostgreSQL", len(servers))

	for i, server := range servers {
		if server.ID == "" || server.Name == "" {
			log.Printf("Skipping server %d: ID or Name is empty", i+1)
			continue
		}

		if server.VersionDetail.Version == "" {
			server.VersionDetail.Version = "0.0.1-seed"
			server.VersionDetail.ReleaseDate = time.Now().Format(time.RFC3339)
			server.VersionDetail.IsLatest = true
		}

		// Prepare JSON data
		repositoryJSON, err := json.Marshal(server.Repository)
		if err != nil {
			log.Printf("Error marshaling repository for server %s: %v", server.ID, err)
			continue
		}

		packagesJSON, err := json.Marshal(server.Packages)
		if err != nil {
			log.Printf("Error marshaling packages for server %s: %v", server.ID, err)
			continue
		}

		remotesJSON, err := json.Marshal(server.Remotes)
		if err != nil {
			log.Printf("Error marshaling remotes for server %s: %v", server.ID, err)
			continue
		}

		releaseDate, err := time.Parse(time.RFC3339, server.VersionDetail.ReleaseDate)
		if err != nil {
			log.Printf("Error parsing release date for server %s: %v", server.ID, err)
			releaseDate = time.Now()
		}

		// Use UPSERT (INSERT ... ON CONFLICT DO UPDATE)
		upsertQuery := `
			INSERT INTO servers (id, name, description, status, repository, version, release_date, is_latest, packages, remotes)
			VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
			ON CONFLICT (id) DO UPDATE SET
				name = EXCLUDED.name,
				description = EXCLUDED.description,
				status = EXCLUDED.status,
				repository = EXCLUDED.repository,
				version = EXCLUDED.version,
				release_date = EXCLUDED.release_date,
				is_latest = EXCLUDED.is_latest,
				packages = EXCLUDED.packages,
				remotes = EXCLUDED.remotes,
				updated_at = NOW()
		`

		_, err = db.conn.Exec(ctx, upsertQuery,
			server.ID,
			server.Name,
			server.Description,
			server.Status,
			repositoryJSON,
			server.VersionDetail.Version,
			releaseDate,
			server.VersionDetail.IsLatest,
			packagesJSON,
			remotesJSON,
		)

		if err != nil {
			log.Printf("Error importing server %s: %v", server.ID, err)
			continue
		}

		log.Printf("[%d/%d] Imported server: %s", i+1, len(servers), server.Name)
	}

	log.Println("PostgreSQL database import completed successfully")
	return nil
}

// Close closes the database connection
func (db *PostgreSQL) Close() error {
	return db.conn.Close(context.Background())
}

// Connection returns information about the database connection
func (db *PostgreSQL) Connection() *ConnectionInfo {
	isConnected := false
	if db.conn != nil {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		err := db.conn.Ping(ctx)
		isConnected = (err == nil)
	}

	return &ConnectionInfo{
		Type:        ConnectionTypePostgreSQL,
		IsConnected: isConnected,
		Raw:         db.conn,
	}
}