package memory

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/application/repository"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/stretchr/testify/require"
)

func TestRecoveryMaintenanceRetainsPolicyAndLeaseChecks(t *testing.T) {
	for _, maintenance := range []string{"topics", "consolidation"} {
		for _, phase := range []string{"before", "after"} {
			t.Run(maintenance+"/"+phase, func(t *testing.T) {
				svc, db, ctx, messages, models, queue := recoveryHarness(t)
				scope, err := ResolveScope(ctx)
				require.NoError(t, err)
				svc.ScheduleExtraction(ctx, "s", "m", "model")
				if maintenance == "topics" {
					_, err := svc.repo.BumpTopic(ctx, scope, "database tuning", types.NormalizeTopicKey("database tuning"), "database tuning")
					require.NoError(t, err)
					models.response = `{"memories":[],"topics":["deployment automation"]}`
					models.responseFor = map[string]string{topicAdjudicationPrompt: `{"resolutions":[{"index":0,"same_as":0}]}`}
				} else {
					for i := range 6 {
						require.NoError(t, svc.repo.CreateItem(ctx, &types.MemoryItem{TenantID: 1, SubjectID: scope.SubjectID,
							Kind: types.MemoryKindFact, Topic: "deployment automation", Content: fmt.Sprintf("deployment automation uses deterministic verification %d", i)}))
					}
					models.responseFor = map[string]string{consolidationSystemPrompt: `{"statement":"deployment automation uses deterministic verification"}`}
				}
				var beforeItems int64
				require.NoError(t, db.Model(&types.MemoryItem{}).Count(&beforeItems).Error)
				before, err := svc.repo.GetSubject(ctx, scope)
				require.NoError(t, err)
				mutate := func() { messages.owners = map[string]interfaces.MemoryScope{} }
				models.onChat = func() {
					if phase == "before" {
						models.onGetChatModel = mutate
					} else {
						models.onChat = mutate
					}
				}
				require.NoError(t, svc.Handle(context.Background(), queue.pop()))
				calls := 1
				if phase == "after" {
					calls = 2
				}
				require.Equal(t, calls, models.callCount())
				var afterItems int64
				require.NoError(t, db.Model(&types.MemoryItem{}).Count(&afterItems).Error)
				require.Equal(t, beforeItems, afterItems)
				after, err := svc.repo.GetSubject(ctx, scope)
				require.NoError(t, err)
				require.Equal(t, before.SessionCursors, after.SessionCursors)
				require.Equal(t, before.PendingSessions, after.PendingSessions)
				require.Zero(t, after.ExtractFailures)
				if maintenance == "topics" {
					topics, err := svc.repo.TopTopics(ctx, scope, 10)
					require.NoError(t, err)
					require.Len(t, topics, 1)
					require.Equal(t, 1, topics[0].Hits, "revoked maintenance cannot fall back to counting the rejected model output")
				}
			})
		}
	}
}

func TestRecoveryWithNoReadableTranscriptSkipsAutomaticMaintenance(t *testing.T) {
	svc, _, ctx, messages, models, queue := recoveryHarness(t)
	svc.ScheduleExtraction(ctx, "s", "m", "model")
	scope, err := ResolveScope(ctx)
	require.NoError(t, err)
	for i := range 6 {
		require.NoError(t, svc.repo.CreateItem(ctx, &types.MemoryItem{TenantID: 1, SubjectID: scope.SubjectID,
			Kind: types.MemoryKindFact, Content: fmt.Sprintf("deployment automation uses deterministic verification %d", i)}))
	}
	messages.owners = map[string]interfaces.MemoryScope{}
	require.NoError(t, svc.Handle(context.Background(), queue.pop()))
	require.Zero(t, models.callCount())
	subject, err := svc.repo.GetSubject(ctx, scope)
	require.NoError(t, err)
	require.Nil(t, subject.ConsolidatedAt)
	require.Nil(t, subject.ExtractCursor)
}

func TestRecoveryRechecksPolicyAndReservationAroundProvider(t *testing.T) {
	for _, phase := range []string{"before", "after", "truncated"} {
		for _, change := range []string{"deleted subject", "disabled subject", "disabled workspace", "revoked owner", "disabled Agent", "revoked context Agent", "replaced attempt", "expired lease"} {
			t.Run(phase+"/"+change, func(t *testing.T) {
				svc, db, ctx, messages, models, queue := recoveryHarness(t)
				scope, err := ResolveScope(ctx)
				require.NoError(t, err)
				base := time.Now().Add(-time.Hour)
				prior := userMessage("s", "earlier preference", base)
				input := userMessage("s", "current preference", base.Add(time.Minute))
				messages.set("s", []*types.Message{prior, input})
				svc.ScheduleExtraction(ctx, "s", "m", "model")
				cursor := base.Add(time.Second)
				require.NoError(t, db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).
					Updates(map[string]interface{}{"session_cursors": types.MemorySessionCursors{"s": cursor}, "extract_cursor": cursor}).Error)
				before, err := svc.repo.GetSubject(ctx, scope)
				require.NoError(t, err)
				models.response = `{"memories":[{"action":"add","kind":"preference","topic":"working style","content":"prefers clear examples"}],"topics":["software engineering"]}`
				mutate := func() {
					switch change {
					case "deleted subject":
						require.NoError(t, db.Where("tenant_id = ?", 1).Delete(&types.MemorySubject{}).Error)
					case "disabled subject":
						require.NoError(t, svc.repo.UpdateSubjectEnabled(ctx, scope, false))
					case "disabled workspace":
						svc.tenantRepo.(*stubTenantRepo).set(1, &types.MemoryConfig{Enabled: false})
					case "revoked owner":
						messages.owners = map[string]interfaces.MemoryScope{}
					case "disabled Agent":
						messages.deniedMessages = map[string]bool{input.ID: true}
					case "revoked context Agent":
						messages.deniedMessages = map[string]bool{prior.ID: true}
					case "replaced attempt":
						require.NoError(t, db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).Update("extract_attempt_id", "replacement").Error)
					case "expired lease":
						require.NoError(t, db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).Update("extract_lease_until", time.Now().Add(-time.Second)).Error)
					}
				}
				if phase == "before" {
					models.onGetChatModel = mutate
				} else {
					models.onChat = mutate
					if phase == "truncated" {
						models.truncateUntilCall = 1
					}
				}
				err = svc.Handle(context.Background(), queue.pop())
				stale := change == "replaced attempt" || change == "expired lease"
				if stale {
					require.ErrorIs(t, err, errExtractionAttemptLost)
				} else {
					require.NoError(t, err)
				}
				wantCalls := 1
				if phase == "before" {
					wantCalls = 0
				}
				require.Equal(t, wantCalls, models.callCount(), "revocation must also block the truncation retry and maintenance")
				var count int64
				require.NoError(t, db.Model(&types.MemoryItem{}).Count(&count).Error)
				require.Zero(t, count, "discard provider output after revocation")
				after, err := svc.repo.GetSubject(ctx, scope)
				require.NoError(t, err)
				if change == "deleted subject" {
					require.Nil(t, after, "worker cannot recreate a deleted subject")
					return
				}
				require.Equal(t, before.PendingSessions, after.PendingSessions)
				require.Equal(t, before.SessionCursors, after.SessionCursors)
				require.True(t, before.ExtractCursor.Equal(*after.ExtractCursor))
				if !stale {
					require.Zero(t, after.ExtractFailures, "policy changes pause without charging the provider retry budget")
				}
				if change == "replaced attempt" {
					require.Equal(t, "replacement", after.ExtractAttemptID)
				}
			})
		}
	}
}

func TestRecoveryRechecksWebAccountLifecycle(t *testing.T) {
	for _, phase := range []string{"queued", "before", "after"} {
		for _, change := range []string{"UPDATE users SET is_active = false", "UPDATE users SET deleted_at = CURRENT_TIMESTAMP", "UPDATE tenant_members SET status = 'suspended'"} {
			t.Run(phase+"/"+change, func(t *testing.T) {
				svc, db, ctx, _, models, queue := recoveryHarness(t)
				require.NoError(t, db.AutoMigrate(&types.Session{}, &types.Message{}))
				for _, sql := range []string{
					`CREATE TABLE users (id TEXT PRIMARY KEY, is_active BOOLEAN, is_system_admin BOOLEAN, deleted_at DATETIME)`,
					`CREATE TABLE tenant_members (user_id TEXT, tenant_id INTEGER, status TEXT, deleted_at DATETIME)`,
					`INSERT INTO users (id, is_active, is_system_admin) VALUES ('alice', true, false)`,
					`INSERT INTO tenant_members (user_id, tenant_id, status) VALUES ('alice', 1, 'active')`,
				} {
					require.NoError(t, db.Exec(sql).Error)
				}
				base := time.Now().Add(-time.Hour).UTC().Truncate(time.Microsecond)
				session := &types.Session{TenantID: 1, UserID: "alice"}
				require.NoError(t, db.Create(session).Error)
				input := &types.Message{SessionID: session.ID, RequestID: "r", Role: "user", Content: "prefers simple examples", CreatedAt: base}
				answer := &types.Message{SessionID: session.ID, RequestID: "r", Role: "assistant", Content: "okay", IsCompleted: true, CreatedAt: base.Add(time.Second)}
				require.NoError(t, db.Create(input).Error)
				require.NoError(t, db.Create(answer).Error)
				svc.messageRepo = repository.NewMessageRepository(db)
				svc.ScheduleExtraction(ctx, session.ID, answer.ID, "model")
				task := queue.pop()
				require.NotNil(t, task)
				scope, err := ResolveScope(ctx)
				require.NoError(t, err)
				before, err := svc.repo.GetSubject(ctx, scope)
				require.NoError(t, err)
				mutate := func() { require.NoError(t, db.Exec(change).Error) }
				models.response = `{"memories":[{"action":"add","content":"prefers simple examples"}],"topics":["examples"]}`
				switch phase {
				case "queued":
					mutate()
				case "before":
					models.onGetChatModel = mutate
				case "after":
					models.onChat = mutate
				}
				require.NoError(t, svc.Handle(context.Background(), task))
				wantCalls := 0
				if phase == "after" {
					wantCalls = 1
				}
				require.Equal(t, wantCalls, models.callCount())
				after, err := svc.repo.GetSubject(ctx, scope)
				require.NoError(t, err)
				require.Equal(t, before.SessionCursors, after.SessionCursors, "revoked owners never replay or advance historical cursors")
				require.Nil(t, after.ExtractCursor)
				require.Zero(t, after.ExtractFailures)
				var count int64
				require.NoError(t, db.Model(&types.MemoryItem{}).Count(&count).Error)
				require.Zero(t, count)
			})
		}
	}
}