package daemon

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

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/store/agent"
)

func TestWorkspaceDeleteRemoteRequestUsesSessionLink(t *testing.T) {
	serverConn, desktopConn := newDaemonProtocolTestPair(t)
	client := newWorkspaceListRemoteReadyClient(serverConn)
	now := time.Date(2026, 6, 30, 9, 0, 0, 0, time.UTC)
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})}

	go func() {
		var assignment RemoteRequestAssignment[json.RawMessage]
		require.NoError(t, desktopConn.ReadJSON(&assignment))
		require.Equal(t, sessionLinkRemoteRequestAssignType, assignment.Type)
		require.Equal(t, remoteEndpointWorkspaceDelete, assignment.EndpointID)
		var payload workspaceDeleteRemoteRequestPayload
		require.NoError(t, json.Unmarshal(assignment.Payload, &payload))
		require.Equal(t, "agt_1", payload.AgentID)
		require.Equal(t, ".agents/skills", payload.Path)
		result := RemoteRequestResult[json.RawMessage]{
			Type:      sessionLinkRemoteRequestResultType,
			RequestID: assignment.RequestID,
			Success:   true,
			Outcome:   "ok",
			Payload: mustMarshalRemoteWorkspaceRawTest(t, workspaceDeleteRemoteRequestResponse{
				AgentID: payload.AgentID,
				Success: true,
			}),
		}
		require.NoError(t, handleSessionLinkRemoteRequestResultCommand(hub, context.Background(), client, mustMarshalRemoteWorkspaceTest(t, result)))
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	require.NoError(t, hub.deleteWorkspacePath(ctx, client, "agt_1", ".agents/skills"))
}

func TestWorkspaceApplySeedFilesRemoteRequestUsesElasticUnderPolicy(t *testing.T) {
	now := time.Date(2026, 6, 30, 9, 10, 0, 0, time.UTC)
	manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: manager}
	client := newElasticReadyClient()
	descriptor, ok := defaultRemoteEndpointRegistry().Lookup(remoteEndpointWorkspaceApplySeedFiles, 1)
	require.True(t, ok)
	open, _, err := manager.CreateLinkGrant(client, descriptor)
	require.NoError(t, err)
	sender := newCaptureElasticSender()
	_, err = manager.AcceptLink(client.workspaceID, open.LinkID, open.Token, sender)
	require.NoError(t, err)

	go func() {
		message := <-sender.messages
		assignment, ok := message.(RemoteRequestAssignment[json.RawMessage])
		require.True(t, ok)
		require.Equal(t, remoteEndpointWorkspaceApplySeedFiles, assignment.EndpointID)
		var payload workspaceApplySeedFilesRemoteRequestPayload
		require.NoError(t, json.Unmarshal(assignment.Payload, &payload))
		require.Equal(t, "agt_1", payload.AgentID)
		require.Equal(t, []agentWorkspaceSeedFile{{Path: "seed.md", Content: "seed"}}, payload.SeedFiles)
		manager.CompleteRemoteRequest(RemoteRequestResult[json.RawMessage]{
			Type:      sessionLinkRemoteRequestResultType,
			RequestID: assignment.RequestID,
			Success:   true,
			Outcome:   "ok",
			Payload: mustMarshalRemoteWorkspaceRawTest(t, workspaceApplySeedFilesRemoteRequestResponse{
				AgentID: payload.AgentID,
				Success: true,
			}),
		})
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	err = hub.applySeedFiles(ctx, client, "agt_1", []agent.WorkspaceSeedFile{{Path: "seed.md", Content: "seed"}})
	require.NoError(t, err)
}

