rabbitmq.go

package goqueue

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"os"
	"sync"
	"time"

	amqp "github.com/rabbitmq/amqp091-go"
)

// RabbitMQBackend implements the Backend interface using RabbitMQ
type RabbitMQBackend struct {
	// Connection management
	conn        *amqp.Connection
	channels    chan *amqp.Channel // Channel pool
	channelsMu  sync.Mutex
	maxChannels int

	// Subscriber management (pattern from Redis backend)
	subscribers  map[string]chan *Envelope
	consumerTags map[string]string
	mu           sync.RWMutex
	closed       bool
	stopChans    map[string]chan struct{}

	// Message tracking for Ack/Nack
	messageTracking map[string]*messageTracking
	trackingMu      sync.RWMutex

	// Dedicated channels per queue
	queueChannels map[string]*amqp.Channel
	queueChanMu   sync.RWMutex

	config *rabbitmqConfig
}

// messageTracking stores delivery tag and queue for Ack/Nack operations
type messageTracking struct {
	deliveryTag uint64
	queue       string
}

// rabbitmqConfig holds all configuration for RabbitMQ backend
type rabbitmqConfig struct {
	// Connection
	url      string
	username string
	password string
	host     string
	port     int
	vhost    string

	// TLS
	tlsConfig *tls.Config

	// Connection pooling
	maxChannels int

	// Queue configuration
	durable    bool
	autoDelete bool
	exclusive  bool
	noWait     bool

	// Message configuration
	persistent    bool // Delivery mode 2
	prefetchCount int  // QoS
	prefetchSize  int

	// Priority queue
	enablePriority bool
	maxPriority    uint8

	// Dead Letter Exchange
	enableDLX     bool
	dlxName       string
	dlqName       string
	dlxRoutingKey string
}

// RabbitMQOption is a function that configures RabbitMQBackend
type RabbitMQOption func(*rabbitmqConfig)

// defaultRabbitMQConfig returns the default configuration
func defaultRabbitMQConfig() *rabbitmqConfig {
	return &rabbitmqConfig{
		host:           "localhost",
		port:           5672,
		vhost:          "/",
		maxChannels:    100,
		durable:        true,
		autoDelete:     false,
		exclusive:      false,
		noWait:         false,
		persistent:     true,
		prefetchCount:  1,
		prefetchSize:   0,
		enablePriority: false,
		maxPriority:    10,
		enableDLX:      true,
		dlxName:        "goqueue.dlx",
		dlqName:        "dead-letter-queue",
		dlxRoutingKey:  "",
	}
}

// WithRabbitMQURL sets the full AMQP URL
func WithRabbitMQURL(url string) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.url = url
	}
}

// WithRabbitMQAuth sets the username and password
func WithRabbitMQAuth(username, password string) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.username = username
		c.password = password
	}
}

// WithRabbitMQHost sets the host and port
func WithRabbitMQHost(host string, port int) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.host = host
		c.port = port
	}
}

// WithRabbitMQVHost sets the virtual host
func WithRabbitMQVHost(vhost string) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.vhost = vhost
	}
}

// WithRabbitMQTLS sets TLS configuration from certificate files
func WithRabbitMQTLS(certFile, keyFile, caFile string) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		cfg := &tls.Config{}

		// Load client certificate if provided
		if certFile != "" && keyFile != "" {
			cert, err := tls.LoadX509KeyPair(certFile, keyFile)
			if err == nil {
				cfg.Certificates = []tls.Certificate{cert}
			}
		}

		// Load CA certificate if provided
		if caFile != "" {
			caCert, err := os.ReadFile(caFile)
			if err == nil {
				caCertPool := x509.NewCertPool()
				caCertPool.AppendCertsFromPEM(caCert)
				cfg.RootCAs = caCertPool
			}
		}

		c.tlsConfig = cfg
	}
}

// WithRabbitMQTLSConfig sets a custom TLS configuration
func WithRabbitMQTLSConfig(tlsConfig *tls.Config) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.tlsConfig = tlsConfig
	}
}

// WithRabbitMQInsecureSkipVerify skips TLS certificate verification (use only for development)
func WithRabbitMQInsecureSkipVerify() RabbitMQOption {
	return func(c *rabbitmqConfig) {
		if c.tlsConfig == nil {
			c.tlsConfig = &tls.Config{}
		}
		c.tlsConfig.InsecureSkipVerify = true
	}
}

