package mcp

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

type targetSettlementCall struct {
	AgentID        string
	TargetKey      string
	SeenUpToSeq    int64
	ReplyMessageID string
}

type targetAwarePartialSendRuntime struct {
	serviceTestMessageRuntime
	call targetSettlementCall
}

func (r *targetAwarePartialSendRuntime) ConfirmVisibleReplyForTarget(agentID, targetKey string, seenUpToSeq int64, replyMessageID string) bool {
	r.call = targetSettlementCall{
		AgentID:        agentID,
		TargetKey:      targetKey,
		SeenUpToSeq:    seenUpToSeq,
		ReplyMessageID: replyMessageID,
	}
	return true
}

func TestMessageSendAcceptedPrefixSurvivesEveryFailure(t *testing.T) {
	for _, test := range []struct {
		name          string
		secondFailure func() (*http.Response, error)
	}{
		{
			name: "transport error",
			secondFailure: func() (*http.Response, error) {
				return nil, errors.New("network lost")
			},
		},
		{
			name: "application Error field",
			secondFailure: func() (*http.Response, error) {
				return messageSendJSONResponse(`{"error":"rejected"}`), nil
			},
		},
		{
			name: "held response",
			secondFailure: func() (*http.Response, error) {
				return messageSendJSONResponse(`{"state":"held","targetKey":"chn_ops","latestSeq":2}`), nil
			},
		},
		{
			name: "unexpected state",
			secondFailure: func() (*http.Response, error) {
				return messageSendJSONResponse(`{"state":"unknown","targetKey":"chn_ops"}`), nil
			},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			recorder := &partialSendRecorder{secondFailure: test.secondFailure}
			service := NewMessageService()
			handler, api := newPartialSendTestHandler(t, service, recorder)
			content := strings.Repeat("a", textMaxBytes) + "b"

			first := handler.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
			if test.name == "held response" {
				require.False(t, first.IsError, first.Content[0].Text)
			} else {
				require.True(t, first.IsError, first.Content[0].Text)
			}

			recorder.retry = true
			rebuilt, _ := newPartialSendTestHandler(t, service, recorder)
			second := rebuilt.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
			require.False(t, second.IsError, second.Content[0].Text)
			require.Equal(t, []int{0, 1, 1}, recorder.chunkIndexes())
		})
	}
}

func TestDynamicHandlerRebuildSharesRecoveryAndLimiterState(t *testing.T) {
	service := NewMessageService()
	recorder := &partialSendRecorder{
		secondFailure: func() (*http.Response, error) {
			return nil, errors.New("network lost")
		},
	}
	first, api := newPartialSendTestHandler(t, service, recorder)
	content := strings.Repeat("a", textMaxBytes) + "b"
	result := first.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
	require.True(t, result.IsError)

	limitKey := sendLimiterKey("server-a", "#ops", "local-agent")
	firstLimiter := service.limits.Get(limitKey)
	require.True(t, firstLimiter.Allow())

	recorder.retry = true
	rebuilt, _ := newPartialSendTestHandler(t, service, recorder)
	result = rebuilt.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
	require.False(t, result.IsError, result.Content[0].Text)
	require.Equal(t, []int{0, 1, 1}, recorder.chunkIndexes())
	require.Same(t, firstLimiter, rebuilt.service.limits.Get(limitKey))
}

func TestMessageSendSameKeyCannotOverwriteAcceptedPrefix(t *testing.T) {
	recorder := &partialSendRecorder{
		blockSecond: make(chan struct{}),
		secondSeen:  make(chan struct{}),
	}
	handler, api := newPartialSendTestHandler(t, NewMessageService(), recorder)
	content := strings.Repeat("a", textMaxBytes) + "b"

	firstDone := make(chan *types.MCPToolResult, 1)
	go func() {
		firstDone <- handler.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
	}()
	<-recorder.secondSeen

	secondDone := make(chan *types.MCPToolResult, 1)
	go func() {
		secondDone <- handler.sendMessageViaAPI(context.Background(), "local-agent", api, "#other", content, nil, nil, "")
	}()
	select {
	case result := <-secondDone:
		t.Fatalf("same-key second send entered before first released: %+v", result)
	case <-time.After(50 * time.Millisecond):
	}

	close(recorder.blockSecond)
	require.False(t, (<-firstDone).IsError)
	require.False(t, (<-secondDone).IsError)
	require.Equal(t, []string{"#ops", "#ops", "#other", "#other"}, recorder.targets())
}

