package main

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/segmentio/kafka-go"
	"github.com/segmentio/kafka-go/sasl"
	"github.com/segmentio/kafka-go/sasl/scram"
	"github.com/tinode/chat/server/logs"
	immetrics "github.com/tinode/chat/server/metrics"
	"github.com/tinode/chat/server/store/types"
)

const (
	kafkaTopicAdminCommand        = "im.admin.command"
	kafkaTopicAdminCommandResult  = "im.admin.command.result"
	kafkaTopicConfigEvent         = "im.config.event"
	kafkaTopicBusinessEvent       = "im.business.event"
	kafkaTopicMessagePersist      = "im.message.persist"
	kafkaTopicGroupMessagePersist = "im.group.message.persist"
	kafkaTopicMessagePush         = "im.message.push"
	kafkaTopicFileSecurity        = "im.file.security"
	kafkaTopicWorkbenchEvent      = "im.workbench.event"
	kafkaTopicCacheInvalidate     = "im.cache.invalidate"
	kafkaTopicAuditEvent          = "im.audit.event"
	kafkaTopicDeadLetter          = "im.dead.letter"
)

type kafkaConfig struct {
	Enabled          bool              `json:"enabled"`
	Brokers          []string          `json:"brokers"`
	ClientID         string            `json:"client_id"`
	ConsumerGroup    string            `json:"consumer_group"`
	Topics           map[string]string `json:"topics"`
	DialTimeoutMS    int               `json:"dial_timeout_ms"`
	WriteTimeoutMS   int               `json:"write_timeout_ms"`
	ReadTimeoutMS    int               `json:"read_timeout_ms"`
	BatchTimeoutMS   int               `json:"batch_timeout_ms"`
	RequiredAcks     int               `json:"required_acks"`
	AllowAutoTopic   bool              `json:"allow_auto_topic_creation"`
	TLSEnabled       bool              `json:"tls_enabled"`
	SASLEnabled      bool              `json:"sasl_enabled"`
	SASLMechanism    string            `json:"sasl_mechanism"`
	SASLUsername     string            `json:"sasl_username"`
	SASLPassword     string            `json:"sasl_password"`
	MaxRetryCount    int               `json:"max_retry_count"`
	RetryBackoffMS   []int             `json:"retry_backoff_ms"`
	MessageBatchSize int               `json:"message_batch_size"`
	MessageBatchWait int               `json:"message_batch_wait_ms"`
	MessageRetries   int               `json:"message_single_retry_count"`
	DeadRetentionSec int               `json:"dead_letter_retention_seconds"`
}

type kafkaEnvelope struct {
	EventID        string          `json:"event_id"`
	TraceID        string          `json:"trace_id,omitempty"`
	TenantCode     string          `json:"tenant_code,omitempty"`
	TenantID       types.TenantID  `json:"tenant_id,omitempty"`
	SourceService  string          `json:"source_service,omitempty"`
	TargetService  string          `json:"target_service,omitempty"`
	EventType      string          `json:"event_type"`
	BizType        string          `json:"biz_type,omitempty"`
	BizID          string          `json:"biz_id,omitempty"`
	IdempotencyKey string          `json:"idempotency_key,omitempty"`
	Version        int64           `json:"version,omitempty"`
	OccurredAt     time.Time       `json:"occurred_at"`
	Actor          *kafkaActor     `json:"actor,omitempty"`
	Payload        json.RawMessage `json:"payload,omitempty"`
	ParentEventID  string          `json:"parent_event_id,omitempty"`
}

type kafkaActor struct {
	Type        string `json:"type,omitempty"`
	ID          string `json:"id,omitempty"`
	DisplayName string `json:"display_name,omitempty"`
}

type kafkaRuntime struct {
	config kafkaConfig
	topics map[string]string

	writer *kafka.Writer
	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup
}

var kafkaBus *kafkaRuntime

var kafkaConsumerGroupUnsafeChars = regexp.MustCompile(`[^A-Za-z0-9._-]+`)