// WithRabbitMQMaxChannels sets the maximum number of channels in the pool
func WithRabbitMQMaxChannels(max int) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.maxChannels = max
	}
}

// WithRabbitMQDurableQueues sets whether queues should be durable
func WithRabbitMQDurableQueues(durable bool) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.durable = durable
	}
}

// WithRabbitMQAutoDelete sets whether queues should auto-delete
func WithRabbitMQAutoDelete(autoDelete bool) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.autoDelete = autoDelete
	}
}

// WithRabbitMQExclusive sets whether queues should be exclusive
func WithRabbitMQExclusive(exclusive bool) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.exclusive = exclusive
	}
}

// WithRabbitMQPersistence sets whether messages should be persistent
func WithRabbitMQPersistence(persistent bool) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.persistent = persistent
	}
}

// WithRabbitMQPrefetch sets QoS prefetch count and size
func WithRabbitMQPrefetch(count, size int) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.prefetchCount = count
		c.prefetchSize = size
	}
}

// WithRabbitMQPriority enables priority queues with the given max priority
func WithRabbitMQPriority(maxPriority uint8) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.enablePriority = true
		c.maxPriority = maxPriority
	}
}

// WithRabbitMQDLX configures Dead Letter Exchange
func WithRabbitMQDLX(dlxName, dlqName string) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.enableDLX = true
		c.dlxName = dlxName
		c.dlqName = dlqName
	}
}

// WithRabbitMQDLXRoutingKey sets the routing key for DLX
func WithRabbitMQDLXRoutingKey(key string) RabbitMQOption {
	return func(c *rabbitmqConfig) {
		c.dlxRoutingKey = key
	}
}

// buildAMQPURL constructs an AMQP URL from components
func buildAMQPURL(c *rabbitmqConfig) string {
	scheme := "amqp"
	if c.tlsConfig != nil {
		scheme = "amqps"
	}

	auth := ""
	if c.username != "" {
		auth = c.username
		if c.password != "" {
			auth += ":" + c.password
		}
		auth += "@"
	}

	return fmt.Sprintf("%s://%s%s:%d%s", scheme, auth, c.host, c.port, c.vhost)
}

// convertMetadata converts Go map to AMQP Table
func convertMetadata(metadata map[string]string) amqp.Table {
	if metadata == nil {
		return nil
	}

	table := make(amqp.Table, len(metadata))
	for k, v := range metadata {
		table[k] = v
	}
	return table
}

// NewRabbitMQBackend creates a new RabbitMQ backend
func NewRabbitMQBackend(url string, opts ...RabbitMQOption) (*RabbitMQBackend, error) {
	// Apply default config
	config := defaultRabbitMQConfig()
	config.url = url

	// Apply options
	for _, opt := range opts {
		opt(config)
	}

	// Build connection URL if components provided
	if config.host != "" && config.url == url {
		config.url = buildAMQPURL(config)
	}

	// Establish connection
	var conn *amqp.Connection
	var err error

	if config.tlsConfig != nil {
		conn, err = amqp.DialTLS(config.url, config.tlsConfig)
	} else {
		conn, err = amqp.Dial(config.url)
	}

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

	// Initialize backend
	b := &RabbitMQBackend{
		conn:            conn,
		channels:        make(chan *amqp.Channel, config.maxChannels),
		maxChannels:     config.maxChannels,
		subscribers:     make(map[string]chan *Envelope),
		consumerTags:    make(map[string]string),
		stopChans:       make(map[string]chan struct{}),
		messageTracking: make(map[string]*messageTracking),
		queueChannels:   make(map[string]*amqp.Channel),
		config:          config,
	}

	// Pre-create channel pool
	if err := b.initChannelPool(); err != nil {
		conn.Close()
		return nil, err
	}

	return b, nil
}

