package main

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

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

type kafkaMessagePersistPayload struct {
	Message        types.Message `json:"message"`
	AttachmentURLs []string      `json:"attachment_urls,omitempty"`
	ReadBySender   bool          `json:"read_by_sender,omitempty"`
	MessageKind    string        `json:"message_kind"`
	ShardMonth     string        `json:"shard_month,omitempty"`
}

var publishKafkaMessagePersist = publishKafkaEnvelope

func startKafkaMessagePersistConsumer() {
	startKafkaBatchConsumer("message_persist", func(ctx context.Context, messages []kafka.Message) error {
		return handleKafkaMessagePersistBatch(ctx, messages, "message_persist")
	})
	startKafkaBatchConsumer("group_message_persist", func(ctx context.Context, messages []kafka.Message) error {
		return handleKafkaMessagePersistBatch(ctx, messages, "group_message_persist")
	})
}

func saveMessageMaybeKafka(ctx context.Context, tenantID types.TenantID, msg *types.Message, attachmentURLs []string, readBySender bool) (*types.MessageSaveResult, error) {
	if !kafkaEnabled() {
		return store.Messages.Save(tenantID, msg, attachmentURLs, readBySender)
	}
	if msg == nil || msg.TenantID != tenantID || msg.Topic == "" || msg.SeqId <= 0 {
		return nil, types.ErrMalformed
	}
	msg.InitTimes()
	if msg.Uid().IsZero() {
		msg.SetUid(store.Store.GetUid())
	}
	isGroupMessage := types.GetTopicCat(msg.Topic) == types.TopicCatGrp
	clientMsgID, err := store.NormalizeClientMsgID(msg.Head)
	if err != nil {
		return nil, err
	}
	msg.ClientMsgID = clientMsgID
	fromUid := types.ParseUid(msg.From)
	if clientMsgID != "" {
		started := time.Now()
		claim, claimed, runtimeReady, claimErr := store.ClaimMessageUnique(tenantID, msg.Topic, fromUid, clientMsgID)
		immetrics.ObserveMessagePublishStage(messageKind(msg.Topic), "redis_idempotency", started, claimErr)
		if claimErr != nil {
			return nil, claimErr
		}
		if runtimeReady && !claimed {
			if claim != nil && claim.Seq > 0 {
				msg.SeqId = claim.Seq
				return &types.MessageSaveResult{SeqId: claim.Seq, Duplicate: true, Pending: claim.Status == "pending"}, nil
			}
			if claim != nil && claim.Status == "pending" {
				return &types.MessageSaveResult{Duplicate: true, Pending: true}, types.ErrDuplicate
			}
			return nil, types.ErrDuplicate
		}
	}
	candidateMonth := msg.CreatedAt.UTC().Format("200601")
	seqStarted := time.Now()
	seq, shardMonth, runtimeReady, seqErr := store.ReserveMessageSeq(tenantID, msg.Topic, msg.SeqId-1, candidateMonth)
	immetrics.ObserveMessagePublishStage(messageKind(msg.Topic), "redis_sequence", seqStarted, seqErr)
	if seqErr != nil {
		if clientMsgID != "" {
			_ = store.MarkMessageUniqueFailed(tenantID, msg.Topic, fromUid, clientMsgID)
		}
		return nil, seqErr
	}
	if runtimeReady {
		msg.SeqId = seq
		if isGroupMessage {
			msg.ShardMonth = shardMonth
		}
	}
	if msg.ShardMonth == "" && isGroupMessage {
		msg.ShardMonth = candidateMonth
	}
	kafkaStarted := time.Now()
	if err = enqueueKafkaMessagePersist(ctx, tenantID, msg, attachmentURLs, readBySender); err != nil {
		immetrics.ObserveMessagePublishStage(messageKind(msg.Topic), "kafka_produce", kafkaStarted, err)
		if clientMsgID != "" {
			if markErr := store.MarkMessageUniquePending(tenantID, msg.Topic, fromUid, clientMsgID, msg.SeqId); markErr != nil {
				logs.Warn.Printf("kafka message persist: failed to mark message unique key pending topic=%s seq=%d: %v", msg.Topic, msg.SeqId, markErr)
			}
		}
		logs.Warn.Printf("kafka message persist: publish state unknown topic=%s seq=%d client_mid=%s err=%v", msg.Topic, msg.SeqId, clientMsgID, err)
		return &types.MessageSaveResult{SeqId: msg.SeqId, Pending: true, MarkedReadBySender: readBySender && !fromUid.IsZero()}, err
	}
	immetrics.ObserveMessagePublishStage(messageKind(msg.Topic), "kafka_produce", kafkaStarted, nil)
	if clientMsgID != "" {
		if err = store.MarkMessageUniquePending(tenantID, msg.Topic, fromUid, clientMsgID, msg.SeqId); err != nil {
			logs.Warn.Printf("kafka message persist: failed to mark message unique key pending topic=%s seq=%d: %v", msg.Topic, msg.SeqId, err)
		}
	}
	if isGroupMessage {
		if cacheErr := store.CachePendingHotMessage(tenantID, msg); cacheErr != nil {
			logs.Warn.Printf("kafka message persist: failed to cache pending hot message after publish topic=%s seq=%d: %v", msg.Topic, msg.SeqId, cacheErr)
		}
	}
	return &types.MessageSaveResult{SeqId: msg.SeqId, Pending: true, MarkedReadBySender: readBySender && !fromUid.IsZero()}, nil
}