func initKafkaRuntime(raw json.RawMessage) error {
	if len(raw) == 0 {
		kafkaBus = nil
		logs.Info.Println("Kafka disabled: no config")
		return nil
	}
	var cfg kafkaConfig
	if err := json.Unmarshal(raw, &cfg); err != nil {
		return err
	}
	if !cfg.Enabled {
		kafkaBus = nil
		logs.Info.Println("Kafka disabled")
		return nil
	}
	normalizeKafkaConfig(&cfg)
	if len(cfg.Brokers) == 0 {
		return errors.New("kafka: enabled but brokers are not configured")
	}
	saslMechanism, err := kafkaSASLMechanism(cfg)
	if err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(context.Background())
	rt := &kafkaRuntime{
		config: cfg,
		topics: kafkaTopics(cfg.Topics),
		ctx:    ctx,
		cancel: cancel,
	}
	transport := &kafka.Transport{
		ClientID: cfg.ClientID,
	}
	if cfg.TLSEnabled {
		transport.TLS = &tls.Config{MinVersion: tls.VersionTLS12}
	}
	transport.SASL = saslMechanism
	rt.writer = &kafka.Writer{
		Addr:                   kafka.TCP(cfg.Brokers...),
		Balancer:               &kafka.Hash{},
		Transport:              transport,
		RequiredAcks:           kafka.RequiredAcks(cfg.RequiredAcks),
		AllowAutoTopicCreation: cfg.AllowAutoTopic,
		BatchTimeout:           time.Duration(cfg.BatchTimeoutMS) * time.Millisecond,
		WriteTimeout:           time.Duration(cfg.WriteTimeoutMS) * time.Millisecond,
		ReadTimeout:            time.Duration(cfg.ReadTimeoutMS) * time.Millisecond,
		MaxAttempts:            cfg.MaxRetryCount,
	}

	if err := rt.ping(context.Background()); err != nil {
		cancel()
		_ = rt.writer.Close()
		return err
	}
	kafkaBus = rt
	logs.Info.Printf("Kafka enabled: brokers=%s client_id=%s group=%s", strings.Join(cfg.Brokers, ","), cfg.ClientID, cfg.ConsumerGroup)
	return nil
}

func kafkaConfigured(raw json.RawMessage) bool {
	var cfg kafkaConfig
	return len(raw) > 0 && json.Unmarshal(raw, &cfg) == nil && cfg.Enabled
}

func normalizeKafkaConfig(cfg *kafkaConfig) {
	if cfg.ClientID == "" {
		cfg.ClientID = "im-server"
	}
	if cfg.ConsumerGroup == "" {
		cfg.ConsumerGroup = "im-server"
	}
	if cfg.DialTimeoutMS <= 0 {
		cfg.DialTimeoutMS = 3000
	}
	if cfg.WriteTimeoutMS <= 0 {
		cfg.WriteTimeoutMS = 3000
	}
	if cfg.ReadTimeoutMS <= 0 {
		cfg.ReadTimeoutMS = 3000
	}
	if cfg.BatchTimeoutMS <= 0 {
		cfg.BatchTimeoutMS = 10
	}
	if cfg.RequiredAcks == 0 {
		cfg.RequiredAcks = int(kafka.RequireAll)
	}
	if cfg.MaxRetryCount <= 0 {
		cfg.MaxRetryCount = 3
	}
	if len(cfg.RetryBackoffMS) == 0 {
		cfg.RetryBackoffMS = []int{10000, 30000, 60000, 300000, 900000}
	}
	if cfg.MessageBatchSize <= 0 {
		cfg.MessageBatchSize = 100
	} else if cfg.MessageBatchSize > 500 {
		cfg.MessageBatchSize = 500
	}
	if cfg.MessageBatchWait <= 0 {
		cfg.MessageBatchWait = 10
	}
	if cfg.MessageRetries <= 0 {
		cfg.MessageRetries = 3
	}
	if cfg.DeadRetentionSec <= 0 {
		cfg.DeadRetentionSec = 90 * 24 * 3600
	}
	if cfg.SASLMechanism == "" {
		cfg.SASLMechanism = "SCRAM-SHA-512"
	}
}

func kafkaSASLMechanism(cfg kafkaConfig) (sasl.Mechanism, error) {
	if !cfg.SASLEnabled && cfg.SASLUsername == "" {
		return nil, nil
	}
	if cfg.SASLUsername == "" || cfg.SASLPassword == "" {
		return nil, errors.New("kafka: sasl enabled but username or password is empty")
	}
	switch strings.ToUpper(strings.ReplaceAll(cfg.SASLMechanism, "_", "-")) {
	case "", "SCRAM-SHA-512":
		return scram.Mechanism(scram.SHA512, cfg.SASLUsername, cfg.SASLPassword)
	case "SCRAM-SHA-256":
		return scram.Mechanism(scram.SHA256, cfg.SASLUsername, cfg.SASLPassword)
	default:
		return nil, errors.New("kafka: unsupported sasl mechanism " + cfg.SASLMechanism)
	}
}