// NewRabbitMQBackendWithConnection creates a new RabbitMQ backend with an existing connection
func NewRabbitMQBackendWithConnection(conn *amqp.Connection, opts ...RabbitMQOption) (*RabbitMQBackend, error) {
	// Apply default config
	config := defaultRabbitMQConfig()

	// Apply options
	for _, opt := range opts {
		opt(config)
	}

	// Initialize backend
	b := &RabbitMQBackend{
		conn:            conn,
		channels:        make(chan *amqp.Channel, config.maxChannels),
		maxChannels:     config.maxChannels,
		subscribers:     make(map[string]chan *Envelope),
		consumerTags:    make(map[string]string),
		stopChans:       make(map[string]chan struct{}),
		messageTracking: make(map[string]*messageTracking),
		queueChannels:   make(map[string]*amqp.Channel),
		config:          config,
	}

	// Pre-create channel pool
	if err := b.initChannelPool(); err != nil {
		return nil, err
	}

	return b, nil
}

// initChannelPool pre-creates channels for the pool
func (b *RabbitMQBackend) initChannelPool() error {
	for i := 0; i < b.maxChannels; i++ {
		ch, err := b.conn.Channel()
		if err != nil {
			return fmt.Errorf("failed to create channel: %w", err)
		}
		b.channels <- ch
	}
	return nil
}

// getChannel retrieves a channel from the pool
func (b *RabbitMQBackend) getChannel() (*amqp.Channel, error) {
	select {
	case ch := <-b.channels:
		return ch, nil
	case <-time.After(5 * time.Second):
		// Pool exhausted, create new channel
		return b.conn.Channel()
	}
}

// returnChannel returns a channel to the pool
func (b *RabbitMQBackend) returnChannel(ch *amqp.Channel) {
	select {
	case b.channels <- ch:
		// Successfully returned to pool
	default:
		// Pool full, close channel
		ch.Close()
	}
}

// setupDLX declares Dead Letter Exchange and Dead Letter Queue
func (b *RabbitMQBackend) setupDLX(ch *amqp.Channel) error {
	// Declare DLX exchange
	err := ch.ExchangeDeclare(
		b.config.dlxName, // name
		"direct",         // type
		b.config.durable, // durable
		false,            // auto-delete
		false,            // internal
		false,            // no-wait
		nil,              // arguments
	)
	if err != nil {
		return fmt.Errorf("failed to declare DLX: %w", err)
	}

	// Declare DLQ queue (no DLX on DLQ itself to avoid loops)
	_, err = ch.QueueDeclare(
		b.config.dlqName, // name
		b.config.durable, // durable
		false,            // auto-delete
		false,            // exclusive
		false,            // no-wait
		nil,              // arguments
	)
	if err != nil {
		return fmt.Errorf("failed to declare DLQ: %w", err)
	}

	// Bind DLQ to DLX
	routingKey := b.config.dlxRoutingKey
	if routingKey == "" {
		routingKey = "#" // Catch all
	}

	err = ch.QueueBind(
		b.config.dlqName, // queue
		routingKey,       // routing key
		b.config.dlxName, // exchange
		false,            // no-wait
		nil,              // arguments
	)
	if err != nil {
		return fmt.Errorf("failed to bind DLQ to DLX: %w", err)
	}

	return nil
}

