package memory

import (
	"context"
	"errors"
	"fmt"
	"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/service/content/migration"
	"github.com/tidbcloud/AgentTeam/internal/service/resource/repos"
	contentstore "github.com/tidbcloud/AgentTeam/internal/store/content"
	migrationstore "github.com/tidbcloud/AgentTeam/internal/store/content/migration"
	memorystore "github.com/tidbcloud/AgentTeam/internal/store/resource/memory"
	reposstore "github.com/tidbcloud/AgentTeam/internal/store/resource/repos"
)

func TestMemoryContentMigrationPublishesExactBranchDetachedBackupAndReplays(t *testing.T) {
	root := migrationTestSHA(1)
	head := migrationTestSHA(2)
	detached := migrationTestSHA(3)
	rootAuthoredAt := time.Date(2025, 1, 1, 10, 0, 0, 0, time.UTC)
	headAuthoredAt := rootAuthoredAt.Add(time.Hour)
	detachedAuthoredAt := rootAuthoredAt.Add(30 * time.Minute)
	source := &fakeMemoryMigrationSource{
		head:    head,
		backups: []MemoryMigrationReference{{Kind: "backup", Label: "detached", CommitSHA: detached}},
		commits: map[string]MemoryMigrationCommit{
			root:     {CommitSHA: root, AuthorName: "Root Author", AuthorEmail: "root@example.test", AuthoredAt: rootAuthoredAt.Format(time.RFC3339Nano), Message: "root message"},
			head:     {CommitSHA: head, FirstParentSHA: root, AuthorName: "Head Author", AuthorEmail: "head@example.test", AuthoredAt: headAuthoredAt.Format(time.RFC3339Nano), Message: "head message"},
			detached: {CommitSHA: detached, FirstParentSHA: root, AuthorName: "Backup Author", AuthorEmail: "backup@example.test", AuthoredAt: detachedAuthoredAt.Format(time.RFC3339Nano), Message: "backup message"},
		},
		snapshots: map[string][]MemoryMigrationFile{
			root:     {{Path: "MEMORY.md", Data: []byte("root\n")}},
			head:     {{Path: "MEMORY.md", Data: []byte("head\n")}},
			detached: {{Path: "MEMORY.md", Data: []byte("detached\n")}},
		},
	}
	service, db := newMemoryMigrationServiceTest(t, source)

	result, err := service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{})
	require.NoError(t, err)
	require.Equal(t, 3, result.ImportedVersions)
	require.Equal(t, 1, result.ImportedPins)
	require.False(t, result.Replayed)
	require.EqualValues(t, 1, result.Generation)

	var content contentstore.VersionedContent
	require.NoError(t, db.First(&content, "id = ?", result.ContentID).Error)
	require.Equal(t, result.HeadVersionID, content.HeadVersionID)
	require.EqualValues(t, memoryContentMaxVersions, content.MaxRetainedVersions)
	require.Equal(t, memoryContentMaxStoredBytes, content.MaxStoredPayloadBytes)
	require.EqualValues(t, memoryContentMaxPins, content.MaxPins)

	mappings := migrationstore.New(db)
	rootMap, err := mappings.GetMapping(context.Background(), "repo_1", root)
	require.NoError(t, err)
	headMap, err := mappings.GetMapping(context.Background(), "repo_1", head)
	require.NoError(t, err)
	detachedMap, err := mappings.GetMapping(context.Background(), "repo_1", detached)
	require.NoError(t, err)
	require.Equal(t, headMap.ContentVersionID, content.HeadVersionID)
	var detachedVersion contentstore.ContentVersion
	require.NoError(t, db.First(&detachedVersion, "id = ?", detachedMap.ContentVersionID).Error)
	require.Equal(t, rootMap.ContentVersionID, detachedVersion.ParentVersionID)
	require.NotEqual(t, detachedVersion.ID, content.HeadVersionID)

	var pin contentstore.VersionPin
	require.NoError(t, db.First(&pin, "content_id = ? AND label = ?", content.ID, "detached").Error)
	require.Equal(t, detachedVersion.ID, pin.VersionID)
	var memory memorystore.Memory
	require.NoError(t, db.First(&memory, "id = ?", "mem_1").Error)
	require.Equal(t, content.ID, memory.ContentID)
	require.Equal(t, content.HeadVersionID, memory.ActiveContentVersionID)
	require.EqualValues(t, 1, memory.ContentGeneration)
	require.EqualValues(t, 3, memory.ContentHistorySequence)

	checkpoints, err := memorystore.New(db).ListMemoryContentCheckpoints(context.Background(), "wsp_1", "mem_1", 10, 0)
	require.NoError(t, err)
	require.Len(t, checkpoints, 3)
	require.Equal(t, headMap.ContentVersionID, checkpoints[0].ContentVersionID)
	require.Equal(t, "Head Author", checkpoints[0].SourceAuthorName)
	require.Equal(t, "head@example.test", checkpoints[0].SourceAuthorEmail)
	require.Equal(t, "head message", checkpoints[0].SourceMessage)
	require.NotNil(t, checkpoints[0].SourceAuthoredAt)
	require.True(t, checkpoints[0].SourceAuthoredAt.Equal(headAuthoredAt))
	require.True(t, checkpoints[0].CreatedAt.Equal(headAuthoredAt))
	require.NotContains(t, checkpoints[0].OperationID, head, "product operation IDs must not expose AGS commit OIDs")
	require.Len(t, checkpoints[0].OperationID, len("memory-import:")+64)

	var versionCount int64
	require.NoError(t, db.Model(&contentstore.ContentVersion{}).Where("content_id = ?", content.ID).Count(&versionCount).Error)
	replayed, err := service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{})
	require.NoError(t, err)
	require.True(t, replayed.Replayed)
	require.Zero(t, replayed.ImportedPins)
	var replayedVersionCount int64
	require.NoError(t, db.Model(&contentstore.ContentVersion{}).Where("content_id = ?", content.ID).Count(&replayedVersionCount).Error)
	require.Equal(t, versionCount, replayedVersionCount)
}

func TestMemoryContentMigrationHistoryOrderIgnoresRegressingAndEqualAuthorClocks(t *testing.T) {
	root := migrationTestSHA(31)
	middle := migrationTestSHA(32)
	head := migrationTestSHA(33)
	detached := migrationTestSHA(34)
	rootAuthoredAt := time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC)
	regressedAuthoredAt := rootAuthoredAt.Add(-time.Hour)
	source := &fakeMemoryMigrationSource{
		head:    head,
		backups: []MemoryMigrationReference{{Kind: "backup", Label: "detached", CommitSHA: detached}},
		commits: map[string]MemoryMigrationCommit{
			root:     {CommitSHA: root, AuthoredAt: rootAuthoredAt.Format(time.RFC3339Nano)},
			middle:   {CommitSHA: middle, FirstParentSHA: root, AuthoredAt: regressedAuthoredAt.Format(time.RFC3339Nano)},
			head:     {CommitSHA: head, FirstParentSHA: middle, AuthoredAt: regressedAuthoredAt.Format(time.RFC3339Nano)},
			detached: {CommitSHA: detached, FirstParentSHA: root, AuthoredAt: regressedAuthoredAt.Format(time.RFC3339Nano)},
		},
		snapshots: map[string][]MemoryMigrationFile{
			root: {{Path: "MEMORY.md", Data: []byte("root\n")}}, middle: {{Path: "MEMORY.md", Data: []byte("middle\n")}},
			head: {{Path: "MEMORY.md", Data: []byte("head\n")}}, detached: {{Path: "MEMORY.md", Data: []byte("detached\n")}},
		},
	}
	service, db := newMemoryMigrationServiceTest(t, source)
	_, err := service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{})
	require.NoError(t, err)

	mappings := migrationstore.New(db)
	versionIDs := make(map[string]string, 4)
	for _, sha := range []string{root, middle, head, detached} {
		mapping, err := mappings.GetMapping(context.Background(), "repo_1", sha)
		require.NoError(t, err)
		versionIDs[sha] = mapping.ContentVersionID
	}
	checkpoints, err := memorystore.New(db).ListMemoryContentCheckpoints(context.Background(), "wsp_1", "mem_1", 10, 0)
	require.NoError(t, err)
	require.Equal(t, []string{versionIDs[head], versionIDs[middle], versionIDs[root], versionIDs[detached]}, []string{
		checkpoints[0].ContentVersionID, checkpoints[1].ContentVersionID, checkpoints[2].ContentVersionID, checkpoints[3].ContentVersionID,
	})
	require.Equal(t, []uint64{4, 3, 2, 1}, []uint64{
		checkpoints[0].HistorySequence, checkpoints[1].HistorySequence, checkpoints[2].HistorySequence, checkpoints[3].HistorySequence,
	})
	require.True(t, checkpoints[0].CreatedAt.Equal(regressedAuthoredAt))
	require.True(t, checkpoints[1].CreatedAt.Equal(regressedAuthoredAt), "equal author clocks must not affect import lineage order")
	require.True(t, checkpoints[2].CreatedAt.Equal(rootAuthoredAt), "source timestamp remains preserved even when it regresses against lineage order")
}

func TestMemoryContentMigrationRejectsUnboundBranchBeforeAGSRead(t *testing.T) {
	root := migrationTestSHA(6)
	source := &fakeMemoryMigrationSource{
		head: root, commits: map[string]MemoryMigrationCommit{root: {CommitSHA: root}},
		snapshots: map[string][]MemoryMigrationFile{root: {{Path: "MEMORY.md", Data: []byte("root\n")}}},
	}
	service, db := newMemoryMigrationServiceTest(t, source)

	_, err := service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{Branch: "detached"})
	require.ErrorContains(t, err, "request branch does not match Memory binding")
	require.Empty(t, source.snapshotReads)
	var contents int64
	require.NoError(t, db.Model(&contentstore.VersionedContent{}).Count(&contents).Error)
	require.Zero(t, contents)
}

func TestMemoryContentMigrationRejectsPrimaryRepositoryChangeBeforeFinalization(t *testing.T) {
	root := migrationTestSHA(7)
	source := &fakeMemoryMigrationSource{
		head: root, commits: map[string]MemoryMigrationCommit{root: {CommitSHA: root}},
		snapshots: map[string][]MemoryMigrationFile{root: {{Path: "MEMORY.md", Data: []byte("root\n")}}},
	}
	service, db := newMemoryMigrationServiceTest(t, source)
	source.snapshotHook = func(read int, _ repos.RepositoryIdentity) {
		if read == 2 {
			require.NoError(t, db.Model(&reposstore.ResourceRepository{}).Where("id = ?", "rr_1").Updates(map[string]any{
				"ags_repo_id": "repo_2", "repo_owner": "other", "repo_name": "other-memory", "repo_full_name": "other/other-memory",
			}).Error)
		}
	}

	_, err := service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{})
	require.ErrorContains(t, err, "primary repository changed after preflight")
	var memory memorystore.Memory
	require.NoError(t, db.First(&memory, "id = ?", "mem_1").Error)
	require.Empty(t, memory.ContentID, "an old-repository import must never become the active Memory projection")
}

func TestMemoryContentMigrationResumesMappedProjectionWithoutRepublishing(t *testing.T) {
	root := migrationTestSHA(11)
	source := &fakeMemoryMigrationSource{
		head: root, commits: map[string]MemoryMigrationCommit{root: {CommitSHA: root}},
		snapshots: map[string][]MemoryMigrationFile{root: {{Path: "MEMORY.md", Data: []byte("root\n")}}},
	}
	service, db := newMemoryMigrationServiceTest(t, source)
	var fail atomic.Bool
	fail.Store(true)
	require.NoError(t, db.Callback().Create().Before("gorm:create").Register("test:fail_memory_import_projection", func(tx *gorm.DB) {
		if tx.Statement != nil && tx.Statement.Schema != nil && tx.Statement.Schema.Table == (memorystore.MemoryContentCheckpoint{}).TableName() && fail.CompareAndSwap(true, false) {
			tx.AddError(errors.New("injected memory import projection failure"))
		}
	}))

	_, err := service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{})
	require.ErrorContains(t, err, "injected memory import projection failure")
	var versionsAfterFailure int64
	require.NoError(t, db.Model(&contentstore.ContentVersion{}).Count(&versionsAfterFailure).Error)
	require.EqualValues(t, 1, versionsAfterFailure)

	result, err := service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{})
	require.NoError(t, err)
	require.True(t, result.Replayed)
	var versionsAfterRetry int64
	require.NoError(t, db.Model(&contentstore.ContentVersion{}).Count(&versionsAfterRetry).Error)
	require.Equal(t, versionsAfterFailure, versionsAfterRetry)
}

func TestMemoryContentMigrationFailsClosedWhenSnapshotChangesAfterPreflight(t *testing.T) {
	root := migrationTestSHA(21)
	source := &changingMemoryMigrationSource{fakeMemoryMigrationSource: fakeMemoryMigrationSource{
		head: root, commits: map[string]MemoryMigrationCommit{root: {CommitSHA: root}},
		snapshots: map[string][]MemoryMigrationFile{root: {{Path: "MEMORY.md", Data: []byte("preflight\n")}}},
	}}
	service, db := newMemoryMigrationServiceTest(t, source)
	_, err := service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{})
	require.ErrorContains(t, err, "snapshot changed after preflight")
	var memory memorystore.Memory
	require.NoError(t, db.First(&memory, "id = ?", "mem_1").Error)
	require.Empty(t, memory.ContentID)
	var versions int64
	require.NoError(t, db.Model(&contentstore.ContentVersion{}).Count(&versions).Error)
	require.Zero(t, versions)
}

func TestMemoryContentMigrationFailsBeforeWritesWhenBranchMovesAfterPreflight(t *testing.T) {
	root := migrationTestSHA(25)
	source := &changingHeadMemoryMigrationSource{fakeMemoryMigrationSource: fakeMemoryMigrationSource{
		head: root, commits: map[string]MemoryMigrationCommit{root: {CommitSHA: root}},
		snapshots: map[string][]MemoryMigrationFile{root: {{Path: "MEMORY.md", Data: []byte("root\n")}}},
	}}
	service, db := newMemoryMigrationServiceTest(t, source)
	_, err := service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{})
	require.ErrorContains(t, err, "source branch changed after preflight")
	var contents int64
	require.NoError(t, db.Model(&contentstore.VersionedContent{}).Count(&contents).Error)
	require.Zero(t, contents)
}

func TestMemoryContentMigrationRejectsNativeAuthorityAndStaleGeneration(t *testing.T) {
	t.Run("native authority", func(t *testing.T) {
		root := migrationTestSHA(31)
		source := &fakeMemoryMigrationSource{
			head: root, commits: map[string]MemoryMigrationCommit{root: {CommitSHA: root}},
			snapshots: map[string][]MemoryMigrationFile{root: {{Path: "MEMORY.md", Data: []byte("root\n")}}},
		}
		service, db := newMemoryMigrationServiceTest(t, source)
		now := time.Now().UTC()
		require.NoError(t, db.Create(&contentstore.Authority{
			ID: "cta_memory", WorkspaceID: "wsp_1", Domain: memoryContentDomain, BindingKey: "mem_1",
			State: contentstore.AuthorityStateNativeActive, Generation: 1, UpdatedBy: "operator", CreatedAt: now, UpdatedAt: now,
		}).Error)
		_, err := service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{})
		require.ErrorIs(t, err, ErrConflict)
		require.Empty(t, source.snapshotReads, "native authority must fail before any AGS payload read")
	})

	t.Run("stale generation", func(t *testing.T) {
		root := migrationTestSHA(32)
		source := &fakeMemoryMigrationSource{
			head: root, commits: map[string]MemoryMigrationCommit{root: {CommitSHA: root}},
			snapshots: map[string][]MemoryMigrationFile{root: {{Path: "MEMORY.md", Data: []byte("root\n")}}},
		}
		service, db := newMemoryMigrationServiceTest(t, source)
		memory, err := service.GetMemory(context.Background(), "wsp_1", "usr_1", "mem_1")
		require.NoError(t, err)
		content, err := service.ensureMemoryMigrationContent(context.Background(), "wsp_1", memory, "usr_1")
		require.NoError(t, err)
		require.NoError(t, db.Model(&contentstore.VersionedContent{}).Where("id = ?", content.ID).Update("generation", 2).Error)
		_, err = service.RunMemoryContentMigration(context.Background(), "wsp_1", "usr_1", "mem_1", MemoryContentMigrationRequest{})
		require.ErrorIs(t, err, ErrConflict)
		var versions int64
		require.NoError(t, db.Model(&contentstore.ContentVersion{}).Count(&versions).Error)
		require.Zero(t, versions)
	})
}

type changingMemoryMigrationSource struct {
	fakeMemoryMigrationSource
	reads atomic.Int64
}

type changingHeadMemoryMigrationSource struct {
	fakeMemoryMigrationSource
	reads atomic.Int64
}

func (source *changingHeadMemoryMigrationSource) ResolveMemoryMigrationHead(context.Context, repos.RepositoryIdentity, string) (string, error) {
	if source.reads.Add(1) == 1 {
		return source.head, nil
	}
	return migrationTestSHA(26), nil
}

func (source *changingMemoryMigrationSource) ReadMemoryMigrationSnapshot(ctx context.Context, repository repos.RepositoryIdentity, sha string, limits contentpayload.Limits) ([]MemoryMigrationFile, error) {
	files, err := source.fakeMemoryMigrationSource.ReadMemoryMigrationSnapshot(ctx, repository, sha, limits)
	if err == nil && source.reads.Add(1) > 1 {
		files[0].Data = []byte("changed\n")
	}
	return files, err
}

func newMemoryMigrationServiceTest(t *testing.T, source MemoryMigrationSource) (*MemoryService, *gorm.DB) {
	t.Helper()
	db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
	require.NoError(t, db.AutoMigrate(
		&memorystore.Memory{}, &memorystore.MemoryContentCheckpoint{}, &memorystore.MemoryGrant{}, &memorystore.MemoryPin{},
		&reposstore.ResourceRepository{}, &contentstore.Authority{},
		&contentstore.VersionedContent{}, &contentstore.ContentVersion{}, &contentstore.VersionOperation{},
		&contentstore.VersionPath{}, &contentstore.VersionPin{}, &contentstore.Outbox{}, &contentstore.StagedPayloadGC{},
		&migrationstore.Mapping{}, &migrationstore.Job{},
	))
	now := time.Now().UTC()
	require.NoError(t, db.Create(&memorystore.Memory{
		ID: "mem_1", Slug: "memory", Name: "Memory", StorageScope: "workspace", IsActive: true, Branch: "main",
		CreatedBy: "usr_1", UpdatedBy: "usr_1", CreatedAt: now, UpdatedAt: now,
	}).Error)
	require.NoError(t, db.Create(&memorystore.MemoryGrant{
		ID: "mgr_1", WorkspaceID: "wsp_1", MemoryID: "mem_1", SubjectType: "user", SubjectID: "usr_1", Role: "owner",
		CreatedBy: "usr_1", UpdatedBy: "usr_1", CreatedAt: now, UpdatedAt: now,
	}).Error)
	require.NoError(t, db.Create(&reposstore.ResourceRepository{
		ID: "rr_1", WorkspaceID: "wsp_1", ResourceType: reposstore.ResourceTypeMemory, ResourceID: "mem_1",
		RepoRole: reposstore.RepoRolePrimary, OwnerKind: "workspace", RepoOwner: "owner", RepoName: "memory",
		RepoFullName: "owner/memory", AGSRepoID: "repo_1", DefaultRef: "main", CreatedBy: "usr_1", UpdatedBy: "usr_1", CreatedAt: now, UpdatedAt: now,
	}).Error)
	tree, err := contentpayload.NewFilesystemTreeBackend(t.TempDir())
	require.NoError(t, err)
	t.Cleanup(func() { require.NoError(t, tree.Close()) })
	var ids atomic.Int64
	newID := func(prefix string) (string, error) { return fmt.Sprintf("%s_%d", prefix, ids.Add(1)), nil }
	core, err := contentservice.NewTree(contentstore.New(db), tree, newID)
	require.NoError(t, err)
	policy, err := NewMemoryContentAccessPolicy(memoryContentWorkspaceVerifier{allowed: true}, memorystore.New(db))
	require.NoError(t, err)
	authorized, err := contentservice.NewAuthorizedService(core, policy)
	require.NoError(t, err)
	migrations := migrationstore.New(db)
	worker, err := migration.NewWorker(migrations, authorized, newID, memoryContentDomain, "memory:test", 3*time.Minute)
	require.NoError(t, err)
	service := NewMemoryServiceWithOptions(memorystore.New(db), reposstore.New(db), nil, nil, MemoryServiceOptions{
		ContentAuthority: contentstore.New(db), MigrationSource: source, MigrationWorker: worker,
		MigrationMappings: migrations, MigrationContent: authorized,
	})
	return service, db
}
