go-caatsm/internal/infra/nats/consumer.go

package nats

import (
	"caatsm/internal/app"
	"caatsm/internal/infra/config"
	obslogging "caatsm/internal/observability/logging"
	obsmetrics "caatsm/internal/observability/metrics"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/nats-io/nats.go"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/metric"
	"go.uber.org/zap"
)

// Consumer handles NATS JetStream message consumption
type Consumer struct {
	conn            *nats.Conn
	js              nats.JetStreamContext
	processor       *app.MessageProcessor
	cfg             *config.Config
	logger          *zap.Logger
	subject         string
	consumerName    string
	mode            string
	streamName      string
	dlqSubject      string
	ackWait         time.Duration
	batchSize       int
	batchTimeout    time.Duration
	monitorInterval time.Duration
	meter           metric.Meter
	ackPending      metric.Int64Histogram
	redelivered     metric.Int64Histogram
	pending         metric.Int64Histogram
	delivered       metric.Int64Histogram

	// simple backpressure / degradation state
	consecutiveProcessErrors int
}

// ProvideConsumer creates a NATS consumer
func ProvideConsumer(
	conn *nats.Conn,
	js nats.JetStreamContext,
	processor *app.MessageProcessor,
	cfg *config.Config,
	logger *zap.Logger,
) (*Consumer, error) {
	subject := cfg.EffectiveSubscriptionTopic()

	consumerName := cfg.NATS.Consumer
	if consumerName == "" {
		consumerName = "telegram-consumer"
	}

	mode := strings.ToLower(cfg.NATS.Mode)
	if mode == "" {
		mode = "jetstream"
	}

	streamName := cfg.NATS.Stream
	if streamName == "" {
		streamName = "TELEGRAM"
	}

	// DLQ routing is only meaningful in JetStream mode. Respect dlq.enabled to allow
	// environments to opt out cleanly even if a subject is configured.
	dlqSubject := ""
	if cfg.DLQ.Enabled {
		dlqSubject = strings.TrimSpace(cfg.DLQ.Subject)
	}

	ackWait := cfg.NATS.ConsumerRules.AckWait
	if ackWait == 0 {
		ackWait = cfg.Timeouts.AckWait
	}
	if ackWait == 0 {
		ackWait = 30 * time.Second
	}

	batchSize := cfg.App.BatchSize
	if batchSize == 0 {
		batchSize = 50
	}

	batchTimeout := cfg.App.BatchTimeout
	if batchTimeout == 0 {
		batchTimeout = 2 * time.Second
	}

	monitorInterval := cfg.App.MonitorInterval
	if monitorInterval <= 0 {
		monitorInterval = 30 * time.Second
	}

	consumer := &Consumer{
		conn:            conn,
		js:              js,
		processor:       processor,
		cfg:             cfg,
		logger:          logger,
		subject:         subject,
		consumerName:    consumerName,
		mode:            mode,
		streamName:      streamName,
		dlqSubject:      dlqSubject,
		ackWait:         ackWait,
		batchSize:       batchSize,
		batchTimeout:    batchTimeout,
		monitorInterval: monitorInterval,
	}
	consumer.initMetrics()

	if consumer.mode == "jetstream" {
		// Create consumer if it doesn't exist
		if err := consumer.ensureConsumer(); err != nil {
			return nil, fmt.Errorf("failed to ensure consumer: %w", err)
		}
		// Validate DLQ configuration early so misconfiguration is visible at startup
		// rather than only when the first poison message appears.
		consumer.validateDLQ()
	} else {
		logger.Info("Running consumer in core NATS mode",
			zap.String("subject", subject),
			zap.String("queue_group", cfg.Subscription.QueueGroup),
		)
	}

	return consumer, nil
}

