package mcp

import (
	"context"
	"fmt"
	"strings"

	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

type messageFreshnessTracker interface {
	MessageFreshnessSnapshot(agentID string, targetKey string) (types.MessageFreshnessFrontier, bool)
	RememberMessageFreshnessObserved(agentID string, messages []types.InboxMessage) bool
	RememberMessageFreshnessHeld(agentID string, targetKey string, seenUpToSeq int64, latestSeq int64) bool
}

type agentSendFreshnessInput struct {
	TargetKey        string   `json:"targetKey"`
	SeenUpToSeq      int64    `json:"seenUpToSeq"`
	SourceMessageIDs []string `json:"sourceMessageIds,omitempty"`
}

type agentSendMessageResponse struct {
	State                string `json:"state"`
	MessageID            string `json:"messageId"`
	TargetKey            string `json:"targetKey"`
	ClientMessageID      string `json:"clientMessageId"`
	SeenUpToSeq          int64  `json:"seenUpToSeq"`
	AcceptedSeq          int64  `json:"acceptedSeq"`
	LatestSeq            int64  `json:"latestSeq"`
	NeedsRefreshUntilSeq int64  `json:"needsRefreshUntilSeq"`
	Error                string `json:"error"`
}

type sendFreshnessContext struct {
	Enabled              bool
	TargetKey            string
	SeenUpToSeq          int64
	LatestKnownSeq       int64
	NeedsRefreshUntilSeq int64
	SourceMessageIDs     []string
}

func (f sendFreshnessContext) localHeldResponse(clientMessageID string) (agentSendMessageResponse, bool) {
	if !f.Enabled {
		return agentSendMessageResponse{}, false
	}
	latestSeq := f.LatestKnownSeq
	if f.NeedsRefreshUntilSeq > latestSeq {
		latestSeq = f.NeedsRefreshUntilSeq
	}
	if latestSeq <= f.SeenUpToSeq {
		return agentSendMessageResponse{}, false
	}
	return agentSendMessageResponse{
		State:                agentSendStateHeld,
		TargetKey:            f.TargetKey,
		ClientMessageID:      strings.TrimSpace(clientMessageID),
		SeenUpToSeq:          f.SeenUpToSeq,
		LatestSeq:            latestSeq,
		NeedsRefreshUntilSeq: latestSeq,
	}, true
}

type acceptedSendChunk struct {
	ChunkIndex  int
	MessageID   string
	AcceptedSeq int64
	State       string
}

const (
	agentSendStateSent        = "sent"
	agentSendStateAlreadySent = "already_sent"
	agentSendStateHeld        = "held"
	// The Agent send API accepts one message body up to 32 KiB. Team
	// invocations are capability-bearing handoffs and must remain one persisted
	// message, rather than repeating the capability across text chunks.
	teamInvokeMessageMaxContentBytes = 32 * 1024
)

func (h *mcpHandler) sendMessageViaAPI(ctx context.Context, agentID string, api *apiclient.Client, target, content string, filePaths []string, attachmentIDs []string, teamInvokeTeamID string) *types.MCPToolResult {
	normalizedTarget, err := normalizeMessageTarget(target)
	if err != nil {
		return toolError(err.Error())
	}
	target = normalizedTarget
	content = normalizeEscapedAgentSendMessageContent(content)
	teamInvokeTeamID = strings.TrimSpace(teamInvokeTeamID)
	if teamInvokeTeamID != "" && len(content) > teamInvokeMessageMaxContentBytes {
		return toolError(fmt.Sprintf("Team invocation message is too large (max %d bytes)", teamInvokeMessageMaxContentBytes))
	}

	// Upload files via API first, collecting attachment IDs.
	uploadedAttachmentIDs := make([]string, 0, len(attachmentIDs)+len(filePaths))
	uploadedAttachmentIDs = append(uploadedAttachmentIDs, attachmentIDs...)
	for _, fp := range filePaths {
		fp = strings.TrimSpace(fp)
		if fp == "" {
			continue
		}
		id, err := h.uploadFileViaAPI(ctx, api, target, fp)
		if err != nil {
			return toolError(err.Error())
		}
		uploadedAttachmentIDs = append(uploadedAttachmentIDs, id)
	}

	chunks := []string{content}
	if teamInvokeTeamID == "" && len(uploadedAttachmentIDs) == 0 {
		chunks = apiclient.SplitMessageText(content, textMaxBytes)
	}
	if len(chunks) == 0 {
		chunks = []string{""}
	}

	freshness, err := h.prepareFreshnessForSend(ctx, agentID, api, target)
	if err != nil {
		return toolError(err.Error())
	}

	serverID := ""
	if h != nil && h.c != nil {
		serverID = h.c.ID()
	}
	lease, err := h.service.acquirePartialSend(ctx, serverID, agentID)
	if err != nil {
		return toolError(err.Error())
	}
	defer lease.Release()

	path := api.AgentRuntimePath("/send")
	logicalClientID := newAgentToolSentMessageID()
	partialSignature := ""
	if freshness.Enabled {
		partialSignature = partialSendSignature(target, chunks, uploadedAttachmentIDs, teamInvokeTeamID)
	}
	pendingLastAcceptedSeq := int64(0)
	// sendSeenUpToSeq is the freshness cursor used to sequence subsequent
	// chunks. It is deliberately separate from the server-authoritative
	// settlement cursor persisted below: AcceptedSeq may advance this request
	// cursor, but it can never become reply-settlement coverage by itself.
	sendSeenUpToSeq := freshness.SeenUpToSeq
	nextChunkIndex := 0
	pendingAcceptedChunks := []acceptedSendChunk(nil)
	pendingSettlementTargetKey := ""
	pendingSettlementSeenUpToSeq := int64(0)
	pendingSettlementMessageID := ""
	recoveredPending := false
	if partialSignature != "" {
		if pending, ok := lease.Pending(); ok &&
			pending.Target == target &&
			pending.Signature == partialSignature &&
			strings.TrimSpace(pending.LogicalClientID) != "" &&
			pending.NextChunkIndex >= 0 &&
			pending.NextChunkIndex <= len(chunks) {
			logicalClientID = pending.LogicalClientID
			pendingLastAcceptedSeq = pending.LastAcceptedSeq
			if pendingLastAcceptedSeq > sendSeenUpToSeq {
				sendSeenUpToSeq = pendingLastAcceptedSeq
			}
			nextChunkIndex = pending.NextChunkIndex
			pendingAcceptedChunks = append([]acceptedSendChunk(nil), pending.AcceptedChunks...)
			pendingSettlementTargetKey = strings.TrimSpace(pending.SettlementTargetKey)
			pendingSettlementSeenUpToSeq = pending.SettlementSeenUpToSeq
			pendingSettlementMessageID = strings.TrimSpace(pending.SettlementMessageID)
			recoveredPending = true
		} else {
			lease.ClearPending()
		}
	} else {
		lease.ClearPending()
	}
	// A fully accepted prefix has no new freshness claim to make. Even if a
	// stale local held hint survived a process restart, use the persisted
	// successful server coverage instead of blocking recovery on that hint.
	fullPrefixRecovered := recoveredPending && nextChunkIndex == len(chunks)
	if !fullPrefixRecovered {
		if result, ok := freshness.localHeldResponse(logicalClientID); ok {
			h.rememberSendHeld(agentID, result)
			refresh := h.heldRefresh(ctx, agentID, api, target, result)
			return toolTextWithRuntimeDeliveryReceipts(formatHeldSendResult(target, result, refresh), refresh.RuntimeDeliveryReceipts, refresh.RuntimeDeliveryCommitOwner)
		}
	}

	// Accepted chunks settle only after a genuine server response carries the
	// canonical target frontier. No local wake or token state is changed here.
	lastAcceptedSeq := pendingLastAcceptedSeq
	acceptedChunks := append([]acceptedSendChunk(nil), pendingAcceptedChunks...)
	firstMessageID := pendingSettlementMessageID
	lastResponse := agentSendMessageResponse{}
	if fullPrefixRecovered {
		// Every chunk was accepted by an earlier successful server response.
		// Reuse only the target/frontier/message identity captured from that
		// response; never turn the last accepted message sequence into a
		// settlement frontier. No chunk POST is necessary on this recovery.
		lastResponse = agentSendMessageResponse{
			State:       agentSendStateAlreadySent,
			MessageID:   firstMessageID,
			TargetKey:   pendingSettlementTargetKey,
			SeenUpToSeq: pendingSettlementSeenUpToSeq,
		}
	}
	for i := nextChunkIndex; i < len(chunks); i++ {
		chunk := chunks[i]
		clientMessageID := logicalClientID
		if len(chunks) > 1 {
			clientMessageID = fmt.Sprintf("%s:%d", logicalClientID, i)
		}
		body := map[string]any{
			"target":        target,
			"content":       chunk,
			"attachmentIds": uploadedAttachmentIDs,
		}
		if teamInvokeTeamID != "" {
			body["teamInvokeTeamId"] = teamInvokeTeamID
		}
		if len(chunks) > 1 {
			body["attachmentIds"] = []string{}
		}
		if freshness.Enabled {
			body["clientMessageId"] = clientMessageID
			body["freshness"] = agentSendFreshnessInput{
				TargetKey:        freshness.TargetKey,
				SeenUpToSeq:      sendSeenUpToSeq,
				SourceMessageIDs: append([]string(nil), freshness.SourceMessageIDs...),
			}
		}

		var result agentSendMessageResponse
		if err := api.Post(ctx, path, body, &result); err != nil {
			return toolError(err.Error())
		}
		if result.Error != "" {
			return toolError(result.Error)
		}
		result.State = normalizeAgentSendState(result.State, result.MessageID)
		if strings.TrimSpace(result.ClientMessageID) == "" {
			result.ClientMessageID = clientMessageID
		}
		lastResponse = result

		switch result.State {
		case agentSendStateSent, agentSendStateAlreadySent:
			if firstMessageID == "" {
				firstMessageID = strings.TrimSpace(result.MessageID)
			}
			acceptedChunks = append(acceptedChunks, acceptedSendChunk{
				ChunkIndex:  i,
				MessageID:   strings.TrimSpace(result.MessageID),
				AcceptedSeq: result.AcceptedSeq,
				State:       result.State,
			})
			if result.AcceptedSeq > lastAcceptedSeq {
				lastAcceptedSeq = result.AcceptedSeq
			}
			if result.AcceptedSeq > sendSeenUpToSeq {
				sendSeenUpToSeq = result.AcceptedSeq
			}
			settlementTargetKey := pendingSettlementTargetKey
			settlementSeenUpToSeq := pendingSettlementSeenUpToSeq
			if targetKey := strings.TrimSpace(result.TargetKey); targetKey != "" {
				// These two fields are the only coverage that may settle a
				// acknowledged wake. They come from this genuine successful server
				// response, never from AcceptedSeq or the local frontier.
				settlementTargetKey = targetKey
				settlementSeenUpToSeq = result.SeenUpToSeq
			}
			pendingSettlementTargetKey = settlementTargetKey
			pendingSettlementSeenUpToSeq = settlementSeenUpToSeq
			pendingSettlementMessageID = firstMessageID
			lease.StorePending(pendingPartialSendState{
				Target:                target,
				Signature:             partialSignature,
				LogicalClientID:       logicalClientID,
				LastAcceptedSeq:       lastAcceptedSeq,
				NextChunkIndex:        i + 1,
				AcceptedChunks:        acceptedChunks,
				SettlementTargetKey:   settlementTargetKey,
				SettlementSeenUpToSeq: settlementSeenUpToSeq,
				SettlementMessageID:   firstMessageID,
			})
			h.rememberSendAccepted(agentID, target, result)
		case agentSendStateHeld:
			h.rememberSendHeld(agentID, result)
			refresh := h.heldRefresh(ctx, agentID, api, target, result)
			if len(chunks) > 1 && (nextChunkIndex > 0 || len(acceptedChunks) > 0) {
				if h.c != nil {
					h.recordSent()
					h.recordAgentSent(agentID)
				}
				// A held response settles no visible-reply obligation, even when
				// an earlier chunk in this partial send was accepted. The
				// refresh/retry context is the only authoritative outcome for this attempt.
				return toolTextWithRuntimeDeliveryReceipts(formatPartialHeldSendResult(target, i, acceptedChunks, result, refresh), refresh.RuntimeDeliveryReceipts, refresh.RuntimeDeliveryCommitOwner)
			}
			return toolTextWithRuntimeDeliveryReceipts(formatHeldSendResult(target, result, refresh), refresh.RuntimeDeliveryReceipts, refresh.RuntimeDeliveryCommitOwner)
		default:
			return toolError("unexpected message_send response state: " + result.State)
		}
	}

	// If the final successful echo is missing its resolved identity, retain
	// the last genuine successful response's coverage captured in the pending
	// record. Never synthesize either field from chunk sequencing or local
	// freshness state.
	if strings.TrimSpace(lastResponse.TargetKey) == "" && pendingSettlementTargetKey != "" {
		lastResponse.TargetKey = pendingSettlementTargetKey
		lastResponse.SeenUpToSeq = pendingSettlementSeenUpToSeq
	}

	// Record sent stats for the API path (doSend path records in outbound.go).
	if h.c != nil {
		h.recordSent()
		h.recordAgentSent(agentID)
	}

	// Phase 2: settle every acknowledged wake covered by the server response's
	// canonical target frontier. This intentionally does not use AcceptedSeq,
	// local snapshots, or a local delivery token as coverage. Keep the pending
	// recovery record until this attempt has had a chance to settle so
	// a process exit between send and confirmation can recover without posting
	// any already accepted chunk again.
	if fullPrefixRecovered || len(acceptedChunks) > 0 {
		h.confirmAcceptedVisibleReply(agentID, target, lastResponse, firstMessageID)
	}
	lease.ClearPending()

	return toolText(formatAcceptedSendResult(target, chunks, acceptedChunks, lastResponse))
}

func (h *mcpHandler) prepareFreshnessForSend(ctx context.Context, agentID string, api *apiclient.Client, target string) (sendFreshnessContext, error) {
	tracker := h.freshnessTracker()
	if tracker == nil {
		return sendFreshnessContext{}, nil
	}
	targetKey, err := h.resolveFreshnessTargetKey(ctx, api, target)
	if err != nil {
		return sendFreshnessContext{}, err
	}
	frontier, _ := tracker.MessageFreshnessSnapshot(agentID, targetKey)
	return sendFreshnessContext{
		Enabled:              true,
		TargetKey:            targetKey,
		SeenUpToSeq:          frontier.ObservedSeq,
		LatestKnownSeq:       frontier.LatestKnownSeq,
		NeedsRefreshUntilSeq: frontier.NeedsRefreshUntilSeq,
		SourceMessageIDs:     append([]string(nil), frontier.SourceMessageIDs...),
	}, nil
}

func (h *mcpHandler) freshnessTracker() messageFreshnessTracker {
	if h == nil || h.c == nil || h.runtime == nil {
		return nil
	}
	return h.runtime
}

func (h *mcpHandler) rememberObservedMessages(agentID string, messages []types.InboxMessage) {
	tracker := h.freshnessTracker()
	if tracker == nil || len(messages) == 0 {
		return
	}
	tracker.RememberMessageFreshnessObserved(agentID, messages)
}

func (h *mcpHandler) rememberContiguousHistoryCoverage(agentID string, targetKey string, coversThroughSeq int64, sourceMessageID string) {
	tracker := h.freshnessTracker()
	targetKey = strings.TrimSpace(targetKey)
	if tracker == nil || targetKey == "" || coversThroughSeq <= 0 {
		return
	}
	tracker.RememberMessageFreshnessObserved(agentID, []types.InboxMessage{{
		MessageID: strings.TrimSpace(sourceMessageID),
		TargetKey: targetKey,
		TargetSeq: coversThroughSeq,
	}})
}

type heldSendRefresh struct {
	Text                                 string
	RuntimeDeliveryReceipts              []string
	RuntimeDeliveryCommitOwner           types.MCPToolResultDeliveryOwner
	RefreshedTargetMessageCount          int
	RefreshedTargetMessageHasMore        bool
	RefreshedTargetMessageRemainingCount int
	TargetContextRefreshComplete         bool
	PendingRuntimeInputCoversTarget      bool
}

func (h *mcpHandler) heldRefresh(ctx context.Context, agentID string, api *apiclient.Client, target string, result agentSendMessageResponse) heldSendRefresh {
	targetKey := strings.TrimSpace(result.TargetKey)
	if targetKey == "" {
		return heldSendRefresh{}
	}
	needsRefreshUntilSeq := result.NeedsRefreshUntilSeq
	if needsRefreshUntilSeq <= 0 {
		needsRefreshUntilSeq = result.LatestSeq
	}

	var sections []string
	maxInboxSeq := int64(0)
	refreshedTargetMessageCount := 0
	refreshedTargetMessageHasMore := false
	refreshedTargetMessageRemainingCount := 0
	pendingRuntimeInputCoversTarget := false
	runtimeDeliveryReceipts := []string(nil)
	var runtimeDeliveryCommitOwner types.MCPToolResultDeliveryOwner
	if h != nil && h.c != nil && h.runtime != nil {
		var (
			page types.InboxReadPage
			err  error
			read bool
		)
		runtimeSessionID := requestHeader(ctx, types.MCPRuntimeSessionHeader)
		if strings.TrimSpace(runtimeSessionID) != "" {
			page, err = h.runtime.InboxMessagesForFreshnessPageForRuntime(
				agentID, targetKey, needsRefreshUntilSeq, inboxMessagesDefaultLimit, runtimeSessionID,
			)
		} else {
			page, err = h.runtime.InboxMessagesForFreshnessPage(agentID, targetKey, needsRefreshUntilSeq, inboxMessagesDefaultLimit)
		}
		read = true
		if read {
			if err == nil {
				messages := page.Messages
				messages = filterFreshnessTargetMessages(messages, targetKey)
				for _, msg := range messages {
					if msg.TargetSeq > maxInboxSeq {
						maxInboxSeq = msg.TargetSeq
					}
				}
				refreshedTargetMessageHasMore = page.HasMore || page.RemainingCount > 0
				refreshedTargetMessageRemainingCount = page.RemainingCount
				if len(messages) > 0 {
					page.Messages = messages
					if formatted := h.formatInboxMessagePage(ctx, agentID, page, api, inboxMessagesDefaultLimit); formatted != nil && !formatted.IsError {
						if text := firstToolText(formatted); strings.TrimSpace(text) != "" && strings.TrimSpace(text) != "No new messages." {
							refreshedTargetMessageCount = len(messages)
							sections = append(sections, "Refreshed target messages:\n\n"+strings.TrimSpace(text))
							runtimeDeliveryReceipts = append(runtimeDeliveryReceipts, types.RuntimeDeliveryReceipts(formatted)...)
							runtimeDeliveryCommitOwner = formatted.DeliveryOwner
						}
					}
				}
			}
		}
		if needsRefreshUntilSeq > 0 {
			pendingRuntimeInputCoversTarget = h.runtime.PendingRuntimeInputCoversFreshness(agentID, targetKey, needsRefreshUntilSeq)
		}
	}
	if needsRefreshUntilSeq <= 0 || maxInboxSeq >= needsRefreshUntilSeq {
		return heldSendRefresh{
			Text:                                 strings.Join(sections, "\n\n"),
			RuntimeDeliveryReceipts:              runtimeDeliveryReceipts,
			RuntimeDeliveryCommitOwner:           runtimeDeliveryCommitOwner,
			RefreshedTargetMessageCount:          refreshedTargetMessageCount,
			RefreshedTargetMessageHasMore:        refreshedTargetMessageHasMore,
			RefreshedTargetMessageRemainingCount: refreshedTargetMessageRemainingCount,
			TargetContextRefreshComplete:         true,
			PendingRuntimeInputCoversTarget:      pendingRuntimeInputCoversTarget,
		}
	}

	return heldSendRefresh{
		Text:                                 strings.Join(sections, "\n\n"),
		RuntimeDeliveryReceipts:              runtimeDeliveryReceipts,
		RuntimeDeliveryCommitOwner:           runtimeDeliveryCommitOwner,
		RefreshedTargetMessageCount:          refreshedTargetMessageCount,
		RefreshedTargetMessageHasMore:        refreshedTargetMessageHasMore,
		RefreshedTargetMessageRemainingCount: refreshedTargetMessageRemainingCount,
		TargetContextRefreshComplete:         false,
		PendingRuntimeInputCoversTarget:      pendingRuntimeInputCoversTarget,
	}
}

// AcknowledgeMCPToolResultDelivery implements the optional MCP delivery
// callback. It is invoked by the HTTP transport only after its response write
// (or with an explicit uncertain/deferred outcome) and never during result
// construction.