proxy.go

package pgpool

import (
	"context"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"go.uber.org/zap"
)

// ProxyServer implements PostgreSQL wire protocol server
// This allows PHP to connect using standard PostgreSQL drivers
type ProxyServer struct {
	// Server address (e.g., "tcp://127.0.0.1:5433")
	address string

	// Reference to main plugin for pool access
	plugin *Plugin

	// Logger
	log *zap.Logger

	// Configuration
	config *ProxyConfig

	// TCP listener
	listener net.Listener

	// Active client connections
	clients sync.Map // map[string]*ClientConnection

	// Connection counter for IDs
	connCounter atomic.Int64

	// Lifecycle
	mu       sync.RWMutex
	ctx      context.Context
	cancel   context.CancelFunc
	stopCh   chan struct{}
	stopped  atomic.Bool
}

// ClientConnection represents a single PHP client connection
type ClientConnection struct {
	// Unique connection ID
	id string

	// Network connection
	conn net.Conn

	// Associated database pool name
	database string

	// Current transaction state
	inTransaction bool
	txConn        interface{} // Holds pgxpool.Conn during transaction

	// Connection metadata
	remoteAddr string
	startTime  time.Time

	// Lifecycle
	ctx    context.Context
	cancel context.CancelFunc
}

// NewProxyServer creates a new PostgreSQL protocol proxy server
func NewProxyServer(
	address string,
	plugin *Plugin,
	log *zap.Logger,
	config *ProxyConfig,
) (*ProxyServer, error) {
	ctx, cancel := context.WithCancel(context.Background())

	return &ProxyServer{
		address: address,
		plugin:  plugin,
		log:     log,
		config:  config,
		ctx:     ctx,
		cancel:  cancel,
		stopCh:  make(chan struct{}),
	}, nil
}

// Start begins accepting client connections
func (p *ProxyServer) Start(ctx context.Context) error {
	// Parse address (e.g., "tcp://127.0.0.1:5433")
	network, addr, err := parseAddress(p.address)
	if err != nil {
		return fmt.Errorf("invalid proxy address: %w", err)
	}

	// Create listener
	listener, err := net.Listen(network, addr)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %w", p.address, err)
	}
	p.listener = listener

	p.log.Info("proxy server listening",
		zap.String("address", p.address),
		zap.String("network", network))

	// Accept connections in background
	go p.acceptLoop()

	return nil
}

// Stop gracefully shuts down the proxy server
func (p *ProxyServer) Stop(ctx context.Context) error {
	if !p.stopped.CompareAndSwap(false, true) {
		return fmt.Errorf("proxy already stopped")
	}

	p.log.Info("stopping proxy server")
	close(p.stopCh)
	p.cancel()

	// Close listener to stop accepting new connections
	if p.listener != nil {
		if err := p.listener.Close(); err != nil {
			p.log.Warn("error closing listener", zap.Error(err))
		}
	}

	// Close all active client connections
	var wg sync.WaitGroup
	p.clients.Range(func(key, value interface{}) bool {
		wg.Add(1)
		go func(client *ClientConnection) {
			defer wg.Done()
			client.Close()
		}(value.(*ClientConnection))
		return true
	})

	// Wait for all clients to close or timeout
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		p.log.Info("all client connections closed")
	case <-ctx.Done():
		p.log.Warn("shutdown timeout exceeded")
	}

	return nil
}

// acceptLoop continuously accepts new client connections
func (p *ProxyServer) acceptLoop() {
	for {
		select {
		case <-p.stopCh:
			return
		default:
		}

		conn, err := p.listener.Accept()
		if err != nil {
			// Check if this is due to listener being closed
			if p.stopped.Load() {
				return
			}
			p.log.Error("failed to accept connection", zap.Error(err))
			continue
		}

		// Handle connection in separate goroutine
		go p.handleConnection(conn)
	}
}

// handleConnection processes a single client connection
func (p *ProxyServer) handleConnection(conn net.Conn) {
	// Generate unique connection ID
	connID := fmt.Sprintf("conn_%d", p.connCounter.Add(1))

	// Create client connection context
	ctx, cancel := context.WithCancel(p.ctx)
	client := &ClientConnection{
		id:         connID,
		conn:       conn,
		database:   "primary", // Default database
		remoteAddr: conn.RemoteAddr().String(),
		startTime:  time.Now(),
		ctx:        ctx,
		cancel:     cancel,
	}

	// Register client
	p.clients.Store(connID, client)
	defer func() {
		p.clients.Delete(connID)
		client.Close()
	}()

	p.log.Info("client connected",
		zap.String("conn_id", connID),
		zap.String("remote_addr", client.remoteAddr))

	// Handle PostgreSQL protocol
	if err := p.handlePostgreSQLProtocol(client); err != nil {
		p.log.Error("protocol error",
			zap.String("conn_id", connID),
			zap.Error(err))
	}

	p.log.Info("client disconnected",
		zap.String("conn_id", connID),
		zap.Duration("session_duration", time.Since(client.startTime)))
}

// handlePostgreSQLProtocol implements the PostgreSQL wire protocol
// This is a simplified implementation - Phase 1 MVP
func (p *ProxyServer) handlePostgreSQLProtocol(client *ClientConnection) error {
	// NOTE: This is a placeholder for Phase 2 implementation
	// Full PostgreSQL wire protocol implementation requires:
	// 1. Startup message handling
	// 2. Authentication (SCRAM-SHA-256, MD5, etc.)
	// 3. Simple query protocol
	// 4. Extended query protocol (Parse, Bind, Execute, Describe)
	// 5. Transaction commands (BEGIN, COMMIT, ROLLBACK)
	// 6. COPY protocol
	// 7. Notification/LISTEN support
	//
	// For Phase 1, we recommend using the RPC interface
	// Phase 2 will implement full protocol using github.com/jeroenrinzema/psql-wire

	p.log.Warn("PostgreSQL wire protocol not yet implemented",
		zap.String("conn_id", client.id),
		zap.String("phase", "Phase 2 - use RPC interface for now"))

	// Send error message to client
	errMsg := "PostgreSQL wire protocol support coming in Phase 2. Please use RPC interface.\n"
	_, _ = client.conn.Write([]byte(errMsg))

	return fmt.Errorf("protocol not implemented")
}

// Close closes a client connection
func (c *ClientConnection) Close() error {
	c.cancel()

	// Rollback any active transaction
	if c.inTransaction && c.txConn != nil {
		// Transaction cleanup would go here in Phase 2
		c.inTransaction = false
		c.txConn = nil
	}

	return c.conn.Close()
}

// parseAddress parses address string into network and address components
// Supports formats: "tcp://host:port", "unix:///path/to/socket"
func parseAddress(addr string) (network, address string, err error) {
	// Simple parser for common formats
	if len(addr) < 6 {
		return "", "", fmt.Errorf("invalid address format")
	}

	if addr[:6] == "tcp://" {
		return "tcp", addr[6:], nil
	}

	if addr[:7] == "unix://" {
		return "unix", addr[7:], nil
	}

	// Default to TCP if no protocol specified
	return "tcp", addr, nil
}