package kb

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/store/content"
	storekb "github.com/tidbcloud/AgentTeam/internal/store/kb"
)

type mutableKBSourceAuthorityReader struct {
	states map[string]string
	err    error
}

func (reader *mutableKBSourceAuthorityReader) ListAuthoritiesForBindings(_ context.Context, workspaceID, domain string, bindingKeys []string) (map[string]content.Authority, error) {
	if reader.err != nil {
		return nil, reader.err
	}
	out := make(map[string]content.Authority)
	for _, bindingKey := range bindingKeys {
		if state, ok := reader.states[bindingKey]; ok {
			out[bindingKey] = content.Authority{WorkspaceID: workspaceID, Domain: domain, BindingKey: bindingKey, State: state}
		}
	}
	return out, nil
}

func TestNativeKBSourceCancelsQueuedMirrorBeforeAGSCalls(t *testing.T) {
	ctx := t.Context()
	syncer := &fakeSourceMirrorSyncer{}
	wikiStore := &fakeAGSRepositoryStore{}
	svc := newSourceMirrorServiceForTest(t, syncer, wikiStore)
	seedSourceMirror(t, svc, true)
	svc.contentAuthority = &mutableKBSourceAuthorityReader{states: map[string]string{
		"src_1": content.AuthorityStateNativeActive,
	}}
	_, err := svc.store.CreateJob(ctx, storekb.Job{
		ID: "job_native", WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", SourceID: "src_1",
		JobType: KBJobTypeSyncSourceMirror, Status: KBJobStatusQueued,
		Progress: `{"source_id":"src_1","ref":"main"}`, CreatedAt: svc.now(), UpdatedAt: svc.now(),
	})
	require.NoError(t, err)

	result, err := svc.RunPendingSourceMirrorJobs(ctx, SourceMirrorRunOptions{Limit: 1, LeaseOwner: "worker"})
	require.NoError(t, err)
	require.Equal(t, SourceMirrorRunResult{Claimed: 1, Succeeded: 1}, result)
	require.Empty(t, syncer.request.Source.ID)
	require.False(t, wikiStore.reconciled)
	job, err := svc.store.GetJob(ctx, "wsp_1", "job_native")
	require.NoError(t, err)
	require.Equal(t, KBJobStatusCancelled, job.Status)
	require.Contains(t, job.Progress, "native authority is active")
}

func TestNativeKBSourceRejectsNewLegacyWikiAndCancelsQueuedGeneration(t *testing.T) {
	generator := &fakeKBWikiGenerator{output: KBWikiGenerationOutput{PublishedWikiCommit: "must-not-publish"}}
	svc, base, revision := seedDirectWikiGenerationSetup(t, generator)
	authority := &mutableKBSourceAuthorityReader{states: map[string]string{}}
	svc.contentAuthority = authority

	_, err := svc.RequestDirectWikiGeneration(t.Context(), "ws_1", "owner", base.ID, "src_1", KBWikiGenerationMutation{
		SourceRevision: revision.ID, IdempotencyKey: "queued-before-cutover",
	})
	require.NoError(t, err)
	authority.states["src_1"] = content.AuthorityStateNativeActive

	_, err = svc.RequestDirectWikiGeneration(t.Context(), "ws_1", "owner", base.ID, "src_1", KBWikiGenerationMutation{
		SourceRevision: revision.ID, IdempotencyKey: "rejected-after-cutover",
	})
	require.Equal(t, "kb-source-native-active", Code(err))
	claimed, err := svc.RunPendingWikiGenerationJobs(t.Context(), WikiGenerationRunOptions{
		Limit: 1, LeaseOwner: "wiki-worker", LeaseDuration: time.Hour,
	})
	require.NoError(t, err)
	require.Equal(t, 1, claimed)
	require.Empty(t, generator.inputs)
	setup, err := svc.GetKnowledgeBaseSetup(t.Context(), "ws_1", "owner", base.ID)
	require.NoError(t, err)
	require.Equal(t, storekb.GenerationRequestStatusStale, setup.Sources[0].WikiState)
}

func TestKBSourceWriteRouteFailsClosedAfterNativeCutover(t *testing.T) {
	svc := &KBService{contentAuthority: &mutableKBSourceAuthorityReader{states: map[string]string{
		"src_1": content.AuthorityStateNativeActive,
	}}}
	route, err := svc.beginKBSourceWrite(t.Context(), "ws_1", "src_1")
	require.NoError(t, err)
	require.True(t, route.native)
	require.Error(t, svc.ensureKBSourceWriteCurrent(t.Context(), "ws_1", "src_1", kbSourceWriteRoute{generation: 1}))
}