func TestMessageSendAcceptedPrefixCanCoverEveryChunk(t *testing.T) {
	recorder := &partialSendRecorder{}
	handler, api := newPartialSendTestHandler(t, NewMessageService(), recorder)
	content := strings.Repeat("a", textMaxBytes) + "b"
	chunks := apiclient.SplitMessageText(content, textMaxBytes)
	lease, err := handler.service.acquirePartialSend(context.Background(), "server-a", "local-agent")
	require.NoError(t, err)
	lease.StorePending(pendingPartialSendState{
		Target:          "#ops",
		Signature:       partialSendSignature("#ops", chunks, nil, ""),
		LogicalClientID: "accepted-client",
		LastAcceptedSeq: 2,
		NextChunkIndex:  len(chunks),
	})
	lease.Release()

	result := handler.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
	require.False(t, result.IsError, result.Content[0].Text)
	require.Empty(t, recorder.chunkIndexes(), "a fully accepted prefix must not be posted again")
}

func TestMessageSendRecoveredFullPrefixSettlesStoredServerFrontier(t *testing.T) {
	recorder := &partialSendRecorder{}
	runtime := &targetAwarePartialSendRuntime{}
	handler, api := newPartialSendTestHandlerWithRuntime(t, NewMessageService(), recorder, runtime)
	content := strings.Repeat("a", textMaxBytes) + "b"
	chunks := apiclient.SplitMessageText(content, textMaxBytes)
	lease, err := handler.service.acquirePartialSend(context.Background(), "server-a", "local-agent")
	require.NoError(t, err)
	lease.StorePending(pendingPartialSendState{
		Target:                "#ops",
		Signature:             partialSendSignature("#ops", chunks, nil, ""),
		LogicalClientID:       "accepted-client",
		LastAcceptedSeq:       2,
		NextChunkIndex:        len(chunks),
		AcceptedChunks:        []acceptedSendChunk{{ChunkIndex: 0, MessageID: "reply-1", AcceptedSeq: 2, State: agentSendStateSent}, {ChunkIndex: 1, MessageID: "reply-2", AcceptedSeq: 3, State: agentSendStateSent}},
		SettlementTargetKey:   "chn_ops",
		SettlementSeenUpToSeq: 7,
		SettlementMessageID:   "reply-1",
	})
	lease.Release()

	result := handler.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
	require.False(t, result.IsError, result.Content[0].Text)
	require.Empty(t, recorder.chunkIndexes(), "full-prefix recovery must not post an already accepted chunk")
	require.Equal(t, targetSettlementCall{
		AgentID:        "local-agent",
		TargetKey:      "chn_ops",
		SeenUpToSeq:    7,
		ReplyMessageID: "reply-1",
	}, runtime.call)

	lease, err = handler.service.acquirePartialSend(context.Background(), "server-a", "local-agent")
	require.NoError(t, err)
	_, ok := lease.Pending()
	require.False(t, ok, "authoritative full-prefix settlement should retire recovery state")
	lease.Release()
}

func TestMessageSendResumedPrefixUsesFinalServerFrontierForSettlement(t *testing.T) {
	recorder := &partialSendRecorder{
		responseSeenUpToSeq: 7,
		secondFailure: func() (*http.Response, error) {
			return nil, errors.New("network lost")
		},
	}
	runtime := &targetAwarePartialSendRuntime{}
	service := NewMessageService()
	handler, api := newPartialSendTestHandlerWithRuntime(t, service, recorder, runtime)
	content := strings.Repeat("a", textMaxBytes) + "b"

	first := handler.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
	require.True(t, first.IsError, first.Content[0].Text)
	require.Equal(t, targetSettlementCall{}, runtime.call, "partial failure must not settle the accepted prefix")

	recorder.retry = true
	rebuilt, _ := newPartialSendTestHandlerWithRuntime(t, service, recorder, runtime)
	second := rebuilt.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
	require.False(t, second.IsError, second.Content[0].Text)
	require.Equal(t, []int{0, 1, 1}, recorder.chunkIndexes())
	require.Equal(t, []int64{0, 1, 1}, recorder.freshnessSeenUpToSeq(), "chunk sequencing may advance its request cursor independently of settlement coverage")
	require.Equal(t, targetSettlementCall{
		AgentID:        "local-agent",
		TargetKey:      "chn_ops",
		SeenUpToSeq:    7,
		ReplyMessageID: "message-0",
	}, runtime.call)
}