func messageKind(topic string) string {
	switch types.GetTopicCat(topic) {
	case types.TopicCatGrp:
		return "group"
	case types.TopicCatP2P:
		return "p2p"
	default:
		return "other"
	}
}

func enqueueKafkaMessagePersist(ctx context.Context, tenantID types.TenantID, msg *types.Message, attachmentURLs []string, readBySender bool) error {
	kind := messageKind(msg.Topic)
	payload, err := json.Marshal(kafkaMessagePersistPayload{
		Message:        *msg,
		AttachmentURLs: attachmentURLs,
		ReadBySender:   readBySender,
		MessageKind:    kind,
		ShardMonth:     msg.ShardMonth,
	})
	if err != nil {
		return err
	}
	eventID := fmt.Sprintf("msg.persist.%d.%s.%d", tenantID, msg.Topic, msg.SeqId)
	envelope := kafkaEnvelope{
		EventID:        eventID,
		TenantID:       tenantID,
		SourceService:  "im-server",
		TargetService:  "im-server",
		EventType:      "message.persist." + kind,
		BizType:        "message",
		BizID:          fmt.Sprintf("%s:%d", msg.Topic, msg.SeqId),
		IdempotencyKey: msg.ClientMsgID,
		Version:        msg.SeqId,
		OccurredAt:     msg.CreatedAt,
		Actor:          &kafkaActor{Type: "user", ID: types.ParseUid(msg.From).UserId()},
		Payload:        payload,
	}
	return publishKafkaMessagePersist(ctx, kafkaMessagePersistTopicName(msg.Topic), msg.Topic, envelope)
}

type kafkaMessagePersistItem struct {
	envelope kafkaEnvelope
	tenantID types.TenantID
	payload  kafkaMessagePersistPayload
}

