package memory

import (
	"context"
	"fmt"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/contentpayload"
	contentservice "github.com/tidbcloud/AgentTeam/internal/service/content"
	contentstore "github.com/tidbcloud/AgentTeam/internal/store/content"
	memorystore "github.com/tidbcloud/AgentTeam/internal/store/resource/memory"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

func TestNativeMemoryCheckpointHistoryCompareBackupRestoreUsesExactTreeVersions(t *testing.T) {
	ctx := context.Background()
	native, db := newNativeMemoryContentTest(t)
	resolved, err := native.Resolve(ctx, "wsp_1", "mem_1", "usr_1")
	require.NoError(t, err)
	require.Equal(t, "ctn_1", resolved.ContentID)
	require.EqualValues(t, 1, resolved.Generation)

	request := NativeMemoryCheckpointRequest{
		NativeMemoryContentRequest: NativeMemoryContentRequest{
			WorkspaceID: "wsp_1", MemoryID: "mem_1", ContentID: "ctn_1", Generation: 1, Actor: "usr_1",
		},
		OperationID: "checkpoint-1", ReplicaID: "replica-1", CoveredFromSequence: 1, CoveredThroughSequence: 2,
		Files: []NativeMemoryContentFile{{Path: "MEMORY.md", Content: []byte("# first\n")}},
	}
	first, err := native.PublishCheckpoint(ctx, request)
	require.NoError(t, err)
	require.Equal(t, "ctn_1", first.ContentID)
	require.Empty(t, first.ParentVersionID)
	require.Len(t, first.RootDigest, 64)

	// A lost response can replay the exact server-owned operation even though
	// the live head has already advanced.
	replayed, err := native.PublishCheckpoint(ctx, request)
	require.NoError(t, err)
	require.Equal(t, first.VersionID, replayed.VersionID)

	request.VersionID = first.VersionID
	request.OperationID = "checkpoint-2"
	request.CoveredFromSequence = 3
	request.CoveredThroughSequence = 3
	request.Files = []NativeMemoryContentFile{
		{Path: "MEMORY.md", Content: []byte("# second\n")},
		{Path: "knowledge/concept.md", Content: []byte("bounded\n")},
	}
	second, err := native.PublishCheckpoint(ctx, request)
	require.NoError(t, err)
	require.Equal(t, first.VersionID, second.ParentVersionID)

	ref := request.NativeMemoryContentRequest
	ref.VersionID = second.VersionID
	history, err := native.ListVersions(ctx, ref, 10)
	require.NoError(t, err)
	require.Len(t, history, 2)
	require.Equal(t, second.VersionID, history[0].ID)

	diff, err := native.Compare(ctx, ref, first.VersionID, second.VersionID)
	require.NoError(t, err)
	require.Len(t, diff, 2)
	require.Equal(t, "MEMORY.md", diff[0].Path)
	require.Equal(t, "modified", diff[0].Change)

	backup, err := native.CreateBackup(ctx, ref, "before restore")
	require.NoError(t, err)
	require.Equal(t, second.VersionID, backup.VersionID)
	backups, err := native.ListBackups(ctx, ref, 64)
	require.NoError(t, err)
	require.Len(t, backups, 1)

	restored, err := native.Restore(ctx, ref, first.VersionID, "restore-1")
	require.NoError(t, err)
	require.Equal(t, second.VersionID, restored.ParentVersionID)
	require.Equal(t, first.RootDigest, restored.RootDigest)

	replayedRestore, err := native.Restore(ctx, ref, first.VersionID, "restore-1")
	require.NoError(t, err)
	require.Equal(t, restored.VersionID, replayedRestore.VersionID)

	var versions int64
	require.NoError(t, db.Model(&contentstore.ContentVersion{}).Where("content_id = ?", "ctn_1").Count(&versions).Error)
	require.EqualValues(t, 3, versions, "operation replay must not publish a second ContentVersion")
	var operations int64
	require.NoError(t, db.Model(&contentstore.VersionOperation{}).Where("content_id = ?", "ctn_1").Count(&operations).Error)
	require.EqualValues(t, 3, operations)
}

func TestNativeMemoryContentFailsClosedAfterGrantRevocationAndForWrongQuotaProfile(t *testing.T) {
	ctx := context.Background()
	native, db := newNativeMemoryContentTest(t)
	ref := NativeMemoryContentRequest{WorkspaceID: "wsp_1", MemoryID: "mem_1", ContentID: "ctn_1", Generation: 1, Actor: "usr_1"}

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

	// A generic tree lineage has wider defaults than the finite Memory quota.
	// The adapter refuses it instead of relying on non-transactional preflights.
	require.NoError(t, db.Create(&memorystore.MemoryGrant{
		ID: "mgr_2", WorkspaceID: "wsp_1", MemoryID: "mem_1", SubjectType: "user", SubjectID: "usr_1", Role: "owner",
		CreatedBy: "usr_1", UpdatedBy: "usr_1", CreatedAt: time.Now().UTC(), UpdatedAt: time.Now().UTC(),
	}).Error)
	require.NoError(t, db.Model(&contentstore.VersionedContent{}).Where("id = ?", "ctn_1").Updates(map[string]any{
		"max_retained_versions":    contentstore.DefaultTreeVersions,
		"max_stored_payload_bytes": contentstore.DefaultTreeBytes,
		"max_pins":                 contentstore.DefaultTreePins,
	}).Error)
	_, err = native.ListVersions(ctx, ref, 10)
	require.ErrorContains(t, err, "memory content quota profile is unavailable")
}

func TestNativeMemoryRestoreSnapshotAuthorizationQueriesAreConstant(t *testing.T) {
	queryCounts := make(map[int]int64)
	for _, fileCount := range []int{1, 64} {
		t.Run(fmt.Sprintf("files_%d", fileCount), func(t *testing.T) {
			ctx := context.Background()
			native, db := newNativeMemoryContentTest(t)
			files := make([]NativeMemoryContentFile, 0, fileCount)
			for index := 0; index < fileCount; index++ {
				files = append(files, NativeMemoryContentFile{Path: fmt.Sprintf("knowledge/%03d.md", index), Content: []byte("bounded\n")})
			}
			published, err := native.PublishCheckpoint(ctx, NativeMemoryCheckpointRequest{
				NativeMemoryContentRequest: NativeMemoryContentRequest{WorkspaceID: "wsp_1", MemoryID: "mem_1", ContentID: "ctn_1", Generation: 1, Actor: "usr_1"},
				OperationID:                "snapshot", ReplicaID: "replica", CoveredFromSequence: 1, CoveredThroughSequence: 1, Files: files,
			})
			require.NoError(t, err)

			var queries atomic.Int64
			require.NoError(t, db.Callback().Query().Before("gorm:query").Register("test:count_restore_snapshot_queries", func(*gorm.DB) {
				queries.Add(1)
			}))
			request := NativeMemoryContentRequest{WorkspaceID: "wsp_1", MemoryID: "mem_1", ContentID: "ctn_1", VersionID: published.VersionID, Generation: 1, Actor: "usr_1"}
			restored, err := native.snapshotFiles(ctx, request, "ctn_1", published.VersionID)
			require.NoError(t, err)
			require.Len(t, restored, fileCount)
			queryCounts[fileCount] = queries.Load()
		})
	}
	require.Equal(t, queryCounts[1], queryCounts[64], "restore authorization and exact-version reads must not grow per file")
}

func TestNativeMemoryBindingIsDormantUntilExactAuthorityActivation(t *testing.T) {
	ctx := context.Background()
	native, db := newNativeMemoryContentTest(t)
	authority := contentstore.New(db)
	service := NewMemoryServiceWithOptions(memorystore.New(db), nil, nil, nil, MemoryServiceOptions{
		ContentAuthority: authority, NativeContent: native, NativeContentEnabled: true,
	})

	enabled, err := service.nativeMemoryRead(ctx, "wsp_1", "mem_1")
	require.NoError(t, err)
	require.False(t, enabled, "missing authority must remain legacy")

	now := time.Now().UTC()
	require.NoError(t, db.Create(&contentstore.Authority{
		ID: "cta_wildcard", WorkspaceID: "wsp_1", Domain: memoryContentDomain,
		BindingKey: contentstore.AuthorityDomainDefaultBinding, State: contentstore.AuthorityStateNativeActive,
		Generation: 1, UpdatedBy: "operator", CreatedAt: now, UpdatedAt: now,
	}).Error)
	enabled, err = service.nativeMemoryRead(ctx, "wsp_1", "mem_1")
	require.NoError(t, err)
	require.False(t, enabled, "a wildcard row must not activate a product binding")

	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)
	enabled, err = service.nativeMemoryRead(ctx, "wsp_1", "mem_1")
	require.NoError(t, err)
	require.True(t, enabled)

	unavailable := NewMemoryServiceWithOptions(memorystore.New(db), nil, nil, nil, MemoryServiceOptions{
		ContentAuthority: authority, NativeContentEnabled: true,
	})
	_, err = unavailable.nativeMemoryRead(ctx, "wsp_1", "mem_1")
	require.ErrorContains(t, err, "native memory content is unavailable")
}

func newNativeMemoryContentTest(t *testing.T) (NativeMemoryContentStore, *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{},
		&contentstore.Authority{},
		&contentstore.VersionedContent{}, &contentstore.ContentVersion{}, &contentstore.VersionOperation{},
		&contentstore.VersionPath{}, &contentstore.VersionPin{}, &contentstore.Outbox{}, &contentstore.StagedPayloadGC{},
	))
	now := time.Now().UTC()
	require.NoError(t, db.Create(&memorystore.Memory{
		ID: "mem_1", Slug: "memory", Name: "Memory", Description: "", StorageScope: "workspace", IsActive: true,
		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)
	contentDB := contentstore.New(db)
	_, err = contentDB.CreateVersionedContent(context.Background(), contentstore.VersionedContent{
		ID: "ctn_1", WorkspaceID: "wsp_1", Domain: memoryContentDomain, BindingKey: "mem_1", Model: contentservice.TreeModel,
		MaxRetainedVersions: memoryContentMaxVersions, MaxStoredPayloadBytes: memoryContentMaxStoredBytes, MaxPins: memoryContentMaxPins,
		CreatedBy: "usr_1", UpdatedBy: "usr_1", 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()) })
	var ids atomic.Int64
	core, err := contentservice.NewTree(contentDB, tree, func(prefix string) (string, error) {
		return fmt.Sprintf("%s_%d", prefix, ids.Add(1)), nil
	})
	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)
	return NativeMemoryContentStore{Content: authorized}, db
}