modelRT/task/queue_producer.go

// Package task provides asynchronous task processing with RabbitMQ integration
package task

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

	"modelRT/config"
	"modelRT/logger"
	"modelRT/mq"

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

const (
	// TaskExchangeName is the name of the exchange for task routing
	TaskExchangeName = "modelrt.tasks.exchange"
	// TaskQueueName is the name of the main task queue
	TaskQueueName = "modelrt.tasks.queue"
	// TaskRoutingKey is the routing key for task messages
	TaskRoutingKey = "modelrt.task"
	// MaxPriority is the maximum priority level for tasks (0-10)
	MaxPriority = 10
	// DefaultMessageTTL is the default time-to-live for task messages (24 hours)
	DefaultMessageTTL = 24 * time.Hour
)

// QueueProducer handles publishing tasks to RabbitMQ
type QueueProducer struct {
	conn *amqp.Connection
	ch   *amqp.Channel
}

// NewQueueProducer creates a new QueueProducer instance
func NewQueueProducer(ctx context.Context, cfg config.RabbitMQConfig) (*QueueProducer, error) {
	// Initialize RabbitMQ connection if not already initialized
	mq.InitRabbitProxy(ctx, cfg)

	conn := mq.GetConn()
	if conn == nil {
		return nil, fmt.Errorf("failed to get RabbitMQ connection")
	}

	ch, err := conn.Channel()
	if err != nil {
		return nil, fmt.Errorf("failed to open channel: %w", err)
	}

	producer := &QueueProducer{
		conn: conn,
		ch:   ch,
	}

	// Declare exchange and queue
	if err := producer.declareInfrastructure(); err != nil {
		ch.Close()
		return nil, fmt.Errorf("failed to declare infrastructure: %w", err)
	}

	return producer, nil
}

// declareInfrastructure declares the exchange, queue, and binds them
func (p *QueueProducer) declareInfrastructure() error {
	// Declare durable direct exchange
	err := p.ch.ExchangeDeclare(
		TaskExchangeName, // name
		"direct",         // type
		true,             // durable
		false,            // auto-deleted
		false,            // internal
		false,            // no-wait
		nil,              // arguments
	)
	if err != nil {
		return fmt.Errorf("failed to declare exchange: %w", err)
	}

	// Declare durable queue with priority support and message TTL
	_, err = p.ch.QueueDeclare(
		TaskQueueName, // name
		true,          // durable
		false,         // delete when unused
		false,         // exclusive
		false,         // no-wait
		amqp.Table{
			"x-max-priority": MaxPriority,               // support priority levels 0-10
			"x-message-ttl":  DefaultMessageTTL.Milliseconds(), // message TTL
		},
	)
	if err != nil {
		return fmt.Errorf("failed to declare queue: %w", err)
	}

	// Bind queue to exchange
	err = p.ch.QueueBind(
		TaskQueueName,    // queue name
		TaskRoutingKey,   // routing key
		TaskExchangeName, // exchange name
		false,            // no-wait
		nil,              // arguments
	)
	if err != nil {
		return fmt.Errorf("failed to bind queue: %w", err)
	}

	return nil
}

// PublishTask publishes a task message to RabbitMQ
func (p *QueueProducer) PublishTask(ctx context.Context, taskID uuid.UUID, taskType TaskType, priority int) error {
	message := NewTaskQueueMessageWithPriority(taskID, taskType, priority)

	// Validate message
	if !message.Validate() {
		return fmt.Errorf("invalid task message: taskID=%s, taskType=%s", taskID, taskType)
	}

	// Convert message to JSON
	body, err := json.Marshal(message)
	if err != nil {
		return fmt.Errorf("failed to marshal task message: %w", err)
	}

	// Prepare publishing options
	publishing := amqp.Publishing{
		ContentType:  "application/json",
		Body:         body,
		DeliveryMode: amqp.Persistent, // Persistent messages survive broker restart
		Timestamp:    time.Now(),
		Priority:     uint8(priority),
		Headers: amqp.Table{
			"task_id":   taskID.String(),
			"task_type": string(taskType),
		},
	}

	// Publish to exchange
	err = p.ch.PublishWithContext(
		ctx,
		TaskExchangeName, // exchange
		TaskRoutingKey,   // routing key
		false,            // mandatory
		false,            // immediate
		publishing,
	)
	if err != nil {
		return fmt.Errorf("failed to publish task message: %w", err)
	}

	logger.Info(ctx, "Task published to queue",
		"task_id", taskID.String(),
		"task_type", taskType,
		"priority", priority,
		"queue", TaskQueueName,
	)

	return nil
}

// PublishTaskWithRetry publishes a task with retry logic
func (p *QueueProducer) PublishTaskWithRetry(ctx context.Context, taskID uuid.UUID, taskType TaskType, priority int, maxRetries int) error {
	var lastErr error
	for i := range maxRetries {
		err := p.PublishTask(ctx, taskID, taskType, priority)
		if err == nil {
			return nil
		}
		lastErr = err

		// Exponential backoff
		backoff := time.Duration(1<<uint(i)) * time.Second
		backoff = min(backoff, 10*time.Second)

		logger.Warn(ctx, "Failed to publish task, retrying",
			"task_id", taskID.String(),
			"attempt", i+1,
			"max_retries", maxRetries,
			"backoff", backoff,
			"error", err,
		)

		select {
		case <-time.After(backoff):
			continue
		case <-ctx.Done():
			return ctx.Err()
		}
	}

	return fmt.Errorf("failed to publish task after %d retries: %w", maxRetries, lastErr)
}

// Close closes the producer's channel
func (p *QueueProducer) Close() error {
	if p.ch != nil {
		return p.ch.Close()
	}
	return nil
}

// GetQueueInfo returns information about the task queue
func (p *QueueProducer) GetQueueInfo() (*amqp.Queue, error) {
	queue, err := p.ch.QueueDeclarePassive(
		TaskQueueName, // name
		true,          // durable
		false,         // delete when unused
		false,         // exclusive
		false,         // no-wait
		amqp.Table{
			"x-max-priority": MaxPriority,
			"x-message-ttl":  DefaultMessageTTL.Milliseconds(),
		},
	)
	if err != nil {
		return nil, fmt.Errorf("failed to inspect queue: %w", err)
	}
	return &queue, nil
}

// PurgeQueue removes all messages from the task queue
func (p *QueueProducer) PurgeQueue() (int, error) {
	return p.ch.QueuePurge(TaskQueueName, false)
}