// ensureConsumer creates the consumer if it doesn't exist; if it already exists, it is reused.
func (c *Consumer) ensureConsumer() error {
	consumerConfig := &nats.ConsumerConfig{
		Durable:       c.consumerName,
		DeliverPolicy: mapDeliverPolicy(c.cfg.NATS.ConsumerRules.DeliverPolicy),
		AckPolicy:     nats.AckExplicitPolicy,
		AckWait:       c.ackWait,
		ReplayPolicy:  mapReplayPolicy(c.cfg.NATS.ConsumerRules.ReplayPolicy),
		MaxDeliver:    c.cfg.NATS.ConsumerRules.MaxDeliver,
		MaxAckPending: c.cfg.NATS.ConsumerRules.MaxAckPending,
		FilterSubject: c.subject,
		BackOff:       c.cfg.NATS.ConsumerRules.Backoff,
	}
	if consumerConfig.DeliverPolicy == nats.DeliverByStartSequencePolicy && c.cfg.NATS.ConsumerRules.StartSequence > 0 {
		consumerConfig.OptStartSeq = c.cfg.NATS.ConsumerRules.StartSequence
	}
	if consumerConfig.DeliverPolicy == nats.DeliverByStartTimePolicy && strings.TrimSpace(c.cfg.NATS.ConsumerRules.StartTime) != "" {
		startTime, err := time.Parse(time.RFC3339, c.cfg.NATS.ConsumerRules.StartTime)
		if err != nil {
			c.logger.Warn("Invalid start time, falling back to deliver policy defaults",
				zap.String("start_time", c.cfg.NATS.ConsumerRules.StartTime),
				zap.Error(err),
			)
		} else {
			consumerConfig.OptStartTime = &startTime
		}
	}

	// First check if the consumer already exists to make this initialization idempotent.
	info, err := c.js.ConsumerInfo(c.streamName, c.consumerName)
	if err == nil && info != nil {
		c.logger.Info("Using existing JetStream consumer",
			zap.String("consumer", c.consumerName),
			zap.String("stream", c.streamName),
			zap.String("subject", c.subject),
		)
		return nil
	}
	if err != nil && !errors.Is(err, nats.ErrConsumerNotFound) {
		return fmt.Errorf("failed to fetch consumer info: %w", err)
	}

	// Consumer does not exist; create it.
	if _, err := c.js.AddConsumer(c.streamName, consumerConfig); err != nil {
		return fmt.Errorf("failed to create consumer: %w", err)
	}

	c.logger.Info("Created JetStream consumer",
		zap.String("consumer", c.consumerName),
		zap.String("stream", c.streamName),
		zap.String("subject", c.subject),
		zap.Duration("ack_wait", c.ackWait),
		zap.String("deliver_policy", c.cfg.NATS.ConsumerRules.DeliverPolicy),
		zap.String("replay_policy", c.cfg.NATS.ConsumerRules.ReplayPolicy),
	)

	return nil
}

// validateDLQ verifies whether DLQ routing should be enabled and, if so, whether
// the configured DLQ subject is bound to a JetStream stream. If validation fails,
// DLQ routing is disabled (by clearing c.dlqSubject) and a warning is logged,
// but the consumer is still allowed to start.
func (c *Consumer) validateDLQ() {
	if c == nil {
		return
	}

	// DLQ routing is only active in JetStream mode.
	if c.mode != "jetstream" {
		return
	}

	// If DLQ is not enabled in config, make sure we don't accidentally route to it.
	if !c.cfg.DLQ.Enabled {
		if strings.TrimSpace(c.dlqSubject) != "" {
			c.logger.Info("DLQ subject configured but dlq.enabled is false; DLQ routing disabled",
				zap.String("dlq_subject", c.dlqSubject),
			)
		}
		c.dlqSubject = ""
		return
	}

	subject := strings.TrimSpace(c.dlqSubject)
	if subject == "" {
		c.logger.Warn("DLQ enabled but dlq.subject is empty; DLQ routing disabled")
		return
	}

	if c.js == nil {
		c.logger.Warn("DLQ enabled but JetStream context is nil; DLQ routing disabled",
			zap.String("dlq_subject", subject),
		)
		c.dlqSubject = ""
		return
	}

	// Ensure the DLQ subject is actually bound to a JetStream stream. This avoids
	// the opaque `nats: no response from stream` error later when publishing.
	obsmetrics.RecordJSAPICall("dlq_validate_stream")
	streamName, err := c.js.StreamNameBySubject(subject)
	if err != nil || strings.TrimSpace(streamName) == "" {
		c.logger.Warn("DLQ subject not bound to any JetStream stream; DLQ routing disabled",
			zap.String("dlq_subject", subject),
			zap.Error(err),
		)
		c.dlqSubject = ""
		return
	}

	c.logger.Info("DLQ configuration validated",
		zap.String("dlq_subject", subject),
		zap.String("dlq_stream", streamName),
	)
}