func TestMessageSendHeldPrefixRemainsNonSettlingUntilRetrySucceeds(t *testing.T) {
	recorder := &partialSendRecorder{
		responseSeenUpToSeq: 7,
		secondFailure: func() (*http.Response, error) {
			return messageSendJSONResponse(`{"state":"held","targetKey":"chn_ops","latestSeq":2}`), nil
		},
	}
	runtime := &targetAwarePartialSendRuntime{}
	service := NewMessageService()
	handler, api := newPartialSendTestHandlerWithRuntime(t, service, recorder, runtime)
	content := strings.Repeat("a", textMaxBytes) + "b"

	held := handler.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
	require.False(t, held.IsError, held.Content[0].Text)
	require.Equal(t, targetSettlementCall{}, runtime.call, "held partial sends must not settle any wake")

	recorder.retry = true
	rebuilt, _ := newPartialSendTestHandlerWithRuntime(t, service, recorder, runtime)
	sent := rebuilt.sendMessageViaAPI(context.Background(), "local-agent", api, "#ops", content, nil, nil, "")
	require.False(t, sent.IsError, sent.Content[0].Text)
	require.Equal(t, targetSettlementCall{
		AgentID:        "local-agent",
		TargetKey:      "chn_ops",
		SeenUpToSeq:    7,
		ReplyMessageID: "message-0",
	}, runtime.call)
}

func TestMessageSendDoesNotUseWakeTokenWhenServerCoverageIsMissing(t *testing.T) {
	runtime := &targetAwarePartialSendRuntime{}
	h := &mcpHandler{
		c:       serviceTestAgentTeamConnection{id: "server-a"},
		runtime: runtime,
	}

	h.confirmAcceptedVisibleReply("local-agent", "#ops", agentSendMessageResponse{
		State:     agentSendStateSent,
		MessageID: "reply-1",
	}, "reply-1")

	require.Equal(t, targetSettlementCall{}, runtime.call)
}

func TestMessageSendTeamInvocationRemainsOneMessageAboveChunkThreshold(t *testing.T) {
	recorder := &partialSendRecorder{}
	handler, api := newPartialSendTestHandler(t, NewMessageService(), recorder)
	content := "@leader " + strings.Repeat("a", textMaxBytes)

	result := handler.sendMessageViaAPI(
		context.Background(), "local-agent", api, "#team_abcdefghijklmn", content,
		nil, nil, "team-1",
	)

	require.False(t, result.IsError, result.Content[0].Text)
	require.Len(t, recorder.requests, 1)
	require.Equal(t, content, recorder.requests[0].Content)
	require.Equal(t, "team-1", recorder.requests[0].TeamInvokeTeamID)
	require.Equal(t, []int{-1}, recorder.chunkIndexes(), "a Team handoff must not receive chunk client ids")
}

func TestMessageSendTeamInvocationAcceptsExactByteLimit(t *testing.T) {
	recorder := &partialSendRecorder{}
	handler, api := newPartialSendTestHandler(t, NewMessageService(), recorder)
	content := strings.Repeat("你", (teamInvokeMessageMaxContentBytes-2)/len("你")) + "ab"
	require.Len(t, []byte(content), teamInvokeMessageMaxContentBytes)

	result := handler.sendMessageViaAPI(
		context.Background(), "local-agent", api, "#team_abcdefghijklmn", content,
		nil, nil, "team-1",
	)

	require.False(t, result.IsError, result.Content[0].Text)
	require.Len(t, recorder.requests, 1)
	require.Equal(t, content, recorder.requests[0].Content)
	require.Equal(t, []int{-1}, recorder.chunkIndexes(), "the byte-limit Team handoff must remain one message")
}

func TestMessageSendOversizeTeamInvocationFailsBeforeAnyMessagePost(t *testing.T) {
	recorder := &partialSendRecorder{}
	handler, api := newPartialSendTestHandler(t, NewMessageService(), recorder)
	content := strings.Repeat("a", teamInvokeMessageMaxContentBytes+1)

	result := handler.sendMessageViaAPI(
		context.Background(), "local-agent", api, "#team_abcdefghijklmn", content,
		nil, nil, "team-1",
	)

	require.True(t, result.IsError)
	require.Contains(t, result.Content[0].Text, "Team invocation message is too large (max 32768 bytes)")
	require.Empty(t, recorder.requests, "oversize Team invocation must not partially persist a carrier")
}

type partialSendRecorder struct {
	mu                  sync.Mutex
	requests            []partialSendRequest
	sendCalls           int
	retry               bool
	responseSeenUpToSeq int64
	secondFailure       func() (*http.Response, error)
	blockSecond         chan struct{}
	secondSeen          chan struct{}
}

type partialSendRequest struct {
	Target           string `json:"target"`
	Content          string `json:"content"`
	ClientMessageID  string `json:"clientMessageId"`
	TeamInvokeTeamID string `json:"teamInvokeTeamId"`
	Freshness        struct {
		SeenUpToSeq int64 `json:"seenUpToSeq"`
	} `json:"freshness"`
}