// Publish sends a message envelope to the specified queue
func (b *RabbitMQBackend) Publish(ctx context.Context, queue string, envelope *Envelope) error {
	// Check if backend is closed
	b.mu.RLock()
	if b.closed {
		b.mu.RUnlock()
		return ErrQueueStopped
	}
	b.mu.RUnlock()

	// Get channel from pool
	ch, err := b.getChannel()
	if err != nil {
		return fmt.Errorf("failed to get channel: %w", err)
	}
	defer b.returnChannel(ch)

	// Prepare queue arguments
	args := amqp.Table{}

	// Add DLX configuration if enabled
	if b.config.enableDLX {
		args["x-dead-letter-exchange"] = b.config.dlxName
		if b.config.dlxRoutingKey != "" {
			args["x-dead-letter-routing-key"] = b.config.dlxRoutingKey
		}
	}

	// Add priority queue support
	if b.config.enablePriority {
		args["x-max-priority"] = b.config.maxPriority
	}

	// Declare queue
	_, err = ch.QueueDeclare(
		queue,                // name
		b.config.durable,     // durable
		b.config.autoDelete,  // auto-delete
		b.config.exclusive,   // exclusive
		b.config.noWait,      // no-wait
		args,                 // arguments
	)
	if err != nil {
		return fmt.Errorf("failed to declare queue: %w", err)
	}

	// Serialize envelope to JSON
	data, err := json.Marshal(envelope)
	if err != nil {
		return fmt.Errorf("failed to serialize envelope: %w", err)
	}

	// Prepare delivery mode
	deliveryMode := amqp.Transient
	if b.config.persistent {
		deliveryMode = amqp.Persistent
	}

	// Calculate priority from retry count (higher retries = lower priority)
	priority := uint8(0)
	if b.config.enablePriority && envelope.RetryCount > 0 {
		maxPriority := int(b.config.maxPriority)
		priority = uint8(max(0, maxPriority-envelope.RetryCount))
	}

	// Publish to queue using default exchange
	err = ch.PublishWithContext(
		ctx,
		"",                   // exchange (default)
		queue,                // routing key (queue name)
		false,                // mandatory
		false,                // immediate
		amqp.Publishing{
			DeliveryMode: deliveryMode,
			ContentType:  "application/json",
			MessageId:    envelope.ID,
			Timestamp:    envelope.CreatedAt,
			Body:         data,
			Priority:     priority,
			Headers:      convertMetadata(envelope.Metadata),
		},
	)

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

	return nil
}

// max returns the maximum of two integers
func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// Subscribe creates a subscription to the specified queue and returns a channel
func (b *RabbitMQBackend) Subscribe(ctx context.Context, queue string) (<-chan *Envelope, error) {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.closed {
		return nil, ErrQueueStopped
	}

	// Check if already subscribed
	if ch, exists := b.subscribers[queue]; exists {
		return ch, nil
	}

	// Create dedicated channel for this queue
	ch, err := b.conn.Channel()
	if err != nil {
		return nil, fmt.Errorf("failed to create channel: %w", err)
	}

	// Set QoS (prefetch count)
	err = ch.Qos(
		b.config.prefetchCount, // prefetch count
		b.config.prefetchSize,  // prefetch size
		false,                  // global
	)
	if err != nil {
		ch.Close()
		return nil, fmt.Errorf("failed to set QoS: %w", err)
	}

	// Prepare queue arguments
	args := amqp.Table{}
	if b.config.enableDLX {
		args["x-dead-letter-exchange"] = b.config.dlxName
		if b.config.dlxRoutingKey != "" {
			args["x-dead-letter-routing-key"] = b.config.dlxRoutingKey
		}
	}
	if b.config.enablePriority {
		args["x-max-priority"] = b.config.maxPriority
	}

	// Declare queue with same configuration as Publish
	_, err = ch.QueueDeclare(
		queue,
		b.config.durable,
		b.config.autoDelete,
		b.config.exclusive,
		b.config.noWait,
		args,
	)
	if err != nil {
		ch.Close()
		return nil, fmt.Errorf("failed to declare queue: %w", err)
	}

	// Setup DLX and DLQ if enabled
	if b.config.enableDLX {
		if err := b.setupDLX(ch); err != nil {
			ch.Close()
			return nil, err
		}
	}

	// Generate consumer tag
	consumerTag := fmt.Sprintf("goqueue-%s-%d", queue, time.Now().UnixNano())

	// Start consuming
	deliveries, err := ch.Consume(
		queue,                // queue
		consumerTag,          // consumer tag
		false,                // auto-ack (false - manual ack)
		b.config.exclusive,   // exclusive
		false,                // no-local
		b.config.noWait,      // no-wait
		nil,                  // args
	)
	if err != nil {
		ch.Close()
		return nil, fmt.Errorf("failed to start consumer: %w", err)
	}

	// Create Go channel and start adapter goroutine
	envChan := make(chan *Envelope, 100)
	stopChan := make(chan struct{})

	b.subscribers[queue] = envChan
	b.consumerTags[queue] = consumerTag
	b.stopChans[queue] = stopChan
	b.queueChannels[queue] = ch

	// Start goroutine to adapt RabbitMQ deliveries to Go channel
	go b.consumeLoop(ctx, queue, ch, deliveries, envChan, stopChan)

	return envChan, nil
}