// Start starts consuming messages
func (c *Consumer) Start(ctx context.Context) error {
	if c.mode == "core" {
		return c.startCore(ctx)
	}

	return c.startJetStream(ctx)
}

func (c *Consumer) startJetStream(ctx context.Context) error {
	// Create pull subscription
	sub, err := c.js.PullSubscribe(c.subject, c.consumerName, nats.Bind(c.streamName, c.consumerName))
	if err != nil {
		return fmt.Errorf("failed to create pull subscription: %w", err)
	}
	defer sub.Unsubscribe()

	c.logger.Info("Started consuming messages",
		zap.String("subject", c.subject),
		zap.String("consumer", c.consumerName),
		zap.String("stream", c.streamName),
	)

	c.logger.Info("Consumer pull configuration",
		zap.Int("batch_size", c.batchSize),
		zap.Duration("batch_timeout", c.batchTimeout),
		zap.Int("max_deliver", c.cfg.NATS.ConsumerRules.MaxDeliver),
		zap.Duration("ack_wait", c.ackWait),
		zap.String("deliver_policy", c.cfg.NATS.ConsumerRules.DeliverPolicy),
		zap.String("replay_policy", c.cfg.NATS.ConsumerRules.ReplayPolicy),
		zap.Int("backoff_steps", len(c.cfg.NATS.ConsumerRules.Backoff)),
	)

	statsCtx, statsCancel := context.WithCancel(ctx)
	defer statsCancel()
	go c.emitConsumerStats(statsCtx)

	for {
		select {
		case <-ctx.Done():
			c.logger.Info("Stopping consumer", zap.Error(ctx.Err()))
			return ctx.Err()
		default:
		}

		// Fetch messages in batch
		msgs, err := sub.Fetch(c.batchSize, nats.MaxWait(c.batchTimeout))
		if err != nil {
			if errors.Is(err, nats.ErrTimeout) {
				// Timeout is expected when no messages are available
				continue
			}
			if errors.Is(err, nats.ErrNoResponders) {
				// JetStream API is currently unavailable (e.g., NATS just restarted or JetStream not ready).
				// Back off a bit to avoid log spam while allowing the system to recover.
				c.logger.Warn("JetStream not available, will retry",
					zap.Error(err),
					zap.String("stream", c.streamName),
					zap.String("consumer", c.consumerName),
				)
				time.Sleep(5 * time.Second)
				continue
			}
			c.logger.Error("Failed to fetch messages", zap.Error(err))
			time.Sleep(time.Second)
			continue
		}

		// Process each message
		// TODO: consider buffering messages to take advantage of Repository.InsertBatch for higher throughput.
		for _, msg := range msgs {
			start := time.Now()

			if err := c.processMessage(ctx, msg); err != nil {
				isPermanent := app.IsPermanent(err)
				elapsed := time.Since(start)

				c.logger.Error("Failed to process message",
					zap.String("subject", msg.Subject),
					zap.Error(err),
					zap.Bool("permanent", isPermanent),
				)

				result := obsmetrics.ResultFail
				if isPermanent {
					result = obsmetrics.ResultPermanentFail
				}
				obsmetrics.RecordMessageHandled(c.streamName, c.consumerName, result, elapsed)

				if isPermanent {
					c.consecutiveProcessErrors = 0
					// Poison/permanent message: route to DLQ if configured, then ACK
					if err := c.routeToDLQ(ctx, msg, err); err != nil {
						c.logger.Error("Failed to route permanent-error message to DLQ", zap.Error(err))
					}
					if ackErr := msg.Ack(); ackErr != nil {
						c.logger.Error("Failed to ACK permanent-error message", zap.Error(ackErr))
					}
					continue
				}

				// Transient error: increment error streak and apply simple backpressure if needed.
				if c.consecutiveProcessErrors < 0 {
					c.consecutiveProcessErrors = 0
				}
				c.consecutiveProcessErrors++
				if c.consecutiveProcessErrors >= 10 {
					// Apply a brief sleep to slow down consumption when the system
					// is failing many messages in a row (e.g. DB unavailable).
					backoff := time.Duration(c.consecutiveProcessErrors) * 100 * time.Millisecond
					if backoff > 5*time.Second {
						backoff = 5 * time.Second
					}
					c.logger.Warn("Applying backpressure due to consecutive processing errors",
						zap.Int("consecutive_errors", c.consecutiveProcessErrors),
						zap.Duration("sleep", backoff),
					)
					time.Sleep(backoff)
				}

				// Transient error: request redelivery with optional delay
				obsmetrics.RecordRetry(c.streamName, c.consumerName, obsmetrics.RetryReasonProcessorError)
				if nakErr := c.nakWithStrategy(msg); nakErr != nil {
					c.logger.Error("Failed to NAK message", zap.Error(nakErr))
				}
				continue
			}

			// Successful processing resets the error streak.
			if c.consecutiveProcessErrors > 0 {
				c.consecutiveProcessErrors = 0
			}

			// ACK the message
			if ackErr := msg.Ack(); ackErr != nil {
				c.logger.Error("Failed to ACK message", zap.Error(ackErr))
			} else {
				elapsed := time.Since(start)
				obsmetrics.RecordMessageHandled(c.streamName, c.consumerName, "ok", elapsed)
			}
		}
	}
}

