package kb

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"

	"github.com/tidbcloud/AgentTeam/internal/contentpayload"
	contentservice "github.com/tidbcloud/AgentTeam/internal/service/content"
	"github.com/tidbcloud/AgentTeam/internal/store/content"
	storekb "github.com/tidbcloud/AgentTeam/internal/store/kb"
	reposstore "github.com/tidbcloud/AgentTeam/internal/store/resource/repos"
)

type recordingKBSourceMigrationAuthorizer struct {
	allowed atomic.Bool
}

func (authorizer *recordingKBSourceMigrationAuthorizer) CanMigrateKBSource(context.Context, string, string, string, string) (bool, error) {
	return authorizer.allowed.Load(), nil
}

type recordingKBSourceMigrationExporter struct {
	result KBSourceContentExport
	calls  atomic.Int64
}

func (exporter *recordingKBSourceMigrationExporter) ExportExactKBSource(context.Context, KBSourceContentExportRequest) (KBSourceContentExport, error) {
	exporter.calls.Add(1)
	return exporter.result, nil
}

type failOnceKBSourceMappingStore struct {
	KBSourceContentMigrationStore
	fail atomic.Bool
}

func (store *failOnceKBSourceMappingStore) ProjectSourceContentMapping(ctx context.Context, params storekb.ProjectSourceContentMappingParams) (storekb.SourceContentMapping, error) {
	if store.fail.CompareAndSwap(true, false) {
		return storekb.SourceContentMapping{}, errors.New("injected source projection failure")
	}
	return store.KBSourceContentMigrationStore.ProjectSourceContentMapping(ctx, params)
}

func TestKBSourceContentMigrationRetryReusesContentOperation(t *testing.T) {
	fixture := newKBSourceContentMigrationFixture(t, true)
	failing := &failOnceKBSourceMappingStore{KBSourceContentMigrationStore: fixture.store}
	failing.fail.Store(true)
	fixture.service = fixture.newService(t, true, failing)

	_, err := fixture.service.RunExact(t.Context(), fixture.request)
	require.ErrorContains(t, err, "injected source projection failure")
	require.Equal(t, int64(1), migrationRowCount(t, fixture.db, &content.ContentVersion{}))
	require.Equal(t, int64(1), migrationRowCount(t, fixture.db, &content.VersionOperation{}))
	require.Equal(t, int64(0), migrationRowCount(t, fixture.db, &storekb.SourceContentMapping{}))
	require.Equal(t, int64(0), migrationRowCount(t, fixture.db, &storekb.SourceSearchDocument{}))

	second, err := fixture.service.RunExact(t.Context(), fixture.request)
	require.NoError(t, err)
	require.True(t, second.Replayed)
	require.Equal(t, int64(2), fixture.exporter.calls.Load(), "stateless retry repeats only the bounded source read")
	require.Equal(t, int64(1), migrationRowCount(t, fixture.db, &content.ContentVersion{}))
	require.Equal(t, int64(1), migrationRowCount(t, fixture.db, &storekb.SourceContentMapping{}))
	require.Equal(t, int64(1), migrationRowCount(t, fixture.db, &storekb.SourceSearchDocument{}))

	third, err := fixture.service.RunExact(t.Context(), fixture.request)
	require.NoError(t, err)
	require.True(t, third.Replayed)
	require.Equal(t, second.ContentVersionID, third.ContentVersionID)
	require.Equal(t, int64(2), fixture.exporter.calls.Load(), "verified mapping replay skips AGS")
	require.Equal(t, int64(0), migrationRowCount(t, fixture.db, &content.Authority{}), "import does not switch authority")
}

func TestBuildKBSourceSearchProjectionAppliesPerDocumentAndTotalBounds(t *testing.T) {
	files := make([]KBSourceContentExportFile, 0, 132)
	files = append(files, KBSourceContentExportFile{Path: "binary.dat", Data: []byte{'a', 0, 'b'}})
	for index := range 131 {
		files = append(files, KBSourceContentExportFile{
			Path: fmt.Sprintf("docs/%03d.md", index), Data: []byte(strings.Repeat("x", storekb.SourceSearchProjectionMaxDocumentBytes+17)),
		})
	}
	projection := buildKBSourceSearchProjection(files)
	require.True(t, projection.Truncated)
	require.NotEmpty(t, projection.Digest)
	require.LessOrEqual(t, len(projection.Documents), storekb.SourceSearchProjectionMaxDocuments)
	require.LessOrEqual(t, projection.TextBytes, storekb.SourceSearchProjectionMaxTextBytes)
	for _, document := range projection.Documents {
		require.LessOrEqual(t, len(document.SearchText), storekb.SourceSearchProjectionMaxDocumentBytes)
		require.NotEqual(t, "binary.dat", document.Path)
	}
}

