package server

import (
	"context"
	"errors"
	"log/slog"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/config"
	"github.com/tidbcloud/AgentTeam/internal/router"
)

func TestManagedWorkspaceSetForProfile(t *testing.T) {
	t.Parallel()

	allSet, err := managedWorkspaceSetForProfile(context.Background(), router.ProfileAll, "", nil)
	require.NoError(t, err)
	require.True(t, allSet.Has("wsp_any"))

	centralSet, err := managedWorkspaceSetForProfile(context.Background(), router.ProfileCentral, "", nil)
	require.NoError(t, err)
	require.False(t, centralSet.Has("wsp_any"))

	loader := &fakeWorkspacePlacementLoader{activeWorkspaceIDs: map[string]bool{"wsp_1": true}}
	workspaceSet, err := managedWorkspaceSetForProfile(context.Background(), router.ProfileWorkspace, " workspace-default ", loader)
	require.NoError(t, err)
	require.True(t, workspaceSet.Has("wsp_1"))
	require.Equal(t, "workspace-default", loader.groupID)
	require.False(t, workspaceSet.Has("wsp_2"))
}

func TestManagedWorkspaceSetForProfileRequiresWorkspaceLoader(t *testing.T) {
	t.Parallel()

	_, err := managedWorkspaceSetForProfile(context.Background(), router.ProfileWorkspace, "workspace-default", &fakeWorkspacePlacementLoader{err: errors.New("boom")})
	require.NoError(t, err)

	_, err = managedWorkspaceSetForProfile(context.Background(), router.ProfileWorkspace, "workspace-default", nil)
	require.ErrorContains(t, err, "initialize managed workspace set")
	require.ErrorContains(t, err, "workspace placement checker is not configured")
}

func TestReconcileMachineConnectionsOnStartupByProfile(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	logger := slog.Default()

	centralAgent := &fakeStartupAgentStore{}
	require.NoError(t, reconcileMachineConnectionsOnStartup(ctx, config.Config{Server: config.ServerConfig{Profile: config.ServerProfileCentral}}, centralAgent, nil, logger))
	require.Zero(t, centralAgent.globalCalls)

	allAgent := &fakeStartupAgentStore{globalUpdated: 2}
	require.NoError(t, reconcileMachineConnectionsOnStartup(ctx, config.Config{Server: config.ServerConfig{Profile: config.ServerProfileAll}}, allAgent, nil, logger))
	require.Equal(t, 1, allAgent.globalCalls)

	workspaceAgent := &fakeStartupAgentStore{scopedUpdated: 1}
	workspaceLoader := &fakeStartupWorkspaceStore{workspaceIDs: []string{"wsp_1", "wsp_2"}}
	require.NoError(t, reconcileMachineConnectionsOnStartup(ctx, config.Config{Server: config.ServerConfig{
		Profile:          config.ServerProfileWorkspace,
		WorkspaceGroupID: " workspace-default ",
	}}, workspaceAgent, workspaceLoader, logger))
	require.Equal(t, "workspace-default", workspaceLoader.groupID)
	require.Equal(t, []string{"wsp_1", "wsp_2"}, workspaceAgent.scopedWorkspaceIDs)
	require.Zero(t, workspaceAgent.globalCalls)
}

type fakeWorkspacePlacementLoader struct {
	activeWorkspaceIDs map[string]bool
	groupID            string
	err                error
}

type fakeStartupAgentStore struct {
	globalCalls        int
	globalUpdated      int64
	scopedWorkspaceIDs []string
	scopedUpdated      int64
}

func (store *fakeStartupAgentStore) MarkMachinesDisconnectedOnStartup(context.Context) (int64, error) {
	store.globalCalls++
	return store.globalUpdated, nil
}

func (store *fakeStartupAgentStore) MarkMachinesDisconnectedOnStartupForWorkspaces(_ context.Context, workspaceIDs []string) (int64, error) {
	store.scopedWorkspaceIDs = append([]string(nil), workspaceIDs...)
	return store.scopedUpdated, nil
}

type fakeStartupWorkspaceStore struct {
	groupID      string
	workspaceIDs []string
}

func (store *fakeStartupWorkspaceStore) ListRoutableWorkspaceIDsForGroup(_ context.Context, groupID string) ([]string, error) {
	store.groupID = groupID
	return append([]string(nil), store.workspaceIDs...), nil
}

func (loader *fakeWorkspacePlacementLoader) ListActiveWorkspaceIDsForGroup(_ context.Context, groupID string) ([]string, error) {
	loader.groupID = groupID
	if loader.err != nil {
		return nil, loader.err
	}
	return nil, nil
}

func (loader *fakeWorkspacePlacementLoader) IsWorkspaceRoutableToGroup(_ context.Context, workspaceID string, groupID string) (bool, error) {
	loader.groupID = groupID
	if loader.err != nil {
		return false, loader.err
	}
	return loader.activeWorkspaceIDs[workspaceID], nil
}