package wire

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/memorywatcher"
)

type memorySyncTestDispatcher struct {
	agentteamTestDispatcher
	status       memorywatcher.Status
	statusKnown  bool
	retriedAgent string
}

func (d *memorySyncTestDispatcher) MemorySyncStatus(string) (memorywatcher.Status, bool) {
	return d.status, d.statusKnown
}

func (d *memorySyncTestDispatcher) RetryMemorySync(_ context.Context, agentID string) {
	d.retriedAgent = agentID
}

func TestRemoteMemorySyncEndpointsExposeDaemonStateAndRetry(t *testing.T) {
	retryAt := time.Date(2026, time.August, 21, 9, 0, 0, 0, time.UTC)
	dispatcher := &memorySyncTestDispatcher{
		statusKnown: true,
		status: memorywatcher.Status{
			State:   memorywatcher.StatusPendingSync,
			Reason:  memorywatcher.StatusReasonPushFailed,
			RetryAt: retryAt,
		},
	}
	protocol := &Protocol{disp: dispatcher}
	assignment := remoteRequestAssignmentMsg{
		Payload: json.RawMessage(`{"agentId":"agt_1"}`),
	}

	raw, outcome, err := handleRemoteRuntimeMemorySyncStatus(t.Context(), protocol, assignment)
	require.NoError(t, err)
	require.Equal(t, "ok", outcome)
	var statusResponse remoteMemorySyncStatusResponse
	require.NoError(t, json.Unmarshal(raw, &statusResponse))
	require.True(t, statusResponse.Available)
	require.NotNil(t, statusResponse.Status)
	require.Equal(t, memorywatcher.StatusPendingSync, statusResponse.Status.State)
	require.Equal(t, memorywatcher.StatusReasonPushFailed, statusResponse.Status.Reason)
	require.Equal(t, retryAt, statusResponse.Status.RetryAt)

	raw, outcome, err = handleRemoteRuntimeMemorySyncRetry(t.Context(), protocol, assignment)
	require.NoError(t, err)
	require.Equal(t, "ok", outcome)
	var retryResponse remoteMemorySyncRetryResponse
	require.NoError(t, json.Unmarshal(raw, &retryResponse))
	require.True(t, retryResponse.RetryRequested)
	require.Equal(t, "agt_1", dispatcher.retriedAgent)
}

func TestRemoteMemorySyncRetryDoesNotCreateUnknownState(t *testing.T) {
	dispatcher := &memorySyncTestDispatcher{}
	protocol := &Protocol{disp: dispatcher}
	assignment := remoteRequestAssignmentMsg{
		Payload: json.RawMessage(`{"agentId":"agt_1"}`),
	}

	raw, outcome, err := handleRemoteRuntimeMemorySyncRetry(t.Context(), protocol, assignment)
	require.NoError(t, err)
	require.Equal(t, "ok", outcome)
	var response remoteMemorySyncRetryResponse
	require.NoError(t, json.Unmarshal(raw, &response))
	require.False(t, response.RetryRequested)
	require.Empty(t, dispatcher.retriedAgent)
}