func (c *Consumer) startCore(ctx context.Context) error {
	queueGroup := c.cfg.Subscription.QueueGroup
	if queueGroup == "" {
		queueGroup = c.consumerName
	}

	handler := func(msg *nats.Msg) {
		if err := c.processMessage(ctx, msg); err != nil {
			isPermanent := app.IsPermanent(err)
			c.logger.Error("Failed to process message (core mode)",
				zap.String("subject", msg.Subject),
				zap.Error(err),
				zap.Bool("permanent", isPermanent),
			)
		}
	}

	sub, err := c.conn.QueueSubscribe(c.subject, queueGroup, handler)
	if err != nil {
		return fmt.Errorf("failed to subscribe to %s: %w", c.subject, err)
	}
	if err := c.conn.Flush(); err != nil {
		return fmt.Errorf("failed to flush NATS connection: %w", err)
	}

	c.logger.Info("Started core NATS subscription",
		zap.String("subject", c.subject),
		zap.String("queue_group", queueGroup),
	)

	<-ctx.Done()
	c.logger.Info("Stopping core NATS consumer", zap.Error(ctx.Err()))

	if err := sub.Drain(); err != nil && !errors.Is(err, nats.ErrConnectionClosed) {
		return fmt.Errorf("failed to drain core subscription: %w", err)
	}

	return ctx.Err()
}

func (c *Consumer) emitConsumerStats(ctx context.Context) {
	ticker := time.NewTicker(c.monitorInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			info, err := c.js.ConsumerInfo(c.streamName, c.consumerName)
			if err != nil {
				c.logger.Warn("Failed to fetch consumer info", zap.Error(err))
				continue
			}

			c.logger.Debug("JetStream consumer metrics",
				zap.String("stream", c.streamName),
				zap.String("consumer", c.consumerName),
				zap.Uint64("num_ack_pending", uint64(info.NumAckPending)),
				zap.Uint64("num_redelivered", uint64(info.NumRedelivered)),
				zap.Uint64("num_pending", uint64(info.NumPending)),
				zap.Uint64("delivered_consumer_seq", uint64(info.Delivered.Consumer)),
				zap.Uint64("delivered_stream_seq", uint64(info.Delivered.Stream)),
			)
			c.recordConsumerMetrics(ctx, info)
		}
	}
}