func TestKBSourceContentMigrationRejectsUnsafeTargetsBeforePublication(t *testing.T) {
	tests := []struct {
		name      string
		configure func(*kbSourceMigrationFixture)
		wantRead  int64
	}{
		{name: "disabled", configure: func(f *kbSourceMigrationFixture) { f.service = f.newService(t, false, f.store) }},
		{name: "revoked", configure: func(f *kbSourceMigrationFixture) { f.authorizer.allowed.Store(false) }},
		{name: "live content ACL revoked", configure: func(f *kbSourceMigrationFixture) { f.workspace.allowed["wsp_1\x00owner"] = false }, wantRead: 1},
		{name: "stale revision", configure: func(f *kbSourceMigrationFixture) { f.request.SourceRevisionID = "rev_stale" }},
		{name: "cross workspace", configure: func(f *kbSourceMigrationFixture) { f.request.WorkspaceID = "wsp_other" }},
		{name: "invalid digest", configure: func(f *kbSourceMigrationFixture) { f.request.ExpectedDigest = strings.Repeat("z", 64) }},
		{name: "digest mismatch", configure: func(f *kbSourceMigrationFixture) { f.request.ExpectedDigest = strings.Repeat("f", 64) }, wantRead: 1},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			fixture := newKBSourceContentMigrationFixture(t, true)
			test.configure(fixture)
			_, err := fixture.service.RunExact(t.Context(), fixture.request)
			require.Error(t, err)
			require.Equal(t, test.wantRead, fixture.exporter.calls.Load())
			require.Equal(t, int64(0), migrationRowCount(t, fixture.db, &content.ContentVersion{}))
			require.Equal(t, int64(0), migrationRowCount(t, fixture.db, &storekb.SourceContentMapping{}))
		})
	}
}

func TestKBSourceContentMigrationMappingReplayVerifiesExactVersionNotHead(t *testing.T) {
	fixture := newKBSourceContentMigrationFixture(t, true)
	first, err := fixture.service.RunExact(t.Context(), fixture.request)
	require.NoError(t, err)

	content, err := fixture.content.GetContent(t.Context(), fixture.request.WorkspaceID, first.ContentID)
	require.NoError(t, err)
	laterFiles := []contentpayload.File{{Path: "later.txt", Size: 5, Content: strings.NewReader("later")}}
	_, err = fixture.content.CommitVersion(t.Context(), contentservice.CommitVersionRequest{
		WorkspaceID: fixture.request.WorkspaceID, ContentID: content.ID, ExpectedHeadVersionID: content.HeadVersionID,
		ExpectedGeneration: content.Generation, Actor: fixture.request.Actor, Model: contentservice.TreeModel,
		Checkpoint: true, Files: laterFiles, TreeLimits: fixture.request.Limits,
	})
	require.NoError(t, err)

	replayed, err := fixture.service.RunExact(t.Context(), fixture.request)
	require.NoError(t, err)
	require.True(t, replayed.Replayed)
	require.Equal(t, first.ContentVersionID, replayed.ContentVersionID)
	require.Equal(t, int64(1), fixture.exporter.calls.Load())
}

func TestKBSourceContentMigrationImportsHistoricalNonCurrentRevision(t *testing.T) {
	fixture := newKBSourceContentMigrationFixture(t, true)
	require.NoError(t, fixture.db.Model(&storekb.SourceRevision{}).Where("id = ?", fixture.request.SourceRevisionID).Update("is_current", false).Error)
	require.NoError(t, fixture.db.Create(&storekb.SourceRevision{
		ID: "rev_current", WorkspaceID: fixture.request.WorkspaceID, KnowledgeBaseID: fixture.request.KnowledgeBaseID,
		SourceID: fixture.request.SourceID, SourceCommitSHA: strings.Repeat("b", 40), IsCurrent: true,
		CreatedAt: time.Date(2026, 8, 21, 6, 0, 0, 0, time.UTC), UpdatedAt: time.Date(2026, 8, 21, 6, 0, 0, 0, time.UTC),
	}).Error)

	result, err := fixture.service.RunExact(t.Context(), fixture.request)
	require.NoError(t, err)
	require.NotEmpty(t, result.ContentVersionID)
	historical, err := fixture.store.GetSourceRevision(t.Context(), fixture.request.WorkspaceID, fixture.request.KnowledgeBaseID, fixture.request.SourceID, fixture.request.SourceRevisionID)
	require.NoError(t, err)
	require.Equal(t, result.ContentVersionID, historical.SourceContentVersionID)
	current, err := fixture.store.GetCurrentSourceRevision(t.Context(), fixture.request.WorkspaceID, fixture.request.KnowledgeBaseID, fixture.request.SourceID)
	require.NoError(t, err)
	require.Equal(t, "rev_current", current.ID)
	require.Empty(t, current.SourceContentVersionID)
}