func handleKafkaMessagePersistBatch(ctx context.Context, messages []kafka.Message, sourceTopicName string) error {
	groups := make(map[string][]*kafkaMessagePersistItem)
	for _, message := range messages {
		item, terminal, err := prepareKafkaMessagePersistItem(ctx, message, sourceTopicName)
		if err != nil {
			return err
		}
		if terminal {
			continue
		}
		key := fmt.Sprintf("%d", item.tenantID)
		if item.payload.MessageKind == "group" {
			key += ":" + item.payload.Message.ShardMonth
		}
		groups[key] = append(groups[key], item)
	}
	for _, items := range groups {
		kind := items[0].payload.MessageKind
		if kind == "" {
			kind = "p2p"
		}
		started := time.Now()
		batch := make([]store.PreparedMessageBatchItem, 0, len(items))
		for _, item := range items {
			batch = append(batch, store.PreparedMessageBatchItem{Message: &item.payload.Message, ReadBySender: item.payload.ReadBySender})
		}
		if err := store.PersistPreparedMessageBatch(items[0].tenantID, batch); err == nil {
			immetrics.KafkaMessagePersistTotal.WithLabelValues(kind, "batch_ok").Add(float64(len(items)))
			immetrics.MessagePersistDuration.WithLabelValues(kind, "batch_ok").Observe(time.Since(started).Seconds())
			for _, item := range items {
				if err = finishKafkaMessagePersist(ctx, item); err != nil {
					if err = persistKafkaMessageIndividually(ctx, item); err != nil {
						return err
					}
				}
			}
			continue
		} else {
			immetrics.KafkaMessagePersistTotal.WithLabelValues(kind, "batch_failed").Add(float64(len(items)))
			immetrics.MessagePersistDuration.WithLabelValues(kind, "batch_failed").Observe(time.Since(started).Seconds())
			logs.Warn.Printf("kafka message batch persist failed, degrading to singles: tenant=%d kind=%s shard=%s size=%d err=%v",
				items[0].tenantID, items[0].payload.MessageKind, items[0].payload.Message.ShardMonth, len(items), err)
		}
		for _, item := range items {
			if err := persistKafkaMessageIndividually(ctx, item); err != nil {
				return err
			}
		}
	}
	return nil
}

func prepareKafkaMessagePersistItem(ctx context.Context, message kafka.Message, sourceTopicName string) (*kafkaMessagePersistItem, bool, error) {
	var envelope kafkaEnvelope
	if err := json.Unmarshal(message.Value, &envelope); err != nil {
		logs.Err.Printf("kafka message malformed envelope skipped: partition=%d offset=%d err=%v", message.Partition, message.Offset, err)
		return nil, true, nil
	}
	tenantID := envelope.TenantID
	var err error
	if tenantID.IsZero() {
		tenantID, err = kafkaTenantID(envelope.TenantCode)
		if err != nil {
			return nil, false, err
		}
	}
	consume, claimed, err := beginKafkaConsumeForTopic(tenantID, envelope, sourceTopicName)
	if err != nil {
		return nil, false, err
	}
	if !claimed {
		return nil, true, nil
	}
	item := &kafkaMessagePersistItem{envelope: envelope, tenantID: tenantID}
	if err = json.Unmarshal(envelope.Payload, &item.payload); err != nil {
		return nil, true, markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
	}
	if item.payload.Message.TenantID == types.ZeroTenantID {
		item.payload.Message.TenantID = tenantID
	}
	if item.payload.Message.ShardMonth == "" {
		item.payload.Message.ShardMonth = item.payload.ShardMonth
	}
	if item.payload.Message.TenantID != tenantID {
		return nil, true, markKafkaDeadLetter(ctx, tenantID, envelope, consume, "TENANT_MISMATCH", "message tenant does not match envelope tenant")
	}
	if err = validateKafkaMessageKind(item.payload); err != nil {
		return nil, true, markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
	}
	return item, false, nil
}

