package agent

import (
	"context"
	"errors"
	"log/slog"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/api/dto"
	"github.com/tidbcloud/AgentTeam/internal/core/domain"
	"github.com/tidbcloud/AgentTeam/internal/runtimecontrol"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/store/message"
	"github.com/tidbcloud/AgentTeam/internal/store/rewind"
)

// MessageRecoveryGC deletes daemon recovery points after an agent message send
// has pruned them from storage.
type MessageRecoveryGC interface {
	GCRecoveryPoints(machineID string, agentID string, points []rewind.MessageRecoveryPoint) error
}

func (h *MessageHandler) sendMessage(ctx context.Context, authenticatedAgent agent.AuthenticatedAgent, params message.SendMessageParams) (string, error) {
	delivery, err := h.store.SendAgentMessage(ctx, authenticatedAgent, params)
	if err != nil {
		return "", err
	}
	if err := h.afterSentMessage(ctx, authenticatedAgent, delivery); err != nil {
		return "", err
	}
	return delivery.MessageID, nil
}

func (h *MessageHandler) sendMessageWithFreshness(ctx context.Context, authenticatedAgent agent.AuthenticatedAgent, params message.SendMessageParams) (message.SendMessageResult, error) {
	result, err := h.store.SendAgentMessageWithFreshness(ctx, authenticatedAgent, params)
	return h.completeFreshnessSend(ctx, authenticatedAgent, result, err)
}

func (h *MessageHandler) sendWorkflowReplyWithFreshness(
	ctx context.Context,
	authenticatedAgent agent.AuthenticatedAgent,
	planID string,
	request dto.AgentSendMessageRequest,
) (message.SendMessageResult, error) {
	result, err := h.workflowReplies.SendAgentWorkflowReply(
		ctx,
		authenticatedAgent,
		planID,
		request.Content,
		strings.TrimSpace(request.ClientMessageID),
		freshnessFromRequest(request),
	)
	return h.completeFreshnessSend(ctx, authenticatedAgent, result, err)
}

func (h *MessageHandler) completeFreshnessSend(
	ctx context.Context,
	authenticatedAgent agent.AuthenticatedAgent,
	result message.SendMessageResult,
	err error,
) (message.SendMessageResult, error) {
	if err != nil {
		return message.SendMessageResult{}, err
	}
	if result.State == "" {
		result.State = message.SendStateSent
	}
	if result.State != message.SendStateSent {
		return result, nil
	}
	if err := h.afterSentMessage(ctx, authenticatedAgent, result.Delivery); err != nil {
		return message.SendMessageResult{}, err
	}
	return result, nil
}

func (h *MessageHandler) afterSentMessage(ctx context.Context, authenticatedAgent agent.AuthenticatedAgent, delivery message.MessageDelivery) error {
	var prunedRecoveryPoints []rewind.MessageRecoveryPoint
	if h.rewindStore != nil {
		if deleted, err := h.rewindStore.DeleteAgentRecoveryPointsOutsideThread(ctx, authenticatedAgent.WorkspaceID, authenticatedAgent.AgentID, delivery.ThreadID); err != nil {
			h.warn("prune recovery points after agent send", slog.String("agent_id", authenticatedAgent.AgentID), slog.Any("error", err))
		} else {
			prunedRecoveryPoints = deleted
		}
	}
	h.fanoutCommittedDelivery(ctx, authenticatedAgent.WorkspaceID, delivery)
	if h.recoveryGC != nil && len(prunedRecoveryPoints) > 0 && strings.TrimSpace(delivery.AuthorMachineID) != "" {
		if err := h.recoveryGC.GCRecoveryPoints(delivery.AuthorMachineID, authenticatedAgent.AgentID, prunedRecoveryPoints); err != nil && !errors.Is(err, runtimecontrol.ErrMachineOffline) {
			h.warn("gc pruned recovery points after agent send", slog.String("agent_id", authenticatedAgent.AgentID), slog.Any("error", err))
		}
	}
	return nil
}

func writeMessageParamsFromRequest(request dto.AgentSendMessageRequest) (message.SendMessageParams, error) {
	target := strings.TrimSpace(request.Target)
	if target == "" {
		return message.SendMessageParams{}, badRequest("target is required")
	}
	reference, err := domain.ParseMessageTargetReference(target)
	if err != nil {
		return message.SendMessageParams{}, err
	}
	return writeMessageParamsFromReference(request, reference)
}

func writeMessageParamsFromReference(
	request dto.AgentSendMessageRequest,
	reference domain.MessageTargetReference,
) (message.SendMessageParams, error) {
	baseParams := message.SendMessageParams{
		Content:          request.Content,
		AttachmentIDs:    request.AttachmentIDs,
		ClientMessageID:  strings.TrimSpace(request.ClientMessageID),
		Freshness:        freshnessFromRequest(request),
		TeamInvokeTeamID: strings.TrimSpace(request.TeamInvokeTeamID),
	}
	switch reference.Kind {
	case domain.MessageTargetKindTask:
		baseParams.TaskScopeID = reference.TaskScopeID
	case domain.MessageTargetKindTeam:
		baseParams.TargetRef = reference.String()
	case domain.MessageTargetKindDM:
		// Plain DMs preserve the Agent message_send policy of ensuring the
		// conversation on first send. Reply-thread DMs use TargetRef below.
		baseParams.DMTo = reference.Thread.PeerMemberID
	case domain.MessageTargetKindChannel, domain.MessageTargetKindReplyThread:
		baseParams.TargetRef = reference.String()
	default:
		return message.SendMessageParams{}, badRequest("unsupported message target")
	}
	if baseParams.DMTo != "" {
		return baseParams, nil
	}
	return baseParams, nil
}

func freshnessFromRequest(request dto.AgentSendMessageRequest) message.SendMessageFreshness {
	if request.Freshness == nil {
		return message.SendMessageFreshness{}
	}
	return message.SendMessageFreshness{
		TargetKey:        strings.TrimSpace(request.Freshness.TargetKey),
		SeenUpToSeq:      request.Freshness.SeenUpToSeq,
		SourceMessageIDs: append([]string(nil), request.Freshness.SourceMessageIDs...),
	}
}