package daemon

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

	"github.com/stretchr/testify/require"

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

func TestMachineWorkspaceScanRemoteRequestUsesElasticUnderPolicy(t *testing.T) {
	now := time.Date(2026, 7, 1, 0, 40, 0, 0, time.UTC)
	manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: manager}
	client := newElasticReadyClient()
	descriptor, ok := defaultRemoteEndpointRegistry().Lookup(remoteEndpointMachineWorkspaceScan, 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, remoteEndpointMachineWorkspaceScan, assignment.EndpointID)
		var payload machineWorkspaceScanRemoteRequestPayload
		require.NoError(t, json.Unmarshal(assignment.Payload, &payload))
		require.Equal(t, 80, payload.Limit)
		manager.CompleteRemoteRequest(RemoteRequestResult[json.RawMessage]{
			Type:      sessionLinkRemoteRequestResultType,
			RequestID: assignment.RequestID,
			Success:   true,
			Outcome:   "ok",
			Payload: mustMarshalRemoteWorkspaceRawTest(t, machineWorkspaceScanRemoteRequestResponse{
				Directories: []agent.WorkspaceDirectory{{DirectoryName: "agt_1", FileCount: 3}},
				Limit:       payload.Limit,
				HasMore:     true,
				NextCursor:  "80",
			}),
		})
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	page, err := hub.scanMachineWorkspace(ctx, client, agent.WorkspaceListOptions{Limit: 80})
	require.NoError(t, err)
	require.Equal(t, []agent.WorkspaceDirectory{{DirectoryName: "agt_1", FileCount: 3}}, page.Directories)
	require.Equal(t, 80, page.Limit)
	require.True(t, page.HasMore)
	require.Equal(t, "80", page.NextCursor)
}

func TestMachineWorkspaceDeleteRemoteRequestDoesNotRetryAfterEndpointError(t *testing.T) {
	serverConn, desktopConn := newDaemonProtocolTestPair(t)
	client := newWorkspaceListRemoteReadyClient(serverConn)
	now := time.Date(2026, 7, 1, 0, 50, 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, remoteEndpointMachineWorkspaceDelete, assignment.EndpointID)
		var payload machineWorkspaceDeleteRemoteRequestPayload
		require.NoError(t, json.Unmarshal(assignment.Payload, &payload))
		require.Equal(t, "agt-delete", payload.DirectoryName)
		require.NoError(t, handleSessionLinkRemoteRequestResultCommand(hub, context.Background(), client, mustMarshalRemoteWorkspaceTest(t, 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()
	deleted, err := hub.deleteMachineWorkspaceDirectory(ctx, client, "agt-delete")
	require.ErrorContains(t, err, "endpoint_error")
	require.False(t, deleted)
}

func TestMachineWorkspaceDeleteRemoteRequestFallsBackForOldDesktop(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{}),
		remoteRequests:            make(map[string]chan RemoteRequestResult[json.RawMessage]),
	}
	now := time.Date(2026, 7, 1, 1, 0, 0, 0, time.UTC)
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})}

	go func() {
		var command machineWorkspaceDeleteCommand
		require.NoError(t, desktopConn.ReadJSON(&command))
		require.Equal(t, sessionLinkCommandMachineWorkspaceDelete, command.Type)
		require.Equal(t, "agt-delete", command.DirectoryName)
		require.NoError(t, handleSessionLinkMachineWorkspaceDeleteResultCommand(nil, context.Background(), client, mustMarshalRemoteWorkspaceTest(t, deleteResultMessage{
			Type:          sessionLinkCommandMachineWorkspaceDeleteResult,
			DirectoryName: command.DirectoryName,
			Success:       true,
		})))
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	deleted, err := hub.deleteMachineWorkspaceDirectory(ctx, client, "agt-delete")
	require.NoError(t, err)
	require.True(t, deleted)
}

func TestMachineDashboardCapabilityAndRemoteRequest(t *testing.T) {
	now := time.Date(2026, 7, 1, 1, 40, 0, 0, time.UTC)
	manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: manager, clients: map[string]*clientState{}}
	client := newElasticReadyClient()
	client.machineID = "mac_dashboard"
	client.workspaceID = "wsp_dashboard"
	hub.clients[client.machineID] = client
	require.True(t, hub.MachineDashboardSupported(client.machineID))

	descriptor, ok := defaultRemoteEndpointRegistry().Lookup(remoteEndpointMachineDashboard, remoteEndpointMachineDashboardVersion)
	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, remoteEndpointMachineDashboard, assignment.EndpointID)
		var payload machineDashboardRemoteRequestPayload
		require.NoError(t, json.Unmarshal(assignment.Payload, &payload))
		require.Equal(t, "wsp_dashboard", payload.WorkspaceID)
		require.Equal(t, "mac_dashboard", payload.MachineID)
		require.Equal(t, "6h", payload.Window)
		require.Equal(t, []string{"agt_a", "agt_b"}, payload.AgentIDs)
		manager.CompleteRemoteRequest(RemoteRequestResult[json.RawMessage]{
			Type:      sessionLinkRemoteRequestResultType,
			RequestID: assignment.RequestID,
			Success:   true,
			Outcome:   "ok",
			Payload: mustMarshalRemoteWorkspaceRawTest(t, machineDashboardRemoteRequestResponse{
				MachineID:    "mac_dashboard",
				WorkspaceID:  "wsp_dashboard",
				Window:       "6h",
				AgentCount:   2,
				ActiveAgents: 1,
				MessageSummary: runtimecontrol.MachineDashboardMessageSummary{
					Received: 7,
					Sent:     5,
					Total:    12,
				},
				TokenSources: runtimecontrol.MachineTokenSources{
					Available: true,
					Trend: runtimecontrol.MachineTokenSourceTrend{
						Series: []runtimecontrol.MachineTokenSourceTrendSeries{
							{Key: "source_0", Label: "grep", SourceKind: "command"},
						},
						Data: []runtimecontrol.MachineTokenSourceTrendPoint{
							{Timestamp: 1, Values: map[string]uint64{"source_0": 12}},
						},
					},
				},
			}),
		})
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	dashboard, err := hub.QueryMachineDashboard(ctx, client.machineID, runtimecontrol.MachineDashboardRequest{
		WorkspaceID: "wsp_dashboard",
		Window:      runtimecontrol.MachineDashboardWindow6Hours,
		AgentIDs:    []string{"agt_a", "agt_b"},
	})
	require.NoError(t, err)
	require.Equal(t, "mac_dashboard", dashboard.MachineID)
	require.Equal(t, "wsp_dashboard", dashboard.WorkspaceID)
	require.Equal(t, uint64(12), dashboard.MessageSummary.Total)
	// tokenSources is part of the remote dashboard contract.
	require.True(t, dashboard.TokenSources.Available)
	require.Len(t, dashboard.TokenSources.Trend.Series, 1)
	require.Equal(t, "source_0", dashboard.TokenSources.Trend.Series[0].Key)
	require.Equal(t, "grep", dashboard.TokenSources.Trend.Series[0].Label)
	require.Equal(t, "command", dashboard.TokenSources.Trend.Series[0].SourceKind)
}

func TestMachineDashboardCapabilityMissReturnsUnsupported(t *testing.T) {
	now := time.Date(2026, 7, 1, 1, 50, 0, 0, time.UTC)
	manager := newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: manager, clients: map[string]*clientState{}}
	client := newElasticReadyClient()
	delete(client.supportedEndpointVersions, remoteEndpointMachineDashboard)
	hub.clients[client.machineID] = client

	require.False(t, hub.MachineDashboardSupported(client.machineID))
	_, err := hub.QueryMachineDashboard(context.Background(), client.machineID, runtimecontrol.MachineDashboardRequest{WorkspaceID: client.workspaceID, Window: runtimecontrol.MachineDashboardWindow1Day})
	require.True(t, errors.Is(err, runtimecontrol.ErrMachineDashboardUnsupported), "got %v", err)
}

func TestTrajectoryHistoryRemoteRequestCapabilityMismatchFallsBack(t *testing.T) {
	serverConn, desktopConn := newDaemonProtocolTestPair(t)
	client := newWorkspaceListRemoteReadyClient(serverConn)
	delete(client.supportedEndpointVersions, remoteEndpointRuntimeTrajectoryHistory)
	now := time.Date(2026, 7, 1, 1, 20, 0, 0, time.UTC)
	hub := &Hub{remoteEndpoints: defaultRemoteEndpointRegistry(), elasticRequestLinks: newTestElasticRequestLinkManager(t, &now, ElasticRequestLinkManagerConfig{})}

	go func() {
		var command agentTrajectoryHistoryCommand
		require.NoError(t, desktopConn.ReadJSON(&command))
		require.Equal(t, sessionLinkCommandAgentTrajectoryHistory, command.Type)
		require.Equal(t, "agt_1", command.AgentID)
		require.NoError(t, handleSessionLinkAgentTrajectoryHistoryResultCommand(nil, context.Background(), client, mustMarshalRemoteWorkspaceTest(t, agentTrajectoryHistoryResultMessage{
			Type:      sessionLinkCommandAgentTrajectoryHistoryResult,
			AgentID:   command.AgentID,
			RequestID: command.RequestID,
			Items:     []runtimecontrol.AgentTrajectoryItem{{AgentID: command.AgentID, Entry: json.RawMessage(`{"kind":"text"}`), Timestamp: 1}},
		})))
	}()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	items, err := hub.trajectoryHistory(ctx, client, "agt_1", 10, 0)
	require.NoError(t, err)
	require.Equal(t, []runtimecontrol.AgentTrajectoryItem{{AgentID: "agt_1", Entry: json.RawMessage(`{"kind":"text"}`), Timestamp: 1}}, items)
}

func TestRestoreRecoveryPointRemoteRequestDoesNotRetryAfterEndpointError(t *testing.T) {
	serverConn, desktopConn := newDaemonProtocolTestPair(t)
	client := newWorkspaceListRemoteReadyClient(serverConn)
	now := time.Date(2026, 7, 1, 1, 30, 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, remoteEndpointRuntimeRestoreRecoveryPoint, assignment.EndpointID)
		var payload runtimeRestoreRecoveryPointRemoteRequestPayload
		require.NoError(t, json.Unmarshal(assignment.Payload, &payload))
		require.Equal(t, "agt_1", payload.AgentID)
		require.Equal(t, "ref_1", payload.RecoveryRef)
		require.NoError(t, handleSessionLinkRemoteRequestResultCommand(hub, context.Background(), client, mustMarshalRemoteWorkspaceTest(t, 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.restoreAgentRecoveryPoint(ctx, client, runtimecontrol.RestoreAgentRecoveryPointRequest{AgentID: "agt_1", Runtime: "codex", Strategy: "turns", RecoveryRef: "ref_1"})
	require.ErrorContains(t, err, "endpoint_error")
}

func TestRuntimeRemoteRequestErrorsAreRedacted(t *testing.T) {
	_, err := runtimeTrajectoryHistoryFromRemoteResult(RemoteRequestResult[json.RawMessage]{
		Type:    sessionLinkRemoteRequestResultType,
		Success: false,
		Outcome: "endpoint_error",
		Error:   "/tmp/secret/trajectory failed",
	})
	require.Error(t, err)
	require.Equal(t, "endpoint_error", err.Error())
}