package memory

import (
	"archive/zip"
	"bytes"
	"context"
	"errors"
	"io"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/store/content"
	memorystore "github.com/tidbcloud/AgentTeam/internal/store/resource/memory"
	"gorm.io/gorm"
)

func TestMemoryNativeAPIProjectionIsExactRestartSafeAndBounded(t *testing.T) {
	ctx := context.Background()
	service, db, seed := newNativeMemoryServiceTest(t)

	input := MemoryCheckpointMutation{
		OperationID: "checkpoint-1", ReplicaID: "replica-1", ExpectedVersionID: seed.VersionID, Generation: 1,
		CoveredFromSequence: 1, CoveredThroughSequence: 2,
		Files: []NativeMemoryContentFile{{Path: "MEMORY.md", Content: []byte("# checkpoint\n")}},
	}
	created, err := service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", input)
	require.NoError(t, err)
	require.Equal(t, seed.VersionID, created.ParentVersionID)
	require.EqualValues(t, 2, created.CoveredThroughSequence)

	var memory memorystore.Memory
	require.NoError(t, db.First(&memory, "id = ?", "mem_1").Error)
	require.Equal(t, created.VersionID, memory.ActiveContentVersionID)
	require.EqualValues(t, 2, memory.ContentHistorySequence)
	require.Equal(t, "replica-1", memory.ActiveReplicaID)
	require.EqualValues(t, 2, memory.AcknowledgedSequence)

	// The exact client retry still resolves after both Content Store publication
	// and the Memory projection have committed.
	replayed, err := service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", input)
	require.NoError(t, err)
	require.Equal(t, created.VersionID, replayed.VersionID)

	changed := input
	changed.Files = []NativeMemoryContentFile{{Path: "MEMORY.md", Content: []byte("changed reuse\n")}}
	_, err = service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", changed)
	require.Error(t, err)

	stale := input
	stale.OperationID = "checkpoint-stale"
	_, err = service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", stale)
	require.ErrorIs(t, err, ErrConflict)

	next := MemoryCheckpointMutation{
		OperationID: "checkpoint-2", ReplicaID: "replica-1", ExpectedVersionID: created.VersionID, Generation: 1,
		CoveredFromSequence: 3, CoveredThroughSequence: 3,
		Files: []NativeMemoryContentFile{
			{Path: "MEMORY.md", Content: []byte("# next\n")},
			{Path: "kb/topic.md", Content: []byte("bounded\n")},
		},
	}
	second, err := service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", next)
	require.NoError(t, err)
	require.NoError(t, db.First(&memory, "id = ?", "mem_1").Error)
	require.EqualValues(t, 3, memory.ContentHistorySequence)

	history, err := service.ListMemoryContentVersions(ctx, "wsp_1", "usr_1", "mem_1", 1, 2)
	require.NoError(t, err)
	require.Len(t, history.Items, 2)
	require.True(t, history.HasMore)
	require.Equal(t, second.VersionID, history.Items[0].VersionID)

	diff, err := service.CompareMemoryContentVersions(ctx, "wsp_1", "usr_1", "mem_1", created.VersionID, second.VersionID)
	require.NoError(t, err)
	require.Len(t, diff, 2)

	backup, err := service.CreateMemoryContentBackup(ctx, "wsp_1", "usr_1", "mem_1", MemoryBackupMutation{
		ExpectedVersionID: second.VersionID, Generation: 1, Label: "before restore",
	})
	require.NoError(t, err)
	require.Equal(t, second.VersionID, backup.VersionID)
	backups, err := service.ListMemoryContentBackups(ctx, "wsp_1", "usr_1", "mem_1", 64)
	require.NoError(t, err)
	require.Len(t, backups, 1)

	archive, err := service.OpenMemoryContentArchive(ctx, "wsp_1", "usr_1", "mem_1", second.VersionID)
	require.NoError(t, err)
	data, err := io.ReadAll(archive.Body)
	require.NoError(t, err)
	require.NoError(t, archive.Body.Close())
	reader, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
	require.NoError(t, err)
	require.Len(t, reader.File, 2)

	restored, err := service.RestoreMemoryContentVersion(ctx, "wsp_1", "usr_1", "mem_1", MemoryRestoreMutation{
		OperationID: "restore-1", ExpectedVersionID: second.VersionID, TargetVersionID: seed.VersionID, Generation: 1,
	})
	require.NoError(t, err)
	require.Equal(t, second.VersionID, restored.ParentVersionID)
	require.NoError(t, db.First(&memory, "id = ?", "mem_1").Error)
	require.EqualValues(t, 4, memory.ContentHistorySequence)
	replayedRestore, err := service.RestoreMemoryContentVersion(ctx, "wsp_1", "usr_1", "mem_1", MemoryRestoreMutation{
		OperationID: "restore-1", ExpectedVersionID: second.VersionID, TargetVersionID: seed.VersionID, Generation: 1,
	})
	require.NoError(t, err)
	require.Equal(t, restored.VersionID, replayedRestore.VersionID)

	var versionCount int64
	require.NoError(t, db.Model(&content.ContentVersion{}).Where("content_id = ?", "ctn_1").Count(&versionCount).Error)
	require.EqualValues(t, 4, versionCount)

	require.NoError(t, db.Where("workspace_id = ? AND memory_id = ? AND subject_id = ?", "wsp_1", "mem_1", "usr_1").Delete(&memorystore.MemoryGrant{}).Error)
	_, err = service.ListMemoryContentVersions(ctx, "wsp_1", "usr_1", "mem_1", 1, 10)
	require.ErrorIs(t, err, ErrForbidden)
}

func TestMemoryNativeCheckpointRejectsSequenceGapReplicaTakeoverAndStaleGeneration(t *testing.T) {
	ctx := context.Background()
	service, _, seed := newNativeMemoryServiceTest(t)
	base := MemoryCheckpointMutation{
		OperationID: "checkpoint-gap", ReplicaID: "replica-1", ExpectedVersionID: seed.VersionID, Generation: 1,
		CoveredFromSequence: 2, CoveredThroughSequence: 2,
		Files: []NativeMemoryContentFile{{Path: "MEMORY.md", Content: []byte("gap\n")}},
	}
	_, err := service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", base)
	require.ErrorIs(t, err, ErrConflict)

	base.CoveredFromSequence = 1
	base.OperationID = "checkpoint-generation"
	base.Generation = 2
	_, err = service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", base)
	require.ErrorIs(t, err, ErrConflict)

	base.Generation = 1
	base.OperationID = "checkpoint-ok"
	first, err := service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", base)
	require.NoError(t, err)
	base.OperationID = "checkpoint-takeover"
	base.ExpectedVersionID = first.VersionID
	base.ReplicaID = "replica-2"
	base.CoveredFromSequence = 2
	base.CoveredThroughSequence = 2
	_, err = service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", base)
	require.ErrorIs(t, err, ErrConflict)
}

func TestMemoryNativeCheckpointRetryFinalizesProjectionWithoutRepublishing(t *testing.T) {
	ctx := context.Background()
	service, db, seed := newNativeMemoryServiceTest(t)
	input := MemoryCheckpointMutation{
		OperationID: "checkpoint-restart", ReplicaID: "replica-1", ExpectedVersionID: seed.VersionID, Generation: 1,
		CoveredFromSequence: 1, CoveredThroughSequence: 1,
		Files: []NativeMemoryContentFile{{Path: "MEMORY.md", Content: []byte("# restart safe\n")}},
	}

	var failProjection atomic.Bool
	failProjection.Store(true)
	require.NoError(t, db.Callback().Create().Before("gorm:create").Register("test:fail_first_memory_projection", func(tx *gorm.DB) {
		if tx.Statement != nil && tx.Statement.Schema != nil && tx.Statement.Schema.Table == (memorystore.MemoryContentCheckpoint{}).TableName() && failProjection.CompareAndSwap(true, false) {
			tx.AddError(errors.New("injected memory projection failure"))
		}
	}))

	var versionsBefore int64
	require.NoError(t, db.Model(&content.ContentVersion{}).Where("content_id = ?", "ctn_1").Count(&versionsBefore).Error)
	_, err := service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", input)
	require.ErrorContains(t, err, "injected memory projection failure")

	var versionsAfterFailure int64
	require.NoError(t, db.Model(&content.ContentVersion{}).Where("content_id = ?", "ctn_1").Count(&versionsAfterFailure).Error)
	require.Equal(t, versionsBefore+1, versionsAfterFailure, "Content Store publication commits before the injected product projection failure")
	stale, err := service.store.GetMemory(ctx, "wsp_1", "mem_1")
	require.NoError(t, err)
	require.Equal(t, seed.VersionID, stale.ActiveContentVersionID)
	require.Zero(t, stale.AcknowledgedSequence)
	_, found, err := service.store.FindMemoryCheckpointByOperation(ctx, "wsp_1", "mem_1", input.OperationID)
	require.NoError(t, err)
	require.False(t, found, "the failed transaction must not leave a product receipt")
	changedMetadata := input
	changedMetadata.CoveredThroughSequence = 2
	_, err = service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", changedMetadata)
	require.ErrorIs(t, err, ErrConflict, "changed acknowledgement metadata must not reuse the already published operation")
	var versionsAfterChangedRetry int64
	require.NoError(t, db.Model(&content.ContentVersion{}).Where("content_id = ?", "ctn_1").Count(&versionsAfterChangedRetry).Error)
	require.Equal(t, versionsAfterFailure, versionsAfterChangedRetry)

	projected, err := service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", input)
	require.NoError(t, err)
	require.NotEqual(t, seed.VersionID, projected.VersionID)
	var versionsAfterRetry int64
	require.NoError(t, db.Model(&content.ContentVersion{}).Where("content_id = ?", "ctn_1").Count(&versionsAfterRetry).Error)
	require.Equal(t, versionsAfterFailure, versionsAfterRetry, "retry must reuse the durable Content Store operation receipt")
	advanced, err := service.store.GetMemory(ctx, "wsp_1", "mem_1")
	require.NoError(t, err)
	require.Equal(t, projected.VersionID, advanced.ActiveContentVersionID)
	require.EqualValues(t, 1, advanced.AcknowledgedSequence)

	replayed, err := service.PublishMemoryCheckpoint(ctx, "wsp_1", "usr_1", "mem_1", input)
	require.NoError(t, err)
	require.Equal(t, projected.VersionID, replayed.VersionID)
	var versionsAfterReplay int64
	require.NoError(t, db.Model(&content.ContentVersion{}).Where("content_id = ?", "ctn_1").Count(&versionsAfterReplay).Error)
	require.Equal(t, versionsAfterRetry, versionsAfterReplay)
}

func newNativeMemoryServiceTest(t *testing.T) (*MemoryService, *gorm.DB, NativeMemoryVersionResult) {
	t.Helper()
	native, db := newNativeMemoryContentTest(t)
	ctx := context.Background()
	seed, err := native.PublishCheckpoint(ctx, NativeMemoryCheckpointRequest{
		NativeMemoryContentRequest: NativeMemoryContentRequest{WorkspaceID: "wsp_1", MemoryID: "mem_1", ContentID: "ctn_1", Generation: 1, Actor: "usr_1"},
		OperationID:                "seed-import", ReplicaID: "migration", CoveredFromSequence: 1, CoveredThroughSequence: 1,
		Files: []NativeMemoryContentFile{{Path: "MEMORY.md", Content: []byte("# seed\n")}},
	})
	require.NoError(t, err)
	require.NoError(t, db.Model(&memorystore.Memory{}).Where("id = ?", "mem_1").Updates(map[string]any{
		"content_id": "ctn_1", "active_content_version_id": seed.VersionID, "content_generation": 1, "content_history_sequence": 1,
	}).Error)
	store := memorystore.New(db)
	seedAuthoredAt := seed.CreatedAt.UTC()
	_, err = store.CreateMemoryContentCheckpoint(ctx, memorystore.MemoryContentCheckpoint{
		ID: "mcp_seed", WorkspaceID: "wsp_1", MemoryID: "mem_1", ContentID: "ctn_1", ContentVersionID: seed.VersionID,
		OperationID: "seed-import", OperationKind: "import", Generation: 1, HistorySequence: 1, RootDigest: seed.RootDigest,
		FileCount: seed.FileCount, SizeBytes: seed.SizeBytes, SourceAuthoredAt: &seedAuthoredAt,
		CreatedBy: "usr_1", CreatedAt: seedAuthoredAt,
	})
	require.NoError(t, err)
	now := time.Now().UTC()
	require.NoError(t, db.Create(&content.Authority{
		ID: "cta_memory_native", WorkspaceID: "wsp_1", Domain: memoryContentDomain, BindingKey: "mem_1",
		State: content.AuthorityStateNativeActive, Generation: 1, UpdatedBy: "operator", CreatedAt: now, UpdatedAt: now,
	}).Error)
	service := NewMemoryServiceWithOptions(store, nil, nil, nil, MemoryServiceOptions{
		ContentAuthority: content.New(db), NativeContent: native, NativeContentEnabled: true,
	})
	return service, db, seed
}