package service

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	apprepo "github.com/Tencent/WeKnora/internal/application/repository"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/hibiken/asynq"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func openWikiCheckpointDB(t *testing.T, path string) *gorm.DB {
	t.Helper()
	db, err := gorm.Open(sqlite.Open(path), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	t.Cleanup(func() { _ = sqlDB.Close() })
	require.NoError(t, db.AutoMigrate(&wikiIngestCheckpoint{}))
	return db
}

func checkpointChunks(count int) []*types.Chunk {
	chunks := make([]*types.Chunk, count)
	for i := range chunks {
		chunks[i] = &types.Chunk{ID: fmt.Sprintf("chunk-%03d", i), TenantID: 7,
			KnowledgeBaseID: "kb", KnowledgeID: "doc", ChunkType: types.ChunkTypeText, IsEnabled: true,
			ChunkIndex: i, Content: fmt.Sprintf("batch-%03d Alpha ", i) + strings.Repeat("text ", 2398)}
	}
	return chunks
}

func TestWikiCheckpointCitationResumesAcrossDatabaseReopen(t *testing.T) {
	path := filepath.Join(t.TempDir(), "checkpoints.sqlite")
	db := openWikiCheckpointDB(t, path)
	svc := &wikiIngestService{checkpointDB: db}
	scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
	chunks := checkpointChunks(wikiCitationSlice + 1)
	var calls atomic.Int32
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		calls.Add(1)
		return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]}}`}, nil
	}}
	got, _, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
	require.ErrorIs(t, err, errWikiCheckpointYield)
	require.Nil(t, got)
	require.Equal(t, wikiCitationSlice, completed)
	require.Equal(t, int32(wikiCitationSlice), calls.Load())
	sqlDB, _ := db.DB()
	require.NoError(t, sqlDB.Close())

	svc = &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
	got, _, completed, err = svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
	require.NoError(t, err)
	require.Len(t, got["entity/alpha"], len(chunks))
	require.Equal(t, len(chunks), completed)
	require.Equal(t, int32(len(chunks)), calls.Load())
	_, _, _, err = svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
	require.NoError(t, err)
	require.Equal(t, int32(len(chunks)), calls.Load(), "restart must not regenerate validated batches")
}

func TestWikiCheckpointCitationFailureKeepsSuccessfulSibling(t *testing.T) {
	db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
	svc := &wikiIngestService{checkpointDB: db}
	scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
	chunks := checkpointChunks(2)
	var calls [2]atomic.Int32
	var failed atomic.Bool
	model := &wikiScriptedChat{id: t.Name(), call: func(ctx context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		if strings.Contains(messages[0].Content, "batch-000") {
			calls[0].Add(1)
			return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]}}`}, nil
		}
		calls[1].Add(1)
		if failed.CompareAndSwap(false, true) {
			// Fail only after the sibling is durably persisted.
			deadline := time.NewTimer(time.Second)
			defer deadline.Stop()
			for {
				var count int64
				db.Model(&wikiIngestCheckpoint{}).Count(&count)
				if count > 0 {
					break
				}
				select {
				case <-ctx.Done():
					return nil, ctx.Err()
				case <-deadline.C:
					return nil, errors.New("checkpoint not persisted")
				case <-time.After(time.Millisecond):
				}
			}
			return nil, errors.New("invalid request")
		}
		return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]}}`}, nil
	}}
	got, _, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
	require.Error(t, err)
	require.Nil(t, got)
	require.Equal(t, 1, completed)
	svc = &wikiIngestService{checkpointDB: db}
	got, _, completed, err = svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
	require.NoError(t, err)
	require.Equal(t, 2, completed)
	require.Len(t, got["entity/alpha"], 2)
	require.Equal(t, int32(1), calls[0].Load())
	require.Equal(t, int32(2), calls[1].Load())
}

func TestWikiCheckpointFreshCitationCannotReplaceInvalidType(t *testing.T) {
	for _, test := range []struct{ name, output string }{
		{"invalid-new-slug", `{"citations":{},"new_slugs":[{"type":"entity","name":"New","slug":"concept/new","source_chunks":["c000"]}]}`},
	} {
		t.Run(test.name, func(t *testing.T) {
			db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
			svc := &wikiIngestService{checkpointDB: db}
			scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
			chunks := []*types.Chunk{{ID: "first", Content: "first", ChunkIndex: 0}, {ID: "second", Content: "second", ChunkIndex: 1}}
			var calls atomic.Int32
			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
				if calls.Add(1) == 1 {
					return &types.ChatResponse{Content: test.output}, nil
				}
				return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]}}`}, nil
			}}
			got, _, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
			require.Error(t, err)
			require.Zero(t, completed)
			require.Nil(t, got)
			require.Equal(t, int32(1), calls.Load())
			var saved int64
			require.NoError(t, db.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "citation").Count(&saved).Error)
			require.Zero(t, saved, "invalid types cannot become checkpoints")
		})
	}
}

func TestWikiCheckpointUnknownSlugCannotUseSiblingDeclarationOrSavePartialBatch(t *testing.T) {
	db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
	svc := &wikiIngestService{checkpointDB: db}
	scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
	chunks := []*types.Chunk{{ID: "first", Content: "first", ChunkIndex: 0}, {ID: "second", Content: "second", ChunkIndex: 1}}
	var calls atomic.Int32
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		call := calls.Add(1)
		if call == 1 {
			return &types.ChatResponse{Content: `{"citations":`, FinishReason: "length"}, nil
		}
		if call == 2 {
			return &types.ChatResponse{Content: `{"citations":{},"new_slugs":[{"type":"concept","name":"New","slug":"concept/new","source_chunks":["c000"]}]}`}, nil
		}
		return &types.ChatResponse{Content: `{"citations":{"concept/new":["c000"]}}`}, nil
	}}
	got, _, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
	require.ErrorContains(t, err, "citation partition 2/2")
	require.ErrorContains(t, err, "undeclared slug")
	require.NotContains(t, err.Error(), "persisted")
	require.Nil(t, got)
	require.Zero(t, completed)
	require.Equal(t, int32(3), calls.Load(), "semantic errors have the same finite retry budget")
	var saved int64
	require.NoError(t, db.Model(&wikiIngestCheckpoint{}).Count(&saved).Error)
	require.Zero(t, saved)
}

func TestWikiCheckpointCorruptCitationDoesNotTriggerModelRepair(t *testing.T) {
	db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
	svc := &wikiIngestService{checkpointDB: db}
	scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
	chunks := checkpointChunks(1)
	batch := splitChunksIntoCitationBatches(chunks)[0]
	key := wikiCheckpointHash([]any{"entity/alpha", "English", wikiCheckpointChunks(batch.chunks)})
	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "citation", key, citationBatchResult{Citations: map[string][]string{"concept/undeclared": {chunks[0].ID}}}))
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		t.Fatal("corrupt persisted work must be investigated without invoking the model")
		return nil, nil
	}}
	got, _, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
	require.ErrorContains(t, err, "invalid persisted citation")
	require.Nil(t, got)
	require.Zero(t, completed)
}

type checkpointModelService struct {
	interfaces.ModelService
	model *types.Model
	chat  chat.Chat
}

func (m *checkpointModelService) GetModelByID(context.Context, string) (*types.Model, error) {
	return m.model, nil
}
func (m *checkpointModelService) GetChatModel(context.Context, string) (chat.Chat, error) {
	return m.chat, nil
}

func TestWikiCheckpointSignatureStabilityAndInvalidation(t *testing.T) {
	db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
	m := &checkpointModelService{model: &types.Model{ID: "model", Name: "test", TenantID: 7}}
	svc := &wikiIngestService{checkpointDB: db, modelService: m}
	model := &wikiScriptedChat{id: "model"}
	payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb"}
	op := WikiPendingOp{KnowledgeID: "doc", dbID: 42}
	chunks := checkpointChunks(2)
	batch := &WikiBatchContext{ExtractionGranularity: types.WikiExtractionStandard}
	signature := func() string {
		scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, "content", "title", "English", batch)
		require.NoError(t, err)
		return scope.Fingerprint
	}
	original := signature()
	m.model.UpdatedAt = time.Now()
	m.model.DisplayName = "new display label"
	chunks[0].UpdatedAt = time.Now()
	chunks[0].Status = int(types.ChunkStatusIndexed)
	chunks[0].IndexStatus = "processing"
	chunks[0], chunks[1] = chunks[1], chunks[0]
	require.Equal(t, original, signature(), "runtime timestamps/order must not invalidate resumable work")
	m.model.Parameters.BaseURL = "https://new-provider.invalid"
	require.NotEqual(t, original, signature())
	m.model.Parameters.BaseURL = ""
	chunks[0].Content += "updated source"
	require.NotEqual(t, original, signature())
	chunks[0].Content = strings.TrimSuffix(chunks[0].Content, "updated source")
	batch.ContentInstructions = "changed rules"
	require.NotEqual(t, original, signature())
	batch.ContentInstructions = ""
	op.dbID++
	require.NotEqual(t, original, signature(), "an explicitly new ingest must reconsider existing contributions")
	chunks[0].TenantID = 8
	_, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, "content", "title", "English", batch)
	require.ErrorContains(t, err, "scope mismatch")
}

func TestWikiCheckpointIsolationAndCorruptCitationRejection(t *testing.T) {
	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))}
	scope := &wikiCheckpointScope{7, "kb", "doc", "v1"}
	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "summary", "key", "complete"))
	for _, other := range []wikiCheckpointScope{{8, "kb", "doc", "v1"}, {7, "other", "doc", "v1"}, {7, "kb", "other", "v1"}, {7, "kb", "doc", "v2"}} {
		var value string
		hit, err := svc.loadWikiCheckpoint(context.Background(), &other, "summary", "key", &value)
		require.NoError(t, err)
		require.False(t, hit)
	}
	batch := splitChunksIntoCitationBatches(checkpointChunks(1))[0]
	allowed := map[string]bool{"entity/alpha": true}
	for _, result := range []citationBatchResult{
		{Citations: map[string][]string{"entity/alpha": {"another-document-chunk"}}},
		{Citations: map[string][]string{"concept/unknown": {"chunk-000"}}},
		{Citations: map[string][]string{}, NewSlugs: []newSlugFromCitation{{Name: "wrong type", Type: "entity", Slug: "concept/wrong"}}},
	} {
		require.Error(t, validateWikiCitationCheckpoint(&result, batch, allowed))
	}
}

type checkpointKnowledge struct{ interfaces.KnowledgeService }

func (*checkpointKnowledge) GetKnowledgeByIDOnly(_ context.Context, id string) (*types.Knowledge, error) {
	return &types.Knowledge{ID: id, TenantID: 7, KnowledgeBaseID: "kb", Title: "title"}, nil
}

type checkpointWikiPages struct {
	interfaces.WikiPageService
	mu         sync.Mutex
	pages      map[string]*types.WikiPage
	writes     int
	publishErr error
}

func (s *checkpointWikiPages) GetPageBySlug(_ context.Context, _, slug string) (*types.WikiPage, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if p := s.pages[slug]; p != nil {
		out := *p
		return &out, nil
	}
	return nil, apprepo.ErrWikiPageNotFound
}
func (s *checkpointWikiPages) ListBySlugs(_ context.Context, _ string, slugs []string) (map[string]*types.WikiPageLite, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	result := make(map[string]*types.WikiPageLite)
	for _, slug := range slugs {
		if page := s.pages[slug]; page != nil {
			result[slug] = &types.WikiPageLite{Slug: page.Slug, Title: page.Title, PageType: page.PageType}
		}
	}
	return result, nil
}
func (s *checkpointWikiPages) FindPagesByNormalizedTitles(_ context.Context, _ string, kind string, titles []string) ([]*types.WikiPageLite, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var result []*types.WikiPageLite
	for _, page := range s.pages {
		for _, title := range titles {
			if page.PageType == kind && normalizeWikiIdentityTitle(page.Title) == title {
				result = append(result, &types.WikiPageLite{Slug: page.Slug, Title: page.Title, PageType: page.PageType})
				break
			}
		}
	}
	return result, nil
}
func (s *checkpointWikiPages) CreatePage(_ context.Context, page *types.WikiPage) (*types.WikiPage, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := *page
	s.pages[page.Slug] = &out
	s.writes++
	return page, nil
}
func (s *checkpointWikiPages) UpdatePage(ctx context.Context, page *types.WikiPage) (*types.WikiPage, error) {
	return s.CreatePage(ctx, page)
}
func (*checkpointWikiPages) ListSlugsBySourceRef(context.Context, string, string) ([]string, error) {
	return nil, nil
}
func (*checkpointWikiPages) ListDistinctCategoryPaths(context.Context, string, int) ([][]string, error) {
	return nil, nil
}
func (s *checkpointWikiPages) UpdatePageMeta(ctx context.Context, page *types.WikiPage) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.publishErr != nil {
		return s.publishErr
	}
	out := *page
	s.pages[page.Slug] = &out
	return nil
}

func TestWikiCheckpointPageSliceResumesAndRejectsChangedPage(t *testing.T) {
	db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
	chunks := checkpointChunks(1)
	pages := &checkpointWikiPages{pages: make(map[string]*types.WikiPage)}
	scope := &wikiCheckpointScope{7, "kb", "doc", "v1"}
	updates := make([]SlugUpdate, wikiPageSlice+1)
	for i := range updates {
		updates[i] = SlugUpdate{KnowledgeID: "doc", SourceRef: "doc", Slug: fmt.Sprintf("summary/%d", i), Type: types.WikiPageTypeSummary, DocTitle: "title", SummaryBody: "complete summary"}
	}
	svc := &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &checkpointChunksRepo{chunks: chunks}}
	run := func() int {
		batch := &WikiBatchContext{}
		batch.checkpointScopes.Store("doc", scope)
		batch.checkpointSourceHashes.Store("doc", wikiEvidenceSourceHash(chunks))
		yielded := 0
		for _, update := range updates {
			_, _, _, err := svc.reduceSlugUpdates(context.Background(), nil, "kb", update.Slug, []SlugUpdate{update}, 7, batch, nil)
			if errors.Is(err, errWikiCheckpointYield) {
				yielded++
			} else {
				require.NoError(t, err)
			}
		}
		return yielded
	}
	require.Equal(t, 1, run())
	require.Equal(t, wikiPageSlice, pages.writes)
	svc = &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &checkpointChunksRepo{chunks: chunks}}
	require.Zero(t, run())
	require.Equal(t, wikiPageSlice+1, pages.writes)
	pages.pages[updates[0].Slug].Content = "changed source contribution"
	require.Zero(t, run())
	require.Equal(t, wikiPageSlice+2, pages.writes)
	require.Equal(t, "complete summary", pages.pages[updates[0].Slug].Content)
	pages.pages[updates[0].Slug].SourceRefs = nil
	require.Zero(t, run())
	require.Equal(t, wikiPageSlice+3, pages.writes)
	require.Contains(t, pages.pages[updates[0].Slug].SourceRefs, "doc")
}

func TestWikiCheckpointYieldBudgetDoesNotHideRealFailures(t *testing.T) {
	ctx, cancel := context.WithCancelCause(context.Background())
	cancel(errWikiCheckpointYield)
	require.ErrorIs(t, wikiCheckpointError(ctx, fmt.Errorf("provider: %w", context.Canceled)), errWikiCheckpointYield)
	failure := errors.New("invalid API key")
	require.ErrorIs(t, wikiCheckpointError(ctx, failure), failure)
	require.ErrorIs(t, wikiCheckpointMapError(ctx, context.Canceled, failure), failure)
	require.ErrorIs(t, wikiCheckpointMapError(ctx, failure, errWikiCheckpointYield), failure)
	require.ErrorIs(t, wikiCheckpointMapError(ctx, nil, errWikiCheckpointYield), errWikiCheckpointYield)
	ordinary, cancelOrdinary := context.WithCancel(context.Background())
	cancelOrdinary()
	require.ErrorIs(t, wikiCheckpointError(ordinary, context.Canceled), context.Canceled)
}

type checkpointChunksRepo struct {
	interfaces.ChunkRepository
	chunks []*types.Chunk
}

func (r *checkpointChunksRepo) ListChunksByKnowledgeID(context.Context, uint64, string) ([]*types.Chunk, error) {
	return r.chunks, nil
}
func (*checkpointChunksRepo) ListChunksByParentIDs(context.Context, uint64, []string) ([]*types.Chunk, error) {
	return nil, nil
}

func (r *checkpointChunksRepo) ListChunksByID(_ context.Context, tenantID uint64, ids []string) ([]*types.Chunk, error) {
	wanted := make(map[string]bool, len(ids))
	for _, id := range ids {
		wanted[id] = true
	}
	var chunks []*types.Chunk
	for _, chunk := range r.chunks {
		if chunk != nil && chunk.TenantID == tenantID && wanted[chunk.ID] {
			chunks = append(chunks, chunk)
		}
	}
	return chunks, nil
}

type checkpointKBService struct {
	interfaces.KnowledgeBaseService
	kb *types.KnowledgeBase
}

func (s *checkpointKBService) GetKnowledgeBaseByIDOnly(context.Context, string) (*types.KnowledgeBase, error) {
	return s.kb, nil
}
func (s *checkpointKBService) GetKnowledgeBaseByID(context.Context, string) (*types.KnowledgeBase, error) {
	return s.kb, nil
}

type checkpointPending struct {
	interfaces.TaskPendingOpsRepository
	rows     []*types.TaskPendingOp
	released []int64
	deleted  []int64
	failures int
}

func (r *checkpointPending) Enqueue(context.Context, *types.TaskPendingOp) error { return nil }

func (r *checkpointPending) PeekBatch(context.Context, string, string, string, int) ([]*types.TaskPendingOp, error) {
	return r.rows, nil
}
func (r *checkpointPending) DeleteByIDs(_ context.Context, ids []int64) error {
	r.deleted = append(r.deleted, ids...)
	remaining := r.rows[:0]
	for _, row := range r.rows {
		removed := false
		for _, id := range ids {
			removed = removed || row.ID == id
		}
		if !removed {
			remaining = append(remaining, row)
		}
	}
	r.rows = remaining
	return nil
}
func (r *checkpointPending) ReleaseByIDs(_ context.Context, ids []int64) error {
	r.released = append(r.released, ids...)
	return nil
}
func (r *checkpointPending) IncrFailCount(context.Context, int64) (int, error) {
	r.failures++
	return r.failures, nil
}
func (r *checkpointPending) PendingCount(context.Context, string, string, string) (int64, error) {
	return int64(len(r.rows)), nil
}

type checkpointFinalization struct {
	interfaces.KnowledgeRepository
	calls int
}

func (r *checkpointFinalization) FinalizeSubtask(context.Context, string) (int, bool, error) {
	r.calls++
	return 0, true, nil
}

func TestWikiCheckpointYieldReleasesClaimWithoutFailureOrCompletion(t *testing.T) {
	db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
	chunks := checkpointChunks(wikiCitationSlice + 1)
	op := WikiPendingOp{Op: WikiOpIngest, KnowledgeID: "doc", DocTitle: "title", Language: "en-US", dbID: 42}
	opBytes, err := json.Marshal(op)
	require.NoError(t, err)
	pending := &checkpointPending{rows: []*types.TaskPendingOp{{ID: 42, Op: WikiOpIngest, TenantID: 7, Payload: opBytes}}}
	queue := &wikiGuardTaskQueue{}
	finalization := &checkpointFinalization{}
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]}}`}, nil
	}}
	modelSvc := &checkpointModelService{model: &types.Model{ID: t.Name(), Name: "test", TenantID: 7}, chat: model}
	kb := &types.KnowledgeBase{ID: "kb", TenantID: 7, SummaryModelID: t.Name(), WikiConfig: &types.WikiConfig{}}
	kb.IndexingStrategy.WikiEnabled = true
	svc := &wikiIngestService{checkpointDB: db, modelService: modelSvc, kbService: &checkpointKBService{kb: kb}, pendingRepo: pending,
		task: queue, knowledgeSvc: &checkpointKnowledge{}, knowledgeRepo: finalization, chunkRepo: &checkpointChunksRepo{chunks: chunks}, wikiService: &checkpointWikiPages{pages: make(map[string]*types.WikiPage)}}
	payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb", Language: "en-US"}
	content := reconstructContent(chunks)
	scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, content, "title", types.ResolveLanguageName(context.Background(), op.Language), svc.newWikiBatchContext("kb", kb.WikiConfig))
	require.NoError(t, err)
	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "candidates", "complete", combinedExtraction{Entities: []extractedItem{{Name: "Alpha", Slug: "entity/alpha"}}}))
	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "summary", wikiCheckpointHash("- [[entity/alpha]] = Alpha\n"), "Summary\n\nComplete document summary"))
	body, err := json.Marshal(payload)
	require.NoError(t, err)
	require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(types.TypeWikiIngest, body)))
	require.Equal(t, []int64{42}, pending.released)
	require.Empty(t, pending.deleted)
	require.Zero(t, pending.failures)
	require.Zero(t, finalization.calls)
	require.Len(t, queue.tasks, 1)
	require.Equal(t, types.TypeWikiIngest, queue.tasks[0].Type())
}

func checkpointMap(pageCount int) wikiMapCheckpoint {
	saved := wikiMapCheckpoint{Result: docIngestResult{KnowledgeID: "doc", DocTitle: "title", Summary: "Complete summary"}}
	for i := 0; i < pageCount; i++ {
		slug := fmt.Sprintf("concept/item-%d", i)
		update := SlugUpdate{KnowledgeID: "doc", SourceRef: "doc", Slug: slug, Type: types.WikiPageTypeConcept,
			DocTitle: "title", Item: extractedItem{Slug: slug, Name: fmt.Sprintf("Item %d", i), Details: "Grounded source details"}}
		if i == 0 {
			update.Slug = "summary/doc"
			update.Type = types.WikiPageTypeSummary
			update.SummaryBody = "Complete document summary"
		} else {
			update.SourceChunks = []string{"chunk-000"}
			update.Item.SourceChunks = []string{"chunk-000"}
		}
		saved.Updates = append(saved.Updates, update)
		saved.Result.Pages = append(saved.Result.Pages, wikiIngestPageRef{Slug: update.Slug, Title: "title"})
	}
	return saved
}

func TestWikiCheckpointMapRejectsIncompleteOrForeignArtifacts(t *testing.T) {
	scope := &wikiCheckpointScope{7, "kb", "doc", "v1"}
	chunks := checkpointChunks(1)
	require.NoError(t, validateWikiMapCheckpoint(checkpointMap(2), scope, chunks))
	for _, mutate := range []func(*wikiMapCheckpoint){
		func(m *wikiMapCheckpoint) { m.Result.KnowledgeID = "other" },
		func(m *wikiMapCheckpoint) { m.Updates[1].KnowledgeID = "other" },
		func(m *wikiMapCheckpoint) { m.Updates[1].SourceRef = "other" },
		func(m *wikiMapCheckpoint) { m.Updates[1].SourceChunks = []string{"foreign-chunk"} },
		func(m *wikiMapCheckpoint) { m.Updates[1].Item.SourceChunks = []string{"foreign-chunk"} },
		func(m *wikiMapCheckpoint) { m.Updates[1].Item.Slug = "entity/other" },
		func(m *wikiMapCheckpoint) { m.Result.Pages = m.Result.Pages[:1] },
		func(m *wikiMapCheckpoint) { m.Result.Pages = append(m.Result.Pages, m.Result.Pages[0]) },
		func(m *wikiMapCheckpoint) { m.Updates = m.Updates[1:]; m.Result.Pages = m.Result.Pages[1:] },
		func(m *wikiMapCheckpoint) { m.Updates[0].SummaryBody = "" },
	} {
		saved := checkpointMap(2)
		mutate(&saved)
		require.Error(t, validateWikiMapCheckpoint(saved, scope, chunks))
	}
}

func TestWikiCheckpointWholeMapResumesAfterDatabaseReopen(t *testing.T) {
	path := filepath.Join(t.TempDir(), "checkpoints.sqlite")
	db := openWikiCheckpointDB(t, path)
	chunks := checkpointChunks(1)
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		t.Error("completed map must not call the model again")
		return nil, errors.New("unexpected model call")
	}}
	modelSvc := &checkpointModelService{model: &types.Model{ID: t.Name(), Name: "test", TenantID: 7}, chat: model}
	svc := &wikiIngestService{checkpointDB: db, modelService: modelSvc, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &checkpointChunksRepo{chunks: chunks}}
	op := WikiPendingOp{Op: WikiOpIngest, KnowledgeID: "doc", Language: "en-US", dbID: 42}
	payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb"}
	batch := svc.newWikiBatchContext("kb", &types.WikiConfig{})
	scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, reconstructContent(chunks), "title", types.ResolveLanguageName(context.Background(), op.Language), batch)
	require.NoError(t, err)
	saved := checkpointMap(2)
	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "map", "complete", saved))
	sqlDB, _ := db.DB()
	require.NoError(t, sqlDB.Close())
	svc.checkpointDB = openWikiCheckpointDB(t, path)
	result, updates, err := svc.mapOneDocument(context.Background(), model, payload, op, batch)
	require.NoError(t, err)
	require.Equal(t, saved.Result.Pages, result.Pages)
	require.Equal(t, saved.Updates, updates)
}

func TestWikiCheckpointPageYieldKeepsClaimAndFinalizationSlot(t *testing.T) {
	path := filepath.Join(t.TempDir(), "checkpoints.sqlite")
	db := openWikiCheckpointDB(t, path)
	chunks := checkpointChunks(1)
	op := WikiPendingOp{Op: WikiOpIngest, KnowledgeID: "doc", DocTitle: "title", Language: "en-US", dbID: 42}
	opBytes, err := json.Marshal(op)
	require.NoError(t, err)
	pending := &checkpointPending{rows: []*types.TaskPendingOp{{ID: 42, Op: WikiOpIngest, TenantID: 7, Payload: opBytes}}}
	queue := &wikiGuardTaskQueue{}
	finalization := &checkpointFinalization{}
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		return &types.ChatResponse{Content: "SUMMARY: Complete page\n\nGrounded page content"}, nil
	}}
	modelSvc := &checkpointModelService{model: &types.Model{ID: t.Name(), Name: "test", TenantID: 7}, chat: model}
	kb := &types.KnowledgeBase{ID: "kb", TenantID: 7, SummaryModelID: t.Name(), WikiConfig: &types.WikiConfig{}}
	kb.IndexingStrategy.WikiEnabled = true
	pages := &checkpointWikiPages{pages: make(map[string]*types.WikiPage)}
	svc := &wikiIngestService{checkpointDB: db, modelService: modelSvc, kbService: &checkpointKBService{kb: kb}, pendingRepo: pending,
		task: queue, knowledgeSvc: &checkpointKnowledge{}, knowledgeRepo: finalization, chunkRepo: &checkpointChunksRepo{chunks: chunks}, wikiService: pages}
	payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb", Language: "en-US"}
	scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, reconstructContent(chunks), "title", types.ResolveLanguageName(context.Background(), op.Language), svc.newWikiBatchContext("kb", kb.WikiConfig))
	require.NoError(t, err)
	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "map", "complete", checkpointMap(wikiPageSlice+1)))
	body, err := json.Marshal(payload)
	require.NoError(t, err)
	require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(types.TypeWikiIngest, body)))
	require.Equal(t, wikiPageSlice, pages.writes)
	require.Equal(t, []int64{42}, pending.released)
	require.Empty(t, pending.deleted)
	require.Zero(t, pending.failures)
	require.Zero(t, finalization.calls)
	require.Len(t, queue.tasks, 1)
	require.Equal(t, types.TypeWikiIngest, queue.tasks[0].Type())
	// Simulate a worker restart after the first bounded page slice. The second
	// run must finish only the missing page and release the finalizing slot once.
	sqlDB, _ := db.DB()
	require.NoError(t, sqlDB.Close())
	svc.checkpointDB = openWikiCheckpointDB(t, path)
	require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(types.TypeWikiIngest, body)))
	require.Equal(t, wikiPageSlice+1, pages.writes)
	require.Equal(t, []int64{42}, pending.deleted)
	require.Zero(t, pending.failures)
	require.Equal(t, 1, finalization.calls)
	for _, page := range pages.pages {
		require.Equal(t, types.WikiPageStatusPublished, page.Status)
	}
}

type checkpointFailingKnowledge struct {
	interfaces.KnowledgeService
	err error
}

func (s *checkpointFailingKnowledge) GetKnowledgeByIDOnly(context.Context, string) (*types.Knowledge, error) {
	return nil, s.err
}

func TestWikiCheckpointTransientSourceFailureCannotCompleteMapOrPage(t *testing.T) {
	want := errors.New("database temporarily unavailable")
	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite")), knowledgeSvc: &checkpointFailingKnowledge{err: want}}
	_, _, err := svc.mapOneDocument(context.Background(), nil, WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb"}, WikiPendingOp{KnowledgeID: "doc"}, &WikiBatchContext{})
	require.ErrorIs(t, err, want)
	batch := &WikiBatchContext{}
	batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "v1"})
	_, _, _, err = svc.reduceSlugUpdates(context.Background(), nil, "kb", "summary/doc", checkpointMap(1).Updates, 7, batch, nil)
	require.ErrorIs(t, err, want)
	var count int64
	require.NoError(t, svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Count(&count).Error)
	require.Zero(t, count)
}

func TestWikiCheckpointPublishFailureKeepsPendingOperation(t *testing.T) {
	db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
	chunks := checkpointChunks(1)
	op := WikiPendingOp{Op: WikiOpIngest, KnowledgeID: "doc", DocTitle: "title", Language: "en-US", dbID: 42}
	opBytes, err := json.Marshal(op)
	require.NoError(t, err)
	pending := &checkpointPending{rows: []*types.TaskPendingOp{{ID: 42, Op: WikiOpIngest, TenantID: 7, Payload: opBytes}}}
	queue := &wikiGuardTaskQueue{}
	finalization := &checkpointFinalization{}
	model := &wikiScriptedChat{id: t.Name()}
	modelSvc := &checkpointModelService{model: &types.Model{ID: t.Name(), Name: "test", TenantID: 7}, chat: model}
	kb := &types.KnowledgeBase{ID: "kb", TenantID: 7, SummaryModelID: t.Name(), WikiConfig: &types.WikiConfig{}}
	kb.IndexingStrategy.WikiEnabled = true
	pages := &checkpointWikiPages{pages: make(map[string]*types.WikiPage), publishErr: errors.New("publish transaction failed")}
	svc := &wikiIngestService{checkpointDB: db, modelService: modelSvc, kbService: &checkpointKBService{kb: kb}, pendingRepo: pending,
		task: queue, knowledgeSvc: &checkpointKnowledge{}, knowledgeRepo: finalization, chunkRepo: &checkpointChunksRepo{chunks: chunks}, wikiService: pages}
	payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb", Language: "en-US"}
	scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, reconstructContent(chunks), "title", types.ResolveLanguageName(context.Background(), op.Language), svc.newWikiBatchContext("kb", kb.WikiConfig))
	require.NoError(t, err)
	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "map", "complete", checkpointMap(1)))
	body, err := json.Marshal(payload)
	require.NoError(t, err)
	require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(types.TypeWikiIngest, body)))
	require.Empty(t, pending.deleted)
	require.Equal(t, []int64{42}, pending.released)
	require.Equal(t, 1, pending.failures)
	require.Zero(t, finalization.calls)
	require.Equal(t, types.WikiPageStatusDraft, pages.pages["summary/doc"].Status)
	pages.publishErr = nil
	require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(types.TypeWikiIngest, body)))
	require.Equal(t, 1, pages.writes, "publication retry must reuse the already written page")
	require.Equal(t, []int64{42}, pending.deleted)
	require.Equal(t, 1, finalization.calls)
	require.Equal(t, types.WikiPageStatusPublished, pages.pages["summary/doc"].Status)
}

type checkpointDeadLetter struct {
	interfaces.TaskDeadLetterRepository
	err error
}

func (r *checkpointDeadLetter) Insert(context.Context, *types.TaskDeadLetter) error { return r.err }

func TestWikiCheckpointDeadLetterFailurePreservesPendingAndFinalization(t *testing.T) {
	want := errors.New("archive transaction failed")
	pending := &checkpointPending{failures: wikiMaxFailRetries}
	finalization := &checkpointFinalization{}
	svc := &wikiIngestService{pendingRepo: pending, deadLetterRepo: &checkpointDeadLetter{err: want}, knowledgeRepo: finalization}
	err := svc.requeueFailedOps(context.Background(), WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb"}, []WikiPendingOp{{Op: WikiOpIngest, KnowledgeID: "doc", dbID: 42}})
	require.ErrorIs(t, err, want)
	require.Empty(t, pending.deleted)
	require.Zero(t, finalization.calls)
}