package mcp

import (
	"context"
	"errors"
	"fmt"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestMessageServiceSerializesSamePartialSendKey(t *testing.T) {
	service := newMessageServiceWithClock(time.Now)
	first, err := service.acquirePartialSend(context.Background(), " server-a ", " agent-a ")
	require.NoError(t, err)
	defer first.Release()

	entered := make(chan *partialSendLease, 1)
	go func() {
		lease, acquireErr := service.acquirePartialSend(context.Background(), "server-a", "agent-a")
		if acquireErr == nil {
			entered <- lease
		}
	}()

	select {
	case lease := <-entered:
		lease.Release()
		t.Fatal("second same-key lease entered before the first was released")
	case <-time.After(50 * time.Millisecond):
	}

	first.Release()
	select {
	case lease := <-entered:
		lease.Release()
	case <-time.After(time.Second):
		t.Fatal("second same-key lease did not enter after release")
	}
}

func TestMessageServiceAllowsDifferentPartialSendKeys(t *testing.T) {
	service := newMessageServiceWithClock(time.Now)
	first, err := service.acquirePartialSend(context.Background(), "server-a", "agent-a")
	require.NoError(t, err)
	defer first.Release()

	second, err := service.acquirePartialSend(context.Background(), "server-a", "agent-b")
	require.NoError(t, err)
	second.Release()
}

func TestMessageServiceRetainsPendingRecoveryForOneHour(t *testing.T) {
	clock := newMessageServiceTestClock()
	service := newMessageServiceWithClock(clock.Now)
	lease, err := service.acquirePartialSend(context.Background(), "server-a", "agent-a")
	require.NoError(t, err)
	want := pendingPartialSendState{
		Target: "#ops", Signature: "sig", LogicalClientID: "client", LastAcceptedSeq: 7, NextChunkIndex: 2,
	}
	lease.StorePending(want)
	lease.Release()

	clock.Advance(time.Hour - time.Nanosecond)
	lease, err = service.acquirePartialSend(context.Background(), "server-a", "agent-a")
	require.NoError(t, err)
	got, ok := lease.Pending()
	require.True(t, ok)
	require.Equal(t, want, got)
	lease.Release()
}

func TestMessageServicePendingReadRefreshesRetention(t *testing.T) {
	clock := newMessageServiceTestClock()
	service := newMessageServiceWithClock(clock.Now)
	lease, err := service.acquirePartialSend(context.Background(), "server-a", "agent-a")
	require.NoError(t, err)
	lease.StorePending(pendingPartialSendState{Signature: "sig"})
	lease.Release()

	clock.Advance(45 * time.Minute)
	lease, err = service.acquirePartialSend(context.Background(), "server-a", "agent-a")
	require.NoError(t, err)
	_, ok := lease.Pending()
	require.True(t, ok)
	lease.Release()

	clock.Advance(45 * time.Minute)
	lease, err = service.acquirePartialSend(context.Background(), "server-a", "agent-a")
	require.NoError(t, err)
	_, ok = lease.Pending()
	require.True(t, ok)
	lease.Release()
}

func TestMessageServiceExpiresIdlePendingRecovery(t *testing.T) {
	clock := newMessageServiceTestClock()
	service := newMessageServiceWithClock(clock.Now)
	lease, err := service.acquirePartialSend(context.Background(), "server-a", "agent-a")
	require.NoError(t, err)
	lease.StorePending(pendingPartialSendState{
		Target: "#ops", Signature: "sig", LogicalClientID: "client", NextChunkIndex: 1,
	})
	lease.Release()

	clock.Advance(time.Hour + time.Nanosecond)
	lease, err = service.acquirePartialSend(context.Background(), "server-a", "agent-a")
	require.NoError(t, err)
	_, ok := lease.Pending()
	require.False(t, ok)
	lease.Release()
}

func TestMessageServiceEvictsOldestEmptyEntryAtCapacity(t *testing.T) {
	clock := newMessageServiceTestClock()
	service := newMessageServiceWithClock(clock.Now)
	service.partialSendCapacity = 2

	oldest, err := service.acquirePartialSend(context.Background(), "server", "oldest")
	require.NoError(t, err)
	oldest.Release()
	clock.Advance(time.Second)
	newer, err := service.acquirePartialSend(context.Background(), "server", "newer")
	require.NoError(t, err)
	newer.Release()

	replacement, err := service.acquirePartialSend(context.Background(), "server", "replacement")
	require.NoError(t, err)
	replacement.Release()

	require.False(t, messageServiceHasEntry(service, "server", "oldest"))
	require.True(t, messageServiceHasEntry(service, "server", "newer"))
	require.True(t, messageServiceHasEntry(service, "server", "replacement"))
}

func TestMessageServiceFailsBusyWhenAll1024EntriesAreLive(t *testing.T) {
	service := newMessageServiceWithClock(time.Now)
	leases := make([]*partialSendLease, 0, maxPartialSendEntries)
	for index := 0; index < maxPartialSendEntries; index++ {
		lease, err := service.acquirePartialSend(context.Background(), "server", fmt.Sprintf("agent-%d", index))
		require.NoError(t, err)
		leases = append(leases, lease)
	}
	t.Cleanup(func() {
		for _, lease := range leases {
			lease.Release()
		}
	})

	_, err := service.acquirePartialSend(context.Background(), "server", "overflow")
	require.ErrorIs(t, err, errMessageServiceBusy)
}

func TestMessageServiceNeverEvictsHeldOrNonExpiredPendingEntry(t *testing.T) {
	clock := newMessageServiceTestClock()
	service := newMessageServiceWithClock(clock.Now)
	service.partialSendCapacity = 2

	held, err := service.acquirePartialSend(context.Background(), "server", "held")
	require.NoError(t, err)
	pending, err := service.acquirePartialSend(context.Background(), "server", "pending")
	require.NoError(t, err)
	pending.StorePending(pendingPartialSendState{
		Target: "#ops", Signature: "sig", LogicalClientID: "client", NextChunkIndex: 1,
	})
	pending.Release()

	_, err = service.acquirePartialSend(context.Background(), "server", "overflow")
	require.ErrorIs(t, err, errMessageServiceBusy)
	require.True(t, messageServiceHasEntry(service, "server", "held"))
	require.True(t, messageServiceHasEntry(service, "server", "pending"))
	held.Release()
}

func TestMessageServiceLeaseReleaseIsIdempotent(t *testing.T) {
	service := newMessageServiceWithClock(time.Now)
	lease, err := service.acquirePartialSend(context.Background(), "server", "agent")
	require.NoError(t, err)
	lease.Release()
	lease.Release()

	next, err := service.acquirePartialSend(context.Background(), "server", "agent")
	require.NoError(t, err)
	next.Release()
}

func TestMessageServiceCancelledWaitDoesNotAcquire(t *testing.T) {
	service := newMessageServiceWithClock(time.Now)
	first, err := service.acquirePartialSend(context.Background(), "server", "agent")
	require.NoError(t, err)
	defer first.Release()

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err = service.acquirePartialSend(ctx, "server", "agent")
	require.ErrorIs(t, err, context.Canceled)
}

func TestMessageServiceNormalizesPartialSendKey(t *testing.T) {
	service := newMessageServiceWithClock(time.Now)
	first, err := service.acquirePartialSend(context.Background(), " server ", " agent ")
	require.NoError(t, err)
	first.StorePending(pendingPartialSendState{Signature: "one"})
	first.Release()

	second, err := service.acquirePartialSend(context.Background(), "server", "agent")
	require.NoError(t, err)
	pending, ok := second.Pending()
	require.True(t, ok)
	require.Equal(t, "one", pending.Signature)
	second.Release()
}

func TestMessageServiceLeaseProtectsStateUntilRelease(t *testing.T) {
	service := newMessageServiceWithClock(time.Now)
	first, err := service.acquirePartialSend(context.Background(), "server", "agent")
	require.NoError(t, err)
	first.StorePending(pendingPartialSendState{Signature: "first", NextChunkIndex: 1})

	var acquired atomic.Bool
	done := make(chan error, 1)
	go func() {
		second, acquireErr := service.acquirePartialSend(context.Background(), "server", "agent")
		if acquireErr != nil {
			done <- acquireErr
			return
		}
		acquired.Store(true)
		state, ok := second.Pending()
		if !ok || state.Signature != "first" || state.NextChunkIndex != 1 {
			done <- errors.New("second lease observed incomplete recovery state")
			second.Release()
			return
		}
		second.Release()
		done <- nil
	}()

	time.Sleep(25 * time.Millisecond)
	require.False(t, acquired.Load())
	first.Release()
	require.NoError(t, <-done)
}

type messageServiceTestClock struct {
	now time.Time
}

func newMessageServiceTestClock() *messageServiceTestClock {
	return &messageServiceTestClock{now: time.Unix(1_700_000_000, 0)}
}

func (c *messageServiceTestClock) Now() time.Time {
	return c.now
}

func (c *messageServiceTestClock) Advance(duration time.Duration) {
	c.now = c.now.Add(duration)
}

func messageServiceHasEntry(service *MessageService, serverID string, agentID string) bool {
	service.mu.Lock()
	defer service.mu.Unlock()
	_, ok := service.partialSends[partialSendKey{serverID: serverID, localAgentID: agentID}]
	return ok
}