func TestKBSourceContentMigrationRequiresOrderedHistoricalHead(t *testing.T) {
	fixture := newKBSourceContentMigrationFixture(t, true)
	first, err := fixture.service.RunExact(t.Context(), fixture.request)
	require.NoError(t, err)

	now := time.Date(2026, 8, 21, 6, 0, 0, 0, time.UTC)
	require.NoError(t, fixture.db.Create(&storekb.SourceRevision{
		ID: "rev_2", WorkspaceID: fixture.request.WorkspaceID, KnowledgeBaseID: fixture.request.KnowledgeBaseID,
		SourceID: fixture.request.SourceID, SourceCommitSHA: strings.Repeat("b", 40), CreatedAt: now, UpdatedAt: now,
	}).Error)
	fixture.request.SourceRevisionID = "rev_2"
	fixture.request.CommitSHA = strings.Repeat("b", 40)
	files := []KBSourceContentExportFile{{Path: "README.md", Data: []byte("next")}}
	digest, size := testKBSourceExportDigest(files)
	fixture.exporter.result = KBSourceContentExport{Files: files, RootDigest: digest, SizeBytes: size, SourceCommittedAt: now}

	_, err = fixture.service.RunExact(t.Context(), fixture.request)
	require.ErrorContains(t, err, "expected head changed")
	require.Equal(t, int64(1), migrationRowCount(t, fixture.db, &content.ContentVersion{}))

	fixture.request.ExpectedHeadVersionID = first.ContentVersionID
	second, err := fixture.service.RunExact(t.Context(), fixture.request)
	require.NoError(t, err)
	version, err := fixture.content.GetVersion(t.Context(), fixture.request.WorkspaceID, second.ContentID, second.ContentVersionID)
	require.NoError(t, err)
	require.Equal(t, first.ContentVersionID, version.ParentVersionID)
}

type kbSourceMigrationFixture struct {
	db         *gorm.DB
	store      *storekb.Store
	repos      *reposstore.Store
	content    *contentservice.Service
	authorized *contentservice.AuthorizedService
	authorizer *recordingKBSourceMigrationAuthorizer
	workspace  *kbContentWorkspaceAccess
	exporter   *recordingKBSourceMigrationExporter
	request    KBSourceContentMigrationRequest
	service    *KBSourceContentMigrationService
	ids        atomic.Int64
}