func TestWorkspaceApplySeedFilesRemoteRequestDoesNotRetryAfterEndpointError(t *testing.T) {
	now := time.Date(2026, 6, 30, 9, 20, 0, 0, time.UTC)
	manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: manager}
	client := newElasticReadyClient()
	descriptor, ok := defaultRemoteEndpointRegistry().Lookup(remoteEndpointWorkspaceApplySeedFiles, 1)
	require.True(t, ok)
	open, _, err := manager.CreateLinkGrant(client, descriptor)
	require.NoError(t, err)
	sender := newCaptureElasticSender()
	_, err = manager.AcceptLink(client.workspaceID, open.LinkID, open.Token, sender)
	require.NoError(t, err)

	go func() {
		message := <-sender.messages
		assignment, ok := message.(RemoteRequestAssignment[json.RawMessage])
		require.True(t, ok)
		manager.CompleteRemoteRequest(RemoteRequestResult[json.RawMessage]{
			Type:      sessionLinkRemoteRequestResultType,
			RequestID: assignment.RequestID,
			Success:   false,
			Outcome:   "endpoint_error",
			Error:     "endpoint_error",
		})
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	err = hub.applySeedFiles(ctx, client, "agt_1", []agent.WorkspaceSeedFile{{Path: "seed.md", Content: "seed"}})
	require.ErrorContains(t, err, "endpoint_error")
	select {
	case message := <-sender.messages:
		t.Fatalf("unexpected retry message after non-idempotent endpoint error: %#v", message)
	case <-time.After(50 * time.Millisecond):
	}
}

func TestWorkspaceDeleteRemoteRequestFallsBackForOldDesktop(t *testing.T) {
	serverConn, desktopConn := newDaemonProtocolTestPair(t)
	client := &clientState{
		machineID:                 "mac_1",
		workspaceID:               "wsp_1",
		conn:                      serverConn,
		capabilities:              map[string]struct{}{},
		supportedEndpointVersions: make(map[string]map[int]struct{}),
		deleteRequests:            make(map[string]chan deleteWorkspacePathResult),
		remoteRequests:            make(map[string]chan RemoteRequestResult[json.RawMessage]),
	}
	now := time.Date(2026, 6, 30, 9, 30, 0, 0, time.UTC)
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})}

	go func() {
		var command agentWorkspaceDeleteCommand
		require.NoError(t, desktopConn.ReadJSON(&command))
		require.Equal(t, sessionLinkCommandAgentWorkspaceDelete, command.Type)
		require.Equal(t, "agt_1", command.AgentID)
		require.Equal(t, ".agents/skills", command.Path)
		require.NoError(t, handleSessionLinkAgentWorkspaceDeleteResultCommand(nil, context.Background(), client, mustMarshalRemoteWorkspaceTest(t, deleteWorkspacePathResultMessage{
			Type:      sessionLinkCommandAgentWorkspaceDeleteResult,
			AgentID:   command.AgentID,
			RequestID: command.RequestID,
			Success:   true,
		})))
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	require.NoError(t, hub.deleteWorkspacePath(ctx, client, "agt_1", ".agents/skills"))
}

func TestWorkspaceReadRemoteRequestCapacityFallsBackToLegacy(t *testing.T) {
	serverConn, desktopConn := newDaemonProtocolTestPair(t)
	client := newWorkspaceListRemoteReadyClient(serverConn)
	now := time.Date(2026, 6, 30, 9, 40, 0, 0, time.UTC)
	manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{MaxPerMachine: 1})
	descriptor, ok := defaultRemoteEndpointRegistry().Lookup(remoteEndpointWorkspaceRead, 1)
	require.True(t, ok)
	_, _, err := manager.CreateLinkGrant(client, descriptor)
	require.NoError(t, err)
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: manager}

	go func() {
		var command agentWorkspaceReadCommand
		require.NoError(t, desktopConn.ReadJSON(&command))
		require.Equal(t, sessionLinkCommandAgentWorkspaceRead, command.Type)
		require.Equal(t, "README.md", command.Path)
		content := "legacy"
		require.NoError(t, handleSessionLinkAgentWorkspaceFileContentCommand(nil, context.Background(), client, mustMarshalRemoteWorkspaceTest(t, fileContentMessage{
			Type:      sessionLinkCommandAgentWorkspaceFileContent,
			AgentID:   command.AgentID,
			RequestID: command.RequestID,
			Content:   &content,
		})))
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	content, err := hub.readWorkspaceFile(ctx, client, "agt_1", "", "README.md")
	require.NoError(t, err)
	require.NotNil(t, content.Content)
	require.Equal(t, "legacy", *content.Content)
}

func mustMarshalRemoteWorkspaceRawTest(t *testing.T, value any) json.RawMessage {
	t.Helper()
	return json.RawMessage(mustMarshalRemoteWorkspaceTest(t, value))
}