func (r *partialSendRecorder) RoundTrip(request *http.Request) (*http.Response, error) {
	switch request.URL.Path {
	case "/api/v1/workspaces/wsp_1/agents/server-agent/runtime/resolve-channel":
		return messageSendJSONResponse(`{"scopeId":"chn_ops","targetKey":"chn_ops"}`), nil
	case "/api/v1/workspaces/wsp_1/agents/server-agent/runtime/send":
		body, err := io.ReadAll(request.Body)
		if err != nil {
			return nil, err
		}
		var sent partialSendRequest
		if err := json.Unmarshal(body, &sent); err != nil {
			return nil, err
		}
		r.mu.Lock()
		r.requests = append(r.requests, sent)
		r.sendCalls++
		call := r.sendCalls
		retry := r.retry
		responseSeenUpToSeq := r.responseSeenUpToSeq
		failure := r.secondFailure
		block := r.blockSecond
		seen := r.secondSeen
		r.mu.Unlock()

		if block != nil && call == 2 {
			if seen != nil {
				close(seen)
			}
			<-block
		}
		if !retry && call == 2 && failure != nil {
			return failure()
		}
		chunkIndex := 0
		if strings.HasSuffix(sent.ClientMessageID, ":1") {
			chunkIndex = 1
		}
		return messageSendJSONResponse(fmt.Sprintf(
			`{"state":"sent","messageId":"message-%d","targetKey":"chn_ops","acceptedSeq":%d,"seenUpToSeq":%d}`,
			chunkIndex, chunkIndex+1, responseSeenUpToSeq,
		)), nil
	default:
		return nil, errors.New("unexpected path: " + request.URL.Path)
	}
}

func (r *partialSendRecorder) chunkIndexes() []int {
	r.mu.Lock()
	defer r.mu.Unlock()
	indexes := make([]int, 0, len(r.requests))
	for _, request := range r.requests {
		switch {
		case strings.HasSuffix(request.ClientMessageID, ":0"):
			indexes = append(indexes, 0)
		case strings.HasSuffix(request.ClientMessageID, ":1"):
			indexes = append(indexes, 1)
		default:
			indexes = append(indexes, -1)
		}
	}
	return indexes
}

func (r *partialSendRecorder) targets() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	targets := make([]string, 0, len(r.requests))
	for _, request := range r.requests {
		targets = append(targets, request.Target)
	}
	return targets
}

func (r *partialSendRecorder) freshnessSeenUpToSeq() []int64 {
	r.mu.Lock()
	defer r.mu.Unlock()
	seen := make([]int64, 0, len(r.requests))
	for _, request := range r.requests {
		seen = append(seen, request.Freshness.SeenUpToSeq)
	}
	return seen
}

func newPartialSendTestHandler(t *testing.T, service *MessageService, transport http.RoundTripper) (*mcpHandler, *apiclient.Client) {
	t.Helper()
	return newPartialSendTestHandlerWithRuntime(t, service, transport, &partialSendRuntime{})
}

type partialSendTestRuntime interface {
	MessageRuntime
}

func newPartialSendTestHandlerWithRuntime(t *testing.T, service *MessageService, transport http.RoundTripper, runtime partialSendTestRuntime) (*mcpHandler, *apiclient.Client) {
	t.Helper()
	connection := partialSendConnection{
		id: "server-a",
		client: &apiclient.Client{
			ServerURL:        "http://agentteam.invalid",
			AgentID:          "server-agent",
			AuthToken:        "token",
			RouteWorkspaceID: "wsp_1",
			HTTPClient:       &http.Client{Transport: transport},
		},
	}
	config := MessageHandlerConfig{
		Connection: connection,
		Runtime:    runtime,
	}
	if provider, ok := runtime.(SendScopeProvider); ok {
		config.SendScopeProvider = provider
	}
	handler, err := service.newMessageHandler(config)
	require.NoError(t, err)
	return handler, connection.client
}

func messageSendJSONResponse(body string) *http.Response {
	return &http.Response{
		StatusCode: http.StatusOK,
		Header:     http.Header{"Content-Type": []string{"application/json"}},
		Body:       io.NopCloser(bytes.NewBufferString(body)),
	}
}

type partialSendConnection struct {
	id     string
	client *apiclient.Client
}

func (c partialSendConnection) ID() string { return c.id }
func (c partialSendConnection) APIClient(...string) *apiclient.Client {
	return c.client
}
func (partialSendConnection) Send(context.Context, types.Metadata, types.OutboundMessage) error {
	return nil
}

type partialSendRuntime struct {
	serviceTestMessageRuntime
}

func (partialSendRuntime) Session(string) (types.SessionSnapshot, bool) {
	return types.SessionSnapshot{ServerAgentID: "server-agent"}, true
}

func (partialSendRuntime) MessageFreshnessSnapshot(string, string) (types.MessageFreshnessFrontier, bool) {
	return types.MessageFreshnessFrontier{}, true
}

func (partialSendRuntime) SendScope(context.Context, string) (messagesend.Scope, bool) {
	return messagesend.Scope{}, true
}