// Shutdown drains the underlying NATS connection gracefully.
func (c *Consumer) Shutdown(ctx context.Context) error {
	if c.conn == nil {
		return nil
	}

	timeout := c.cfg.Timeouts.Close
	if timeout <= 0 {
		timeout = 10 * time.Second
	}

	closeCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	errCh := make(chan error, 1)
	go func() {
		errCh <- c.conn.Drain()
	}()

	select {
	case err := <-errCh:
		c.conn.Close()
		return err
	case <-closeCtx.Done():
		c.conn.Close()
		return fmt.Errorf("nats drain timeout: %w", closeCtx.Err())
	}
}

func (c *Consumer) initMetrics() {
	meter := otel.Meter("caatsm/nats")
	c.meter = meter

	if hist, err := meter.Int64Histogram("caatsm_nats_consumer_ack_pending"); err == nil {
		c.ackPending = hist
	}
	if hist, err := meter.Int64Histogram("caatsm_nats_consumer_redelivered"); err == nil {
		c.redelivered = hist
	}
	if hist, err := meter.Int64Histogram("caatsm_nats_consumer_pending"); err == nil {
		c.pending = hist
	}
	if hist, err := meter.Int64Histogram("caatsm_nats_consumer_delivered"); err == nil {
		c.delivered = hist
	}
}

func (c *Consumer) recordConsumerMetrics(ctx context.Context, info *nats.ConsumerInfo) {
	if info == nil {
		return
	}
	if c.ackPending != nil {
		c.ackPending.Record(ctx, int64(info.NumAckPending))
	}
	if c.redelivered != nil {
		c.redelivered.Record(ctx, int64(info.NumRedelivered))
	}
	if c.pending != nil {
		c.pending.Record(ctx, int64(info.NumPending))
	}
	if c.delivered != nil {
		c.delivered.Record(ctx, int64(info.Delivered.Stream))
	}
}

// routeToDLQ publishes a copy of the failed message to the configured DLQ subject,
// including useful metadata for offline analysis. If DLQ is not configured or the
// consumer is not running in JetStream mode, this is a no-op.
func (c *Consumer) routeToDLQ(ctx context.Context, msg *nats.Msg, cause error) error {
	if c == nil || c.js == nil {
		return nil
	}
	if c.mode != "jetstream" {
		return nil
	}
	if strings.TrimSpace(c.dlqSubject) == "" {
		return nil
	}

	meta, _ := msg.Metadata()
	jsSeq := uint64(0)
	deliveries := uint64(0)
	if meta != nil {
		jsSeq = meta.Sequence.Stream
		deliveries = meta.NumDelivered
	}

	payload := map[string]interface{}{
		"transport_msg_id": msg.Header.Get("Nats-Msg-Id"),
		"subject":          msg.Subject,
		"stream":           c.streamName,
		"consumer":         c.consumerName,
		"nats_sequence":    jsSeq,
		"deliveries":       deliveries,
		"error":            fmt.Sprint(cause),
		"received_at":      time.Now().UTC(),
		"body":             string(msg.Data),
	}

	data, err := json.Marshal(payload)
	if err != nil {
		return fmt.Errorf("marshal dlq payload: %w", err)
	}

	if _, err := c.js.Publish(c.dlqSubject, data); err != nil {
		// nats.ErrNoResponders typically means that no JetStream stream is
		// configured to receive this subject, or JetStream is temporarily
		// unavailable. Surface this explicitly to make operational diagnosis
		// easier.
		if errors.Is(err, nats.ErrNoResponders) {
			obsmetrics.RecordDLQPublishFailure(c.streamName, c.consumerName)
			return fmt.Errorf("publish to dlq subject %s: no JetStream stream found for subject or JetStream unavailable: %w", c.dlqSubject, err)
		}
		obsmetrics.RecordDLQPublishFailure(c.streamName, c.consumerName)
		return fmt.Errorf("publish to dlq subject %s: %w", c.dlqSubject, err)
	}

	obsmetrics.RecordDLQMessage(c.streamName, c.consumerName)

	return nil
}

func (c *Consumer) nakWithStrategy(msg *nats.Msg) error {
	backoff := c.cfg.NATS.ConsumerRules.Backoff
	if len(backoff) == 0 {
		return msg.Nak()
	}

	meta, err := msg.Metadata()
	if err != nil {
		c.logger.Warn("Failed to read metadata for backoff strategy", zap.Error(err))
		return msg.Nak()
	}

	attempt := int(meta.NumDelivered)
	index := attempt - 1
	if index < 0 {
		index = 0
	}
	if index >= len(backoff) {
		index = len(backoff) - 1
	}
	delay := backoff[index]
	if delay <= 0 {
		return msg.Nak()
	}

	return msg.NakWithDelay(delay)
}

// processMessage processes a single message
func (c *Consumer) processMessage(ctx context.Context, msg *nats.Msg) error {
	ctx, span := otel.Tracer("caatsm/nats").Start(ctx, "Consumer.processMessage")
	defer span.End()
	span.SetAttributes(attribute.String("nats.subject", msg.Subject))

	msgID, source, err := c.resolveMsgID(msg)
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
		return fmt.Errorf("unable to resolve message id: %w", err)
	}
	if source != "header" {
		c.logger.Warn("Message missing NATS id header; using fallback",
			zap.String("subject", msg.Subject),
			zap.String("msg_id_source", source),
			zap.String("msg_id", msgID),
		)
	}

	// Attach structured logging context including stream/consumer and NATS metadata.
	jsSeq := uint64(0)
	if meta, metaErr := msg.Metadata(); metaErr == nil {
		jsSeq = meta.Sequence.Stream
		span.SetAttributes(
			attribute.Int64("nats.js.stream_seq", int64(meta.Sequence.Stream)),
			attribute.Int64("nats.js.consumer_seq", int64(meta.Sequence.Consumer)),
		)
	}

	msgLogger := obslogging.WithMessageContext(c.logger, obslogging.MessageFields{
		Service:        "caatsm-consumer",
		TransportMsgID: msgID,
		Stream:         c.streamName,
		Consumer:       c.consumerName,
		Subject:        msg.Subject,
		JSSequence:     jsSeq,
	})

	msgLogger.Debug("Processing message",
		zap.Int("data_size", len(msg.Data)),
		zap.String("msg_id_source", source),
	)

	// Call processor
	if err := c.processor.Handle(ctx, msg.Data, msgID); err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, err.Error())
		return fmt.Errorf("processor error: %w", err)
	}

	span.SetAttributes(attribute.String("telegram.msg_id", msgID))
	return nil
}

func (c *Consumer) resolveMsgID(msg *nats.Msg) (string, string, error) {
	if id := msg.Header.Get("Nats-Msg-Id"); id != "" {
		return id, "header", nil
	}

	if c.mode == "core" {
		return uuid.NewString(), "generated", nil
	}

	meta, err := msg.Metadata()
	if err != nil {
		return "", "", fmt.Errorf("fetch metadata: %w", err)
	}

	return fmt.Sprintf("js-%d", meta.Sequence.Stream), "metadata", nil
}

func mapDeliverPolicy(value string) nats.DeliverPolicy {
	switch strings.ToLower(value) {
	case "new":
		return nats.DeliverNewPolicy
	case "last":
		return nats.DeliverLastPolicy
	case "last_per_subject":
		return nats.DeliverLastPerSubjectPolicy
	case "sequence":
		return nats.DeliverByStartSequencePolicy
	case "time":
		return nats.DeliverByStartTimePolicy
	default:
		return nats.DeliverAllPolicy
	}
}

func mapReplayPolicy(value string) nats.ReplayPolicy {
	switch strings.ToLower(value) {
	case "original":
		return nats.ReplayOriginalPolicy
	default:
		return nats.ReplayInstantPolicy
	}
}