package main

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

	"github.com/segmentio/kafka-go"
	"github.com/tinode/chat/server/logs"
	"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() {
	startKafkaConsumer("message_persist", handleKafkaMessagePersist)
	startKafkaConsumer("group_message_persist", handleKafkaGroupMessagePersist)
}

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 != "" {
		claim, claimed, runtimeReady, claimErr := store.ClaimMessageUnique(tenantID, msg.Topic, fromUid, clientMsgID)
		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")
	if seq, shardMonth, runtimeReady, seqErr := store.ReserveMessageSeq(tenantID, msg.Topic, msg.SeqId-1, candidateMonth); seqErr != nil {
		if clientMsgID != "" {
			_ = store.MarkMessageUniqueFailed(tenantID, msg.Topic, fromUid, clientMsgID)
		}
		return nil, seqErr
	} else if runtimeReady {
		msg.SeqId = seq
		if isGroupMessage {
			msg.ShardMonth = shardMonth
		}
	}
	if msg.ShardMonth == "" && isGroupMessage {
		msg.ShardMonth = candidateMonth
	}
	if err = enqueueKafkaMessagePersist(ctx, tenantID, msg, attachmentURLs, readBySender); err != nil {
		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
	}
	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 enqueueKafkaMessagePersist(ctx context.Context, tenantID types.TenantID, msg *types.Message, attachmentURLs []string, readBySender bool) error {
	kind := "p2p"
	if types.GetTopicCat(msg.Topic) == types.TopicCatGrp {
		kind = "group"
	}
	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)
}

func handleKafkaMessagePersist(ctx context.Context, envelope kafkaEnvelope, _ kafka.Message) error {
	return handleKafkaMessagePersistFromTopic(ctx, envelope, "message_persist")
}

func handleKafkaGroupMessagePersist(ctx context.Context, envelope kafkaEnvelope, _ kafka.Message) error {
	return handleKafkaMessagePersistFromTopic(ctx, envelope, "group_message_persist")
}

func handleKafkaMessagePersistFromTopic(ctx context.Context, envelope kafkaEnvelope, sourceTopicName string) error {
	tenantID := envelope.TenantID
	var err error
	if tenantID.IsZero() {
		tenantID, err = kafkaTenantID(envelope.TenantCode)
		if err != nil {
			return err
		}
	}
	consume, claimed, err := beginKafkaConsumeForTopic(tenantID, envelope, sourceTopicName)
	if err != nil || !claimed {
		return err
	}
	var payload kafkaMessagePersistPayload
	if err = json.Unmarshal(envelope.Payload, &payload); err != nil {
		return markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
	}
	if payload.Message.TenantID == types.ZeroTenantID {
		payload.Message.TenantID = tenantID
	}
	if payload.Message.ShardMonth == "" {
		payload.Message.ShardMonth = payload.ShardMonth
	}
	if payload.Message.TenantID != tenantID {
		return markKafkaDeadLetter(ctx, tenantID, envelope, consume, "TENANT_MISMATCH", "message tenant does not match envelope tenant")
	}
	if err = validateKafkaMessageKind(payload); err != nil {
		return markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
	}
	if _, err = store.PersistPreparedMessage(tenantID, &payload.Message, nil, payload.ReadBySender); err != nil && !errors.Is(err, types.ErrDuplicate) {
		failed, markErr := store.HTTPInfrastructure.MarkKafkaConsumeFailed(tenantID, envelope.EventID, kafkaBus.config.ConsumerGroup,
			kafkaCommandErrorReason(err), err.Error(), time.Now().UTC().Round(time.Millisecond))
		if markErr != nil {
			return markErr
		}
		if failed != nil && failed.RetryCount >= kafkaBus.config.MaxRetryCount {
			return markKafkaDeadLetter(ctx, tenantID, envelope, failed, kafkaCommandErrorReason(err), err.Error())
		}
		sleepKafkaRetryBackoff(failed.RetryCount)
		return err
	}
	if len(payload.AttachmentURLs) > 0 {
		if err = enqueueKafkaFileAttachmentsBind(ctx, tenantID, &payload.Message, payload.AttachmentURLs); err != nil {
			failed, markErr := store.HTTPInfrastructure.MarkKafkaConsumeFailed(tenantID, envelope.EventID, kafkaBus.config.ConsumerGroup,
				"KAFKA_FILE_LIFECYCLE_FAILED", err.Error(), time.Now().UTC().Round(time.Millisecond))
			if markErr != nil {
				return markErr
			}
			if failed != nil && failed.RetryCount >= kafkaBus.config.MaxRetryCount {
				return markKafkaDeadLetter(ctx, tenantID, envelope, failed, "KAFKA_FILE_LIFECYCLE_FAILED", err.Error())
			}
			sleepKafkaRetryBackoff(failed.RetryCount)
			return err
		}
	}
	return store.HTTPInfrastructure.MarkKafkaConsumeSucceeded(tenantID, envelope.EventID, kafkaBus.config.ConsumerGroup, time.Now().UTC().Round(time.Millisecond))
}

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
}