package daemon

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

type extensionCatalogSyncStore struct {
	recordingStartStore
	agents []agent.MachineAgent
}

func (store *extensionCatalogSyncStore) ListAgentsForMachine(
	context.Context,
	string,
) ([]agent.MachineAgent, error) {
	return append([]agent.MachineAgent(nil), store.agents...), nil
}

func TestSyncMachineCatalogWritesProductionExtensionHint(t *testing.T) {
	serverConn, desktopConn := newDaemonProtocolTestPair(t)
	store := &extensionCatalogSyncStore{agents: []agent.MachineAgent{{
		ID: "agt_1", Name: "agent-1", Runtime: "codex", Status: "idle",
	}}}
	var readWorkspaceID string
	var readAgentID string
	hub := &Hub{
		store: store,
		hintReader: func(
			_ context.Context,
			workspaceID string,
			agentID string,
		) (*ExtensionCatalogHint, error) {
			readWorkspaceID = workspaceID
			readAgentID = agentID
			return &ExtensionCatalogHint{Version: 7}, nil
		},
		clients: map[string]*clientState{
			"mac_1": {
				machineID: "mac_1", workspaceID: "wsp_1", conn: serverConn,
				capabilities: map[string]struct{}{},
			},
		},
	}

	require.NoError(t, hub.SyncMachineCatalog(context.Background(), "mac_1"))

	var response machineAgentListResponseCommand
	require.NoError(t, desktopConn.ReadJSON(&response))
	require.Equal(t, sessionLinkCommandMachineAgentListResponse, response.Type)
	require.Equal(t, "wsp_1", readWorkspaceID)
	require.Equal(t, "agt_1", readAgentID)
	require.Len(t, response.Agents, 1)
	require.Equal(t, &ExtensionCatalogHint{Version: 7}, response.Agents[0].ExtensionCatalog)
}

func TestSyncMachineCatalogHintFailureSendsNothing(t *testing.T) {
	serverConn, desktopConn := newDaemonProtocolTestPair(t)
	store := &extensionCatalogSyncStore{agents: []agent.MachineAgent{{
		ID: "agt_1", Name: "agent-1", Runtime: "codex", Status: "idle",
	}}}
	hub := &Hub{
		store: store,
		hintReader: func(
			context.Context,
			string,
			string,
		) (*ExtensionCatalogHint, error) {
			return nil, errors.New("catalog row corrupt")
		},
		clients: map[string]*clientState{
			"mac_1": {
				machineID: "mac_1", workspaceID: "wsp_1", conn: serverConn,
				capabilities: map[string]struct{}{},
			},
		},
	}

	err := hub.SyncMachineCatalog(context.Background(), "mac_1")
	require.ErrorContains(t, err, "read extension catalog hint")

	require.NoError(t, desktopConn.SetReadDeadline(time.Now().Add(50*time.Millisecond)))
	var response machineAgentListResponseCommand
	readErr := desktopConn.ReadJSON(&response)
	require.Error(t, readErr)
	require.ErrorContains(t, readErr, "i/o timeout")
}