func newKBSourceContentMigrationFixture(t *testing.T, enabled bool) *kbSourceMigrationFixture {
	t.Helper()
	db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(
		&storekb.KnowledgeBase{}, &storekb.KnowledgeBaseGrant{}, &storekb.KnowledgeBaseSourceRepositoryAccessOverride{}, &storekb.Source{}, &storekb.SourceMirror{}, &storekb.SourceRevision{}, &storekb.SourceContentMapping{}, &storekb.SourceSearchDocument{}, &storekb.Job{},
		&reposstore.ResourceRepository{}, &content.Authority{}, &content.VersionedContent{}, &content.ContentVersion{},
		&content.VersionOperation{}, &content.VersionPath{}, &content.VersionPin{}, &content.Outbox{}, &content.StagedPayloadGC{},
	))
	sqlDB, err := db.DB()
	require.NoError(t, err)
	t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
	now := time.Date(2026, 8, 21, 5, 0, 0, 0, time.UTC)
	store := storekb.New(db)
	repos := reposstore.New(db)
	require.NoError(t, db.Create(&storekb.KnowledgeBase{
		ID: "kb_1", WorkspaceID: "wsp_1", Name: "Runbooks", Description: "ops", IsActive: true,
		PublicationState: "official", DefaultWorkspaceMemberRole: "view", DefaultRepositoryAccessProfile: storekb.KnowledgeBaseRepositoryAccessProfileRead,
		CreatedBy: "owner", UpdatedBy: "owner", CreatedAt: now, UpdatedAt: now,
	}).Error)
	require.NoError(t, db.Create(&storekb.Source{ID: "src_1", WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", SourceType: "codebase", Status: "ready", Name: "AgentTeam", DefaultRef: "main", CreatedBy: "owner", UpdatedBy: "owner", CreatedAt: now, UpdatedAt: now}).Error)
	require.NoError(t, db.Create(&storekb.SourceRevision{ID: "rev_1", WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", SourceID: "src_1", SourceCommitSHA: strings.Repeat("a", 40), IsCurrent: true, CreatedAt: now, UpdatedAt: now}).Error)
	_, err = repos.Upsert(t.Context(), reposstore.ResourceRepository{
		ID: "rr_1", WorkspaceID: "wsp_1", ResourceType: reposstore.ResourceTypeKBSource, ResourceID: "src_1",
		RepoRole: reposstore.RepoRolePrimary, OwnerKind: "workspace", RepoOwner: "org", RepoName: "agentteam", RepoFullName: "org/agentteam",
		AGSRepoID: "42", DefaultRef: "main", CreatedBy: "owner", UpdatedBy: "owner", CreatedAt: now, UpdatedAt: now,
	})
	require.NoError(t, err)
	tree, err := contentpayload.NewFilesystemTreeBackend(t.TempDir())
	require.NoError(t, err)
	t.Cleanup(func() { require.NoError(t, tree.Close()) })
	fixture := &kbSourceMigrationFixture{db: db, store: store, repos: repos}
	newID := func(prefix string) (string, error) { return fmt.Sprintf("%s_%d", prefix, fixture.ids.Add(1)), nil }
	fixture.content, err = contentservice.NewTree(content.New(db), tree, newID)
	require.NoError(t, err)
	fixture.workspace = &kbContentWorkspaceAccess{
		allowed: map[string]bool{"wsp_1\x00owner": true},
		admins:  map[string]bool{},
	}
	policy, err := NewKBSourceContentAccessPolicy(fixture.workspace, fixture.workspace, store)
	require.NoError(t, err)
	fixture.authorized, err = contentservice.NewAuthorizedService(fixture.content, policy)
	require.NoError(t, err)
	fixture.authorizer = &recordingKBSourceMigrationAuthorizer{}
	fixture.authorizer.allowed.Store(true)
	files := []KBSourceContentExportFile{{Path: "README.md", Data: []byte("current")}}
	digest, size := testKBSourceExportDigest(files)
	fixture.exporter = &recordingKBSourceMigrationExporter{result: KBSourceContentExport{
		Files: files, RootDigest: digest, SizeBytes: size, SourceCommittedAt: time.Date(2026, 8, 20, 1, 2, 3, 0, time.UTC),
	}}
	fixture.request = KBSourceContentMigrationRequest{
		WorkspaceID: "wsp_1", KnowledgeBaseID: "kb_1", SourceID: "src_1", SourceRevisionID: "rev_1",
		RepositoryRole: reposstore.RepoRolePrimary, RepositoryID: "42", CanonicalBranch: "main",
		CommitSHA: strings.Repeat("a", 40), Actor: "owner", Limits: testKBSourceMigrationLimits(),
	}
	fixture.service = fixture.newService(t, enabled, store)
	return fixture
}

func (fixture *kbSourceMigrationFixture) newService(t *testing.T, enabled bool, store KBSourceContentMigrationStore) *KBSourceContentMigrationService {
	t.Helper()
	service, err := NewKBSourceContentMigrationService(KBSourceContentMigrationOptions{
		Enabled: enabled, Store: store, Repositories: fixture.repos, Authorizer: fixture.authorizer,
		Content: fixture.authorized, Exporter: fixture.exporter,
		NewID: func(prefix string) (string, error) { return fmt.Sprintf("%s_%d", prefix, fixture.ids.Add(1)), nil },
		Now:   func() time.Time { return time.Date(2026, 8, 21, 5, 30, int(fixture.ids.Load()), 0, time.UTC) },
	})
	require.NoError(t, err)
	return service
}

func testKBSourceExportDigest(files []KBSourceContentExportFile) (string, int64) {
	manifest := make([]contentpayload.ManifestEntry, 0, len(files))
	var total int64
	for _, file := range files {
		sum := sha256.Sum256(file.Data)
		manifest = append(manifest, contentpayload.ManifestEntry{Path: file.Path, Digest: contentpayload.Digest(hex.EncodeToString(sum[:])), SizeBytes: int64(len(file.Data)), Executable: file.Executable})
		total += int64(len(file.Data))
	}
	encoded, _ := json.Marshal(manifest)
	sum := sha256.Sum256(encoded)
	return hex.EncodeToString(sum[:]), total
}

func migrationRowCount(t *testing.T, db *gorm.DB, model any) int64 {
	t.Helper()
	var count int64
	require.NoError(t, db.Model(model).Count(&count).Error)
	return count
}