func persistKafkaMessageIndividually(ctx context.Context, item *kafkaMessagePersistItem) error {
	kind := item.payload.MessageKind
	if kind == "" {
		kind = "p2p"
	}
	started := time.Now()
	operation := func() error {
		err := store.PersistPreparedMessageBatch(item.tenantID, []store.PreparedMessageBatchItem{{
			Message: &item.payload.Message, ReadBySender: item.payload.ReadBySender,
		}})
		if errors.Is(err, types.ErrUnsupported) {
			_, err = store.PersistPreparedMessage(item.tenantID, &item.payload.Message, nil, item.payload.ReadBySender)
		}
		if err != nil && !errors.Is(err, types.ErrDuplicate) {
			return err
		}
		if err = finishKafkaMessagePersist(ctx, item); err != nil {
			return err
		}
		immetrics.KafkaMessagePersistTotal.WithLabelValues(kind, "single_ok").Inc()
		immetrics.MessagePersistDuration.WithLabelValues(kind, "single_ok").Observe(time.Since(started).Seconds())
		return nil
	}
	markFailed := func(err error) (*types.KafkaConsumeRecord, error) {
		immetrics.KafkaMessagePersistTotal.WithLabelValues(kind, "retry").Inc()
		return store.HTTPInfrastructure.MarkKafkaConsumeFailed(item.tenantID, item.envelope.EventID, kafkaBus.config.ConsumerGroup,
			kafkaMessagePersistFailureReason(err), err.Error(), time.Now().UTC().Round(time.Millisecond))
	}
	markDead := func(failed *types.KafkaConsumeRecord, err error) error {
		immetrics.KafkaMessagePersistTotal.WithLabelValues(kind, "dead").Inc()
		immetrics.MessagePersistDuration.WithLabelValues(kind, "dead").Observe(time.Since(started).Seconds())
		reason := kafkaMessagePersistFailureReason(err)
		immetrics.KafkaDeadLettersTotal.WithLabelValues(kind, reason).Inc()
		return markKafkaDeadLetter(ctx, item.tenantID, item.envelope, failed, reason, err.Error())
	}
	return runKafkaMessageSingleRetries(ctx, kafkaBus.config.MessageRetries, operation, markFailed, markDead, waitKafkaRetry)
}

func runKafkaMessageSingleRetries(ctx context.Context, retryCount int, operation func() error,
	markFailed func(error) (*types.KafkaConsumeRecord, error), markDead func(*types.KafkaConsumeRecord, error) error,
	wait kafkaRetryWaiter) error {
	var lastErr error
	var failed *types.KafkaConsumeRecord
	for attempt := 1; attempt <= retryCount+1; attempt++ {
		if lastErr = operation(); lastErr == nil {
			return nil
		}
		var err error
		if failed, err = markFailed(lastErr); err != nil {
			return err
		}
		if attempt <= retryCount {
			if err = wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
				return err
			}
		}
	}
	return markDead(failed, lastErr)
}

func finishKafkaMessagePersist(ctx context.Context, item *kafkaMessagePersistItem) error {
	if len(item.payload.AttachmentURLs) > 0 {
		if err := enqueueKafkaFileAttachmentsBind(ctx, item.tenantID, &item.payload.Message, item.payload.AttachmentURLs); err != nil {
			return fmt.Errorf("file lifecycle enqueue failed: %w", err)
		}
	}
	return store.HTTPInfrastructure.MarkKafkaConsumeSucceeded(item.tenantID, item.envelope.EventID,
		kafkaBus.config.ConsumerGroup, time.Now().UTC().Round(time.Millisecond))
}

func kafkaMessagePersistFailureReason(err error) string {
	if err != nil && strings.Contains(err.Error(), "file lifecycle") {
		return "KAFKA_FILE_LIFECYCLE_FAILED"
	}
	return kafkaCommandErrorReason(err)
}

func kafkaMessagePersistTopicName(topic string) string {
	if types.GetTopicCat(topic) == types.TopicCatGrp {
		return "group_message_persist"
	}
	return "message_persist"
}

func validateKafkaMessageKind(payload kafkaMessagePersistPayload) error {
	if payload.Message.Topic == "" || payload.Message.SeqId <= 0 {
		return types.ErrMalformed
	}
	isGroup := types.GetTopicCat(payload.Message.Topic) == types.TopicCatGrp
	switch payload.MessageKind {
	case "group":
		if !isGroup {
			return types.ErrMalformed
		}
	case "p2p", "":
		if isGroup {
			return types.ErrMalformed
		}
	default:
		return types.ErrMalformed
	}
	return nil
}