// consumeLoop adapts RabbitMQ deliveries to Go channels
func (b *RabbitMQBackend) consumeLoop(
	ctx context.Context,
	queue string,
	ch *amqp.Channel,
	deliveries <-chan amqp.Delivery,
	envChan chan<- *Envelope,
	stopChan <-chan struct{},
) {
	defer close(envChan)

	for {
		select {
		case <-stopChan:
			return
		case <-ctx.Done():
			return
		case d, ok := <-deliveries:
			if !ok {
				// Channel closed, exit
				return
			}

			// Deserialize envelope
			var envelope Envelope
			if err := json.Unmarshal(d.Body, &envelope); err != nil {
				// Nack and discard malformed message
				d.Nack(false, false)
				continue
			}

			// Store delivery tag mapping
			b.trackingMu.Lock()
			b.messageTracking[envelope.ID] = &messageTracking{
				deliveryTag: d.DeliveryTag,
				queue:       queue,
			}
			b.trackingMu.Unlock()

			// Send to worker pool
			select {
			case envChan <- &envelope:
				// Successfully sent
			case <-stopChan:
				return
			case <-ctx.Done():
				return
			}
		}
	}
}

// Ack acknowledges successful processing of a message
func (b *RabbitMQBackend) Ack(ctx context.Context, messageID string) error {
	b.mu.RLock()
	if b.closed {
		b.mu.RUnlock()
		return ErrQueueStopped
	}
	b.mu.RUnlock()

	// Get tracking info
	b.trackingMu.Lock()
	tracking, exists := b.messageTracking[messageID]
	if !exists {
		b.trackingMu.Unlock()
		return ErrMessageNotFound
	}
	delete(b.messageTracking, messageID)
	b.trackingMu.Unlock()

	// Get channel for queue
	b.queueChanMu.RLock()
	ch, exists := b.queueChannels[tracking.queue]
	b.queueChanMu.RUnlock()

	if !exists {
		return fmt.Errorf("channel not found for queue %s", tracking.queue)
	}

	// Acknowledge message
	return ch.Ack(tracking.deliveryTag, false)
}

// Nack indicates that a message failed to process
func (b *RabbitMQBackend) Nack(ctx context.Context, messageID string) error {
	b.mu.RLock()
	if b.closed {
		b.mu.RUnlock()
		return ErrQueueStopped
	}
	b.mu.RUnlock()

	// Get tracking info
	b.trackingMu.Lock()
	tracking, exists := b.messageTracking[messageID]
	if !exists {
		b.trackingMu.Unlock()
		return ErrMessageNotFound
	}
	delete(b.messageTracking, messageID)
	b.trackingMu.Unlock()

	// Get channel for queue
	b.queueChanMu.RLock()
	ch, exists := b.queueChannels[tracking.queue]
	b.queueChanMu.RUnlock()

	if !exists {
		return fmt.Errorf("channel not found for queue %s", tracking.queue)
	}

	// Nack with requeue=false (routes to DLX if configured)
	return ch.Nack(tracking.deliveryTag, false, false)
}

// Close releases all resources held by the backend
func (b *RabbitMQBackend) Close() error {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.closed {
		return nil
	}

	b.closed = true

	// Stop all consumer loops
	for _, stopChan := range b.stopChans {
		close(stopChan)
	}

	// Cancel all consumers and close queue channels
	b.queueChanMu.Lock()
	for queue, ch := range b.queueChannels {
		if tag, exists := b.consumerTags[queue]; exists {
			ch.Cancel(tag, false)
		}
		ch.Close()
	}
	b.queueChanMu.Unlock()

	// Close channel pool
	b.channelsMu.Lock()
	close(b.channels)
	for ch := range b.channels {
		ch.Close()
	}
	b.channelsMu.Unlock()

	// Close connection
	return b.conn.Close()
}

// Ping checks the connection to RabbitMQ
func (b *RabbitMQBackend) Ping(ctx context.Context) error {
	b.mu.RLock()
	defer b.mu.RUnlock()

	if b.closed {
		return ErrQueueStopped
	}

	// Try to create a temporary channel to verify connection
	ch, err := b.conn.Channel()
	if err != nil {
		return fmt.Errorf("failed to ping RabbitMQ: %w", err)
	}
	defer ch.Close()

	return nil
}