func kafkaTopics(overrides map[string]string) map[string]string {
	topics := map[string]string{
		"admin_command":         kafkaTopicAdminCommand,
		"admin_command_result":  kafkaTopicAdminCommandResult,
		"config_event":          kafkaTopicConfigEvent,
		"business_event":        kafkaTopicBusinessEvent,
		"message_persist":       kafkaTopicMessagePersist,
		"group_message_persist": kafkaTopicGroupMessagePersist,
		"message_push":          kafkaTopicMessagePush,
		"file_security":         kafkaTopicFileSecurity,
		"workbench_external":    kafkaTopicWorkbenchEvent,
		"cache_invalidate":      kafkaTopicCacheInvalidate,
		"audit_event":           kafkaTopicAuditEvent,
		"dead_letter":           kafkaTopicDeadLetter,
	}
	for key, value := range overrides {
		if value != "" {
			topics[key] = value
		}
	}
	return topics
}

func (rt *kafkaRuntime) ping(ctx context.Context) error {
	timeout := time.Duration(rt.config.DialTimeoutMS) * time.Millisecond
	dialer := rt.dialer()
	pingCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	conn, err := dialer.DialContext(pingCtx, "tcp", rt.config.Brokers[0])
	if err != nil {
		return err
	}
	return conn.Close()
}

func kafkaEnabled() bool {
	return kafkaBus != nil
}

func kafkaTopic(name string) string {
	if kafkaBus == nil {
		return ""
	}
	return kafkaBus.topics[name]
}

func publishKafkaEnvelope(ctx context.Context, topicName string, key string, envelope kafkaEnvelope) error {
	if kafkaBus == nil {
		return nil
	}
	if envelope.OccurredAt.IsZero() {
		envelope.OccurredAt = time.Now().UTC()
	}
	topic := kafkaBus.topics[topicName]
	if topic == "" {
		topic = topicName
	}
	value, err := json.Marshal(envelope)
	if err != nil {
		return err
	}
	return kafkaBus.writer.WriteMessages(ctx, kafka.Message{
		Topic: topic,
		Key:   []byte(key),
		Value: value,
		Time:  envelope.OccurredAt,
	})
}

type kafkaMessageHandler func(context.Context, kafkaEnvelope, kafka.Message) error

type kafkaMessageCommitter func(context.Context, kafka.Message) error
type kafkaRetryWaiter func(context.Context, time.Duration) error
type kafkaBatchHandler func(context.Context, []kafka.Message) error
type kafkaBatchCommitter func(context.Context, ...kafka.Message) error
type kafkaMessageFetcher func(context.Context) (kafka.Message, error)

func startKafkaConsumer(topicName string, handler kafkaMessageHandler) {
	startKafkaConsumerWithGroup(topicName, "", handler)
}

func startKafkaBroadcastConsumer(topicName string, handler kafkaMessageHandler) string {
	if kafkaBus == nil || handler == nil {
		return ""
	}
	groupID := kafkaBroadcastConsumerGroup(topicName)
	startKafkaConsumerWithGroup(topicName, groupID, handler)
	return groupID
}

func startKafkaConsumerWithGroup(topicName, groupID string, handler kafkaMessageHandler) {
	if kafkaBus == nil || handler == nil {
		return
	}
	topic := kafkaBus.topics[topicName]
	if topic == "" {
		topic = topicName
	}
	groupOverride := groupID != ""
	if groupID == "" {
		groupID = kafkaBus.config.ConsumerGroup
	}
	if topicName != "" && !groupOverride {
		groupID += "." + topicName
	}
	reader := kafka.NewReader(kafka.ReaderConfig{
		Brokers:     kafkaBus.config.Brokers,
		GroupID:     groupID,
		Topic:       topic,
		StartOffset: kafka.FirstOffset,
		MinBytes:    1,
		MaxBytes:    10e6,
		MaxWait:     time.Second,
		Dialer:      kafkaBus.dialer(),
	})
	readerMetrics := startKafkaReaderMetrics(kafkaBus.ctx, reader, topic, groupID)
	kafkaBus.wg.Add(1)
	go func() {
		defer kafkaBus.wg.Done()
		defer reader.Close()
		for {
			msg, err := reader.FetchMessage(kafkaBus.ctx)
			if err != nil {
				if errors.Is(err, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
					return
				}
				logs.Err.Println("Kafka fetch failed:", err)
				continue
			}
			readerMetrics.observe(msg.Time)
			commit := func(ctx context.Context, message kafka.Message) error {
				return reader.CommitMessages(ctx, message)
			}
			if err = consumeKafkaMessage(kafkaBus.ctx, msg, handler, commit, waitKafkaRetry); err != nil {
				if errors.Is(err, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
					return
				}
				logs.Err.Printf("Kafka message stopped without commit: topic=%s partition=%d offset=%d err=%v", msg.Topic, msg.Partition, msg.Offset, err)
			}
			readerMetrics.clear()
		}
	}()
	logs.Info.Printf("Kafka consumer started: topic=%s group=%s", topic, groupID)
}

func startKafkaBatchConsumer(topicName string, handler kafkaBatchHandler) {
	if kafkaBus == nil || handler == nil {
		return
	}
	topic := kafkaBus.topics[topicName]
	if topic == "" {
		topic = topicName
	}
	groupID := kafkaBus.config.ConsumerGroup + "." + topicName
	reader := kafka.NewReader(kafka.ReaderConfig{
		Brokers: kafkaBus.config.Brokers, GroupID: groupID, Topic: topic, StartOffset: kafka.FirstOffset,
		MinBytes: 1, MaxBytes: 10e6, MaxWait: time.Second, Dialer: kafkaBus.dialer(),
	})
	readerMetrics := startKafkaReaderMetrics(kafkaBus.ctx, reader, topic, groupID)
	kafkaBus.wg.Add(1)
	go func() {
		defer kafkaBus.wg.Done()
		defer reader.Close()
		for {
			messages, fetchErr := fetchKafkaMessageBatch(kafkaBus.ctx, reader, kafkaBus.config.MessageBatchSize,
				time.Duration(kafkaBus.config.MessageBatchWait)*time.Millisecond)
			if fetchErr != nil && len(messages) == 0 {
				if errors.Is(fetchErr, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
					return
				}
				logs.Err.Println("Kafka batch fetch failed:", fetchErr)
				continue
			}
			if fetchErr != nil {
				logs.Warn.Printf("Kafka batch fetch interrupted; processing retained messages: topic=%s size=%d err=%v",
					topic, len(messages), fetchErr)
			}
			if len(messages) > 0 {
				readerMetrics.observe(messages[0].Time)
			}
			if err := consumeKafkaMessageBatch(kafkaBus.ctx, messages, handler, reader.CommitMessages, waitKafkaRetry); err != nil {
				if errors.Is(err, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
					return
				}
				logs.Err.Printf("Kafka batch stopped without commit: topic=%s size=%d err=%v", topic, len(messages), err)
			}
			readerMetrics.clear()
		}
	}()
	logs.Info.Printf("Kafka batch consumer started: topic=%s group=%s size=%d wait_ms=%d", topic, groupID, kafkaBus.config.MessageBatchSize, kafkaBus.config.MessageBatchWait)
}

type kafkaReaderMetrics struct {
	reader     *kafka.Reader
	topic      string
	group      string
	oldestNano atomic.Int64
}

func startKafkaReaderMetrics(ctx context.Context, reader *kafka.Reader, topic, group string) *kafkaReaderMetrics {
	metrics := &kafkaReaderMetrics{reader: reader, topic: topic, group: group}
	kafkaBus.wg.Add(1)
	go func() {
		defer kafkaBus.wg.Done()
		ticker := time.NewTicker(5 * time.Second)
		defer ticker.Stop()
		metrics.update()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				metrics.update()
			}
		}
	}()
	return metrics
}

func (m *kafkaReaderMetrics) observe(messageTime time.Time) {
	if !messageTime.IsZero() {
		m.oldestNano.Store(messageTime.UnixNano())
	}
	m.update()
}

func (m *kafkaReaderMetrics) clear() {
	m.oldestNano.Store(0)
	m.update()
}

func (m *kafkaReaderMetrics) update() {
	stats := m.reader.Stats()
	if stats.Lag >= 0 {
		immetrics.KafkaConsumerLag.WithLabelValues(clusterPodName(), m.topic, m.group).Set(float64(stats.Lag))
	}
	age := float64(0)
	if oldest := m.oldestNano.Load(); oldest > 0 {
		age = time.Since(time.Unix(0, oldest)).Seconds()
		if age < 0 {
			age = 0
		}
	}
	immetrics.KafkaBacklogOldestAge.WithLabelValues(clusterPodName(), m.topic, m.group).Set(age)
}

func fetchKafkaMessageBatch(ctx context.Context, reader *kafka.Reader, maxSize int, maxWait time.Duration) ([]kafka.Message, error) {
	return fetchKafkaMessageBatchWith(ctx, maxSize, maxWait, reader.FetchMessage)
}

func fetchKafkaMessageBatchWith(ctx context.Context, maxSize int, maxWait time.Duration, fetch kafkaMessageFetcher) ([]kafka.Message, error) {
	first, err := fetch(ctx)
	if err != nil {
		return nil, err
	}
	messages := []kafka.Message{first}
	deadline := time.Now().Add(maxWait)
	for len(messages) < maxSize {
		remaining := time.Until(deadline)
		if remaining <= 0 {
			break
		}
		fetchCtx, cancel := context.WithTimeout(ctx, remaining)
		message, fetchErr := fetch(fetchCtx)
		cancel()
		if errors.Is(fetchErr, context.DeadlineExceeded) && ctx.Err() == nil {
			break
		}
		if fetchErr != nil {
			// FetchMessage advances the reader's internal position before offsets are
			// committed. Preserve a partial batch so a later higher-offset commit cannot
			// skip messages already fetched into this batch.
			return messages, fetchErr
		}
		messages = append(messages, message)
	}
	return messages, nil
}

func consumeKafkaMessageBatch(ctx context.Context, messages []kafka.Message, handler kafkaBatchHandler, commit kafkaBatchCommitter, wait kafkaRetryWaiter) error {
	for attempt := 1; ; attempt++ {
		if err := handler(ctx, messages); err == nil {
			return commitKafkaMessageBatch(ctx, messages, commit, wait)
		}
		if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
			return err
		}
	}
}

func consumeKafkaMessage(ctx context.Context, msg kafka.Message, handler kafkaMessageHandler, commit kafkaMessageCommitter, wait kafkaRetryWaiter) error {
	var envelope kafkaEnvelope
	if err := json.Unmarshal(msg.Value, &envelope); err != nil {
		return commitKafkaMessage(ctx, msg, commit, wait)
	}
	for attempt := 1; ; attempt++ {
		if err := handler(ctx, envelope, msg); err == nil {
			return commitKafkaMessage(ctx, msg, commit, wait)
		}
		if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
			return err
		}
	}
}

func commitKafkaMessage(ctx context.Context, msg kafka.Message, commit kafkaMessageCommitter, wait kafkaRetryWaiter) error {
	for attempt := 1; ; attempt++ {
		if err := commit(ctx, msg); err == nil {
			return nil
		}
		if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
			return err
		}
	}
}

func commitKafkaMessageBatch(ctx context.Context, messages []kafka.Message, commit kafkaBatchCommitter, wait kafkaRetryWaiter) error {
	for attempt := 1; ; attempt++ {
		if err := commit(ctx, messages...); err == nil {
			return nil
		}
		if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
			return err
		}
	}
}

func kafkaRetryBackoff(attempt int) time.Duration {
	if kafkaBus == nil || len(kafkaBus.config.RetryBackoffMS) == 0 {
		return time.Second
	}
	if attempt < 1 {
		attempt = 1
	}
	index := attempt - 1
	if index >= len(kafkaBus.config.RetryBackoffMS) {
		index = len(kafkaBus.config.RetryBackoffMS) - 1
	}
	return time.Duration(kafkaBus.config.RetryBackoffMS[index]) * time.Millisecond
}

func waitKafkaRetry(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func kafkaBroadcastConsumerGroup(topicName string) string {
	if kafkaBus == nil {
		return ""
	}
	base := strings.TrimSpace(kafkaBus.config.ConsumerGroup)
	if base == "" {
		base = "im-server"
	}
	nodeID := strings.TrimSpace(currentNodeID())
	if nodeID == "" {
		nodeID = "local"
	}
	nodeID = kafkaConsumerGroupUnsafeChars.ReplaceAllString(nodeID, "-")
	return base + "." + topicName + "." + nodeID
}

func (rt *kafkaRuntime) dialer() *kafka.Dialer {
	dialer := &kafka.Dialer{
		ClientID: rt.config.ClientID,
		Timeout:  time.Duration(rt.config.DialTimeoutMS) * time.Millisecond,
	}
	if rt.config.TLSEnabled {
		dialer.TLS = &tls.Config{MinVersion: tls.VersionTLS12}
	}
	if mechanism, err := kafkaSASLMechanism(rt.config); err != nil {
		logs.Err.Println("Kafka SASL config failed:", err)
	} else {
		dialer.SASLMechanism = mechanism
	}
	return dialer
}

func stopKafkaRuntime() {
	if kafkaBus == nil {
		return
	}
	kafkaBus.cancel()
	kafkaBus.wg.Wait()
	if kafkaBus.writer != nil {
		_ = kafkaBus.writer.Close()
	}
	logs.Info.Println("Stopped Kafka runtime")
	kafkaBus = nil
}