package devsearch

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

type fakeSearchReadinessStore struct {
	workspaceIDs  []string
	states        map[string]message.ScopeSearchReadiness
	backfillCalls map[string]int
	loadErr       map[string]error
	finalizeErr   map[string]error
}

func (store *fakeSearchReadinessStore) ListWorkspaceIDs(context.Context) ([]string, error) {
	return append([]string(nil), store.workspaceIDs...), nil
}

func (store *fakeSearchReadinessStore) LoadWorkspaceGlobalReadiness(_ context.Context, workspaceID string) (message.ScopeSearchReadiness, error) {
	if err := store.loadErr[workspaceID]; err != nil {
		return message.ScopeSearchReadiness{}, err
	}
	state, ok := store.states[workspaceID]
	if !ok {
		return message.ScopeSearchReadiness{}, message.ErrNotFound
	}
	return state, nil
}

func (store *fakeSearchReadinessStore) BackfillWorkspaceGlobal(_ context.Context, workspaceID string, batchSize int, updatedAt time.Time) (message.ScopeSearchBackfillResult, error) {
	_ = batchSize
	state, ok := store.states[workspaceID]
	if !ok {
		state = message.ScopeSearchReadiness{WorkspaceID: workspaceID}
	}
	if store.backfillCalls == nil {
		store.backfillCalls = make(map[string]int)
	}
	store.backfillCalls[workspaceID]++
	call := store.backfillCalls[workspaceID]
	state.WorkspaceID = workspaceID
	state.MigrationStatus = "complete"
	state.UpdatedAt = updatedAt
	if call == 1 {
		state.BackfillStatus = "running"
		state.FinalizeStatus = "pending"
		state.BackfillCursorGlobalSeq = 1
		store.states[workspaceID] = state
		return message.ScopeSearchBackfillResult{ProcessedRows: 1, State: state}, nil
	}
	state.BackfillStatus = "complete"
	state.FinalizeStatus = "pending"
	state.BackfillCursorGlobalSeq = 2
	store.states[workspaceID] = state
	return message.ScopeSearchBackfillResult{ProcessedRows: 1, State: state}, nil
}

func (store *fakeSearchReadinessStore) FinalizeWorkspaceGlobal(_ context.Context, workspaceID string, updatedAt time.Time) (message.ScopeSearchReadiness, error) {
	if err := store.finalizeErr[workspaceID]; err != nil {
		state := store.states[workspaceID]
		return state, err
	}
	state := store.states[workspaceID]
	state.WorkspaceID = workspaceID
	state.MigrationStatus = "complete"
	state.BackfillStatus = "complete"
	state.FinalizeStatus = "complete"
	state.UpdatedAt = updatedAt
	store.states[workspaceID] = state
	return state, nil
}

func (store *fakeSearchReadinessStore) PromoteWorkspaceGlobal(_ context.Context, workspaceID string, updatedAt time.Time) (message.ScopeSearchReadiness, error) {
	state := store.states[workspaceID]
	state.WorkspaceID = workspaceID
	state.MigrationStatus = "complete"
	state.BackfillStatus = "complete"
	state.FinalizeStatus = "complete"
	state.ReadReady = true
	state.UpdatedAt = updatedAt
	store.states[workspaceID] = state
	return state, nil
}

func TestEnsureAllWorkspacesSearchReadyMixesFastPathAndWork(t *testing.T) {
	store := &fakeSearchReadinessStore{
		workspaceIDs: []string{"wsp_fast", "wsp_slow"},
		states: map[string]message.ScopeSearchReadiness{
			"wsp_fast": {
				WorkspaceID:     "wsp_fast",
				MigrationStatus: "complete",
				BackfillStatus:  "complete",
				FinalizeStatus:  "complete",
				ReadReady:       true,
			},
		},
		finalizeErr: map[string]error{},
	}

	result, err := EnsureAllWorkspacesSearchReady(context.Background(), store, 1, nil)
	require.NoError(t, err)
	require.True(t, result.Ready)
	require.False(t, result.FastPath)
	require.Equal(t, 2, result.WorkspaceCount)
	require.Equal(t, 2, result.ProcessedCount)
	require.Len(t, result.Workspaces, 2)
	require.True(t, result.Workspaces[0].FastPath)
	require.Equal(t, "wsp_fast", result.Workspaces[0].WorkspaceID)
	require.Equal(t, "wsp_slow", result.Workspaces[1].WorkspaceID)
	require.True(t, result.Workspaces[1].ReadReady)
	require.Equal(t, 2, store.backfillCalls["wsp_slow"])
}

func TestEnsureAllWorkspacesSearchReadyReturnsPartialResultOnFailure(t *testing.T) {
	store := &fakeSearchReadinessStore{
		workspaceIDs: []string{"wsp_fail"},
		states:       map[string]message.ScopeSearchReadiness{},
		finalizeErr: map[string]error{
			"wsp_fail": fmt.Errorf("boom"),
		},
	}

	result, err := EnsureAllWorkspacesSearchReady(context.Background(), store, 2, nil)
	require.Error(t, err)
	require.False(t, result.Ready)
	require.Equal(t, 1, result.ProcessedCount)
	require.Equal(t, "finalize", result.CurrentPhase)
	require.Contains(t, result.Error, "boom")
	require.Len(t, result.Workspaces, 1)
	require.Equal(t, "wsp_fail", result.Workspaces[0].WorkspaceID)
	require.Equal(t, "complete", result.Workspaces[0].BackfillStatus)
}

func TestEnsureAllWorkspacesSearchReadyReportsLoadFailurePhase(t *testing.T) {
	store := &fakeSearchReadinessStore{
		workspaceIDs: []string{"wsp_fail"},
		states:       map[string]message.ScopeSearchReadiness{},
		loadErr: map[string]error{
			"wsp_fail": fmt.Errorf("load boom"),
		},
		finalizeErr: map[string]error{},
	}

	result, err := EnsureAllWorkspacesSearchReady(context.Background(), store, 2, nil)
	require.Error(t, err)
	require.False(t, result.Ready)
	require.Equal(t, 1, result.ProcessedCount)
	require.Equal(t, "load", result.CurrentPhase)
	require.Contains(t, result.Error, "load boom")
	require.Len(t, result.Workspaces, 1)
	require.Equal(t, "wsp_fail", result.Workspaces[0].WorkspaceID)
}