package main

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

	"github.com/confluentinc/confluent-kafka-go/v2/kafka"
	"github.com/sirupsen/logrus"

	"github.com/tidbcloud/lake-ingest-kafka/config"
	"github.com/tidbcloud/lake-ingest-kafka/message"
)

type BatchReader interface {
	ReadBatch(ctx context.Context) (*message.MessagesBatch, error)

	Close() error
}

func NewBatchReader(cfg *config.Config) BatchReader {
	if cfg.MockData != "" {
		messageData := message.MessageData{}
		err := json.Unmarshal([]byte(cfg.MockData), &messageData)
		if err != nil {
			logrus.Fatalf("Failed to parse mock data: %v", err)
		}
		return NewMockBatchReader(messageData, cfg.BatchSize)
	}
	return NewKafkaBatchReader(cfg)
}

type MockBatchReader struct {
	sampleData message.MessageData
	batchSize  int
}

func NewMockBatchReader(sampleData message.MessageData, batchSize int) *MockBatchReader {
	return &MockBatchReader{
		sampleData: sampleData,
		batchSize:  batchSize,
	}
}

func (r *MockBatchReader) ReadBatch(ctx context.Context) (*message.MessagesBatch, error) {
	messages := make([]message.MessageData, 0, r.batchSize)
	for i := 0; i < r.batchSize; i++ {
		messages = append(messages, r.sampleData)
	}
	return &message.MessagesBatch{
		Messages:           messages,
		CommitFunc:         func(_ context.Context) error { return nil },
		FirstMessageOffset: -1,
		LastMessageOffset:  -1,
	}, nil
}

func (r *MockBatchReader) Close() error {
	return nil
}

type KafkaBatchReader struct {
	consumer         *kafka.Consumer
	topic            string
	batchSize        int
	maxBatchInterval int
	statsRecorder    *LakeConsumeStatsRecorder
}

func BuildKafkaConfigMap(cfg *config.Config) *kafka.ConfigMap {
	configMap := &kafka.ConfigMap{
		"bootstrap.servers":  cfg.KafkaBootstrapServers,
		"group.id":           cfg.KafkaConsumerGroup,
		"auto.offset.reset":  "earliest",
		"enable.auto.commit": false,
		"fetch.min.bytes":    cfg.MinBytes,
		"fetch.wait.max.ms":  cfg.MaxWait * 1000,
	}

	if cfg.EnableRebalanceOptimization {
		// Pick up topic changes (e.g. added partitions) without a restart and
		// rebalance the consumer group accordingly.
		strategy := cfg.PartitionAssignmentStrategy
		if strategy == "" {
			strategy = "range,roundrobin"
		}
		_ = configMap.SetKey("partition.assignment.strategy", strategy)

		refreshMs := cfg.TopicMetadataRefreshIntervalMs
		if refreshMs <= 0 {
			refreshMs = 60000
		}
		_ = configMap.SetKey("topic.metadata.refresh.interval.ms", refreshMs)
	}

	if cfg.IsSASL {
		_ = configMap.SetKey("sasl.mechanism", "PLAIN")
		_ = configMap.SetKey("sasl.username", cfg.SaslUser)
		_ = configMap.SetKey("sasl.password", cfg.SaslPassword)
		if cfg.DisableTLS {
			_ = configMap.SetKey("security.protocol", "SASL_PLAINTEXT")
		} else {
			_ = configMap.SetKey("security.protocol", "SASL_SSL")
		}
	} else {
		if cfg.DisableTLS {
			_ = configMap.SetKey("security.protocol", "PLAINTEXT")
		} else {
			_ = configMap.SetKey("security.protocol", "SSL")
		}
	}

	return configMap
}

func NewKafkaBatchReader(cfg *config.Config) *KafkaBatchReader {
	configMap := BuildKafkaConfigMap(cfg)

	consumer, err := kafka.NewConsumer(configMap)
	if err != nil {
		logrus.Fatalf("Failed to create Kafka consumer: %v", err)
	}

	var rebalanceCb kafka.RebalanceCb
	if cfg.EnableRebalanceOptimization {
		rebalanceCb = logRebalance
	}
	err = consumer.Subscribe(cfg.KafkaTopic, rebalanceCb)
	if err != nil {
		logrus.Fatalf("Failed to subscribe to topic %s: %v", cfg.KafkaTopic, err)
	}

	return &KafkaBatchReader{
		consumer:         consumer,
		topic:            cfg.KafkaTopic,
		batchSize:        cfg.BatchSize,
		maxBatchInterval: cfg.BatchMaxInterval,
		statsRecorder:    NewLakeConsumeStatsRecorder(),
	}
}

// logRebalance logs consumer group rebalance events (partition assignment and
// revocation). It intentionally does not call Assign/Unassign so that the
// confluent-kafka-go library performs the assignment itself, using the
// configured partition.assignment.strategy.
func logRebalance(c *kafka.Consumer, event kafka.Event) error {
	l := logrus.WithField("kafka_rebalance", true)
	switch e := event.(type) {
	case kafka.AssignedPartitions:
		l.Infof("Partitions assigned: %v", e.Partitions)
	case kafka.RevokedPartitions:
		l.Infof("Partitions revoked: %v", e.Partitions)
	}
	return nil
}

func (br *KafkaBatchReader) Close() error {
	return br.consumer.Close()
}

func (br *KafkaBatchReader) ReadBatch(ctx context.Context) (*message.MessagesBatch, error) {
	l := logrus.WithFields(logrus.Fields{"kafka_batch_reader": "ReadBatch"})

	startFetchBatchTime := time.Now()
	var (
		lastMessageOffset  int64 = -1
		firstMessageOffset int64 = -1
		allByteSize        int
		batch              = make([]message.MessageData, 0, br.batchSize)
		allMessages        = make(map[int32]*kafka.Message, br.batchSize)
		lastMessageTime    = time.Now()
		noMessageCount     = 0
		maxEmptyAttempts   = 3
		maxMessageInterval = 2 * time.Second
	)

	processMessage := func(m *kafka.Message) message.MessageData {
		data := strings.TrimSpace(strings.ReplaceAll(strings.ReplaceAll(string(m.Value), "\t", ""), "\n", ""))
		return message.MessageData{
			Data:       data,
			DataOffset: int64(m.TopicPartition.Offset),
			Partition:  int(m.TopicPartition.Partition),
			Key:        string(m.Key),
			CreateTime: m.Timestamp,
		}
	}

	batchDeadline := time.Now().Add(time.Duration(br.maxBatchInterval) * time.Second)
	pollTimeoutMs := 500

	for {
		select {
		case <-ctx.Done():
			if len(batch) > 0 {
				break
			}
			return nil, ctx.Err()
		default:
		}

		if time.Now().After(batchDeadline) && len(batch) > 0 {
			l.Infof("Batch timeout reached with %d messages", len(batch))
			break
		}

		if len(batch) > 0 && time.Since(lastMessageTime) > maxMessageInterval {
			l.Infof("Message interval exceeded %v, ending batch with %d messages",
				maxMessageInterval, len(batch))
			break
		}

		ev := br.consumer.Poll(pollTimeoutMs)
		if ev == nil {
			if len(batch) > 0 {
				noMessageCount++
				if noMessageCount >= maxEmptyAttempts {
					l.Infof("No new messages after %d attempts, returning current batch",
						maxEmptyAttempts)
					break
				}
			}
			continue
		}

		switch e := ev.(type) {
		case *kafka.Message:
			noMessageCount = 0
			lastMessageTime = time.Now()

			if firstMessageOffset == -1 {
				firstMessageOffset = int64(e.TopicPartition.Offset)
			}
			lastMessageOffset = int64(e.TopicPartition.Offset)

			batch = append(batch, processMessage(e))
			allByteSize += len(e.Value)
			allMessages[e.TopicPartition.Partition] = e

			if len(batch) >= br.batchSize {
				l.Infof("Batch full with %d messages in %v", len(batch),
					time.Since(startFetchBatchTime))
				break
			}
			continue

		case kafka.Error:
			if e.IsFatal() {
				if len(batch) > 0 {
					l.WithError(e).Warn("Fatal Kafka error, processing current batch")
					break
				}
				l.WithError(e).Error("Fatal Kafka error with empty batch")
				return nil, e
			}
			l.WithError(e).Warn("Transient Kafka error, continuing")
			continue

		default:
			continue
		}
		break
	}

	batchDuration := time.Since(startFetchBatchTime)
	batchSizeHist.Observe(float64(len(batch)))
	batchFillDuration.Observe(batchDuration.Seconds())
	l.WithFields(logrus.Fields{
		"batch_size":         len(batch),
		"bytes_size":         allByteSize,
		"first_offset":       firstMessageOffset,
		"last_offset":        lastMessageOffset,
		"duration_ms":        batchDuration.Milliseconds(),
		"messages_per_ms":    float64(len(batch)) / float64(batchDuration.Milliseconds()+1),
		"last_message_delay": time.Since(lastMessageTime).Milliseconds(),
	}).Info("Batch complete")

	if len(batch) == 0 {
		return &message.MessagesBatch{
			Messages:           batch,
			FirstMessageOffset: -1,
			LastMessageOffset:  -1,
			CommitFunc:         func(context.Context) error { return nil },
		}, nil
	}

	return &message.MessagesBatch{
		Messages:           batch,
		CommitFunc:         br.createCommitFunc(allMessages, l),
		FirstMessageOffset: firstMessageOffset,
		LastMessageOffset:  lastMessageOffset,
	}, nil
}

func (br *KafkaBatchReader) createCommitFunc(messages map[int32]*kafka.Message, l *logrus.Entry) func(context.Context) error {
	if len(messages) == 0 {
		return func(_ context.Context) error { return nil }
	}

	return func(ctx context.Context) error {
		for partition, msg := range messages {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
				if _, err := br.consumer.CommitMessage(msg); err != nil {
					if isMembershipError(err) {
						// The consumer's group membership is gone (rebalance or
						// expulsion): commits keep failing until it rejoins, and
						// committing now could rewind progress another member has
						// already made on this partition. Skip instead of retry;
						// the messages will be re-consumed from the last committed
						// offset after rejoin (at-least-once).
						l.WithError(err).WithFields(logrus.Fields{
							"partition": partition,
							"offset":    msg.TopicPartition.Offset,
						}).Warn("Group membership lost, skip commit; batch will be re-consumed after rejoin")
						return nil
					}
					l.WithError(err).WithFields(logrus.Fields{
						"partition": partition,
						"offset":    msg.TopicPartition.Offset,
					}).Error("Failed to commit message")
					return fmt.Errorf("commit message failed: %w", err)
				}
			}
		}
		return nil
	}
}

// isMembershipError reports whether err means the consumer's group membership
// is gone (rebalance in progress, illegal generation, or unknown member).
// Such errors cannot be fixed by retrying the commit: the consumer must
// rejoin the group first.
func isMembershipError(err error) bool {
	var kerr kafka.Error
	if errors.As(err, &kerr) {
		switch kerr.Code() {
		case kafka.ErrUnknownMemberID, kafka.ErrIllegalGeneration, kafka.ErrRebalanceInProgress:
			return true
		}
	}
	return false
}