package service

import (
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
	"unicode/utf8"

	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/require"
)

func wikiCatalogTestCandidates() []extractedItem {
	return []extractedItem{
		{Name: "Alpha", Slug: "entity/alpha", Aliases: []string{strings.Repeat("甲", 13000)}},
		{Name: "Beta", Slug: "entity/beta", Aliases: []string{strings.Repeat("乙", 13000)}},
	}
}

func wikiCatalogTestChunks(n int) []*types.Chunk {
	var chunks []*types.Chunk
	for i := 0; i < n; i++ {
		chunks = append(chunks, &types.Chunk{ID: fmt.Sprintf("source-%02d", i), ChunkIndex: i,
			Content: fmt.Sprintf("SOURCE_%02d\n", i) + strings.Repeat("原", 7000)})
	}
	return chunks
}

func TestWikiCandidateCatalogWholeEntriesStableAndInvalidated(t *testing.T) {
	entities := wikiCatalogTestCandidates()
	entities[0].Details = "Audit detail is retained but not rendered"
	entities[0].Description = "Audit description is retained but not rendered"
	entities[0].SourceChunks = []string{"original-source"}
	concepts := []extractedItem{{Name: "Topic", Slug: "concept/topic", Aliases: []string{"one\ntwo", "three"}}}
	shards, err := buildWikiCandidateShards(entities, concepts, wikiReadCandidateRunes)
	require.NoError(t, err)
	require.GreaterOrEqual(t, len(shards), 2)
	var got []wikiCandidateEntry
	for _, shard := range shards {
		require.Equal(t, shard.XML, renderWikiCandidateEntries(shard.Entries))
		require.LessOrEqual(t, utf8.RuneCountInString(shard.XML), wikiReadCandidateRunes)
		got = append(got, shard.Entries...)
	}
	require.ElementsMatch(t, []wikiCandidateEntry{{"entity", entities[0]}, {"entity", entities[1]}, {"concept", concepts[0]}}, got)
	reordered, err := buildWikiCandidateShards([]extractedItem{entities[1], entities[0]}, concepts, wikiReadCandidateRunes)
	require.NoError(t, err)
	require.Equal(t, shards, reordered)
	duplicates, err := buildWikiCandidateShards([]extractedItem{entities[0], entities[0]}, nil, wikiReadCandidateRunes)
	require.NoError(t, err)
	require.Len(t, duplicates, 2, "even duplicate input entries are not silently discarded by planning")
	require.NotEqual(t, duplicates[0].Hash, duplicates[1].Hash, "identical shards retain distinct occurrence checkpoints")
	chunks := wikiCatalogTestChunks(3)
	plan, err := buildWikiCitationPlan(entities, concepts, chunks, "Chinese")
	require.NoError(t, err)
	reorderedPlan, err := buildWikiCitationPlan([]extractedItem{entities[1], entities[0]}, concepts, []*types.Chunk{chunks[2], chunks[0], chunks[1]}, "Chinese")
	require.NoError(t, err)
	require.Equal(t, plan, reorderedPlan)
	changed, err := buildWikiCitationPlan(entities, concepts, chunks, "English")
	require.NoError(t, err)
	require.NotEqual(t, plan.Hash, changed.Hash)
	entities[0].Details += " changed audit field"
	changed, err = buildWikiCitationPlan(entities, concepts, chunks, "Chinese")
	require.NoError(t, err)
	require.NotEqual(t, plan.Hash, changed.Hash)
	chunks[2].Content += "tail condition"
	changedSource, err := buildWikiCitationPlan(entities, concepts, chunks, "Chinese")
	require.NoError(t, err)
	require.NotEqual(t, changed.Hash, changedSource.Hash)
	_, err = buildWikiCandidateShards([]extractedItem{{Name: "Oversize", Slug: "entity/oversize", Aliases: []string{strings.Repeat("名", wikiReadCandidateRunes)}}}, nil, wikiReadCandidateRunes)
	require.ErrorContains(t, err, "single-entry budget")
	_, err = buildWikiCandidateShards([]extractedItem{{Slug: "entity/no-name"}}, nil, wikiReadCandidateRunes)
	require.ErrorContains(t, err, "empty identity")
}

func TestWikiCandidateCatalogSourcesPreserveExactRangesAndWrappers(t *testing.T) {
	chunks := []*types.Chunk{
		{ID: "second", ChunkIndex: 1, Content: "\nsecond\n\n"},
		{ID: "first", ChunkIndex: 0, Content: strings.Repeat("原文😀\n", 9100) + " trailing qualification \n"},
	}
	sources, err := buildWikiCitationSources(chunks, maxRunesPerCitationBatch)
	require.NoError(t, err)
	require.GreaterOrEqual(t, len(sources), 3)
	joined := make(map[string]string)
	ends := make(map[string]int)
	for _, source := range sources {
		require.LessOrEqual(t, utf8.RuneCountInString(renderChunksXML(source.Batch)), maxRunesPerCitationBatch)
		for i, span := range source.Ranges {
			text := source.Batch.chunks[i].Content
			require.Equal(t, ends[span.ChunkID], span.StartRune)
			require.Equal(t, span.EndRune-span.StartRune, utf8.RuneCountInString(text))
			require.Equal(t, wikiCheckpointHash(text), span.TextHash)
			joined[span.ChunkID] += text
			ends[span.ChunkID] = span.EndRune
		}
	}
	for _, chunk := range chunks {
		require.Equal(t, chunk.Content, joined[chunk.ID])
	}
	require.NotEqual(t, sources[0].Hash, sources[1].Hash, "source ranges remain distinct even for repeating content")
	_, err = buildWikiCitationSources([]*types.Chunk{chunks[0], chunks[0]}, maxRunesPerCitationBatch)
	require.ErrorContains(t, err, "duplicate")
	_, err = buildWikiCitationSources(chunks, 1)
	require.ErrorContains(t, err, "wrapper")
}

func TestWikiCandidateCatalogAllPairsResumeAndLateEvidence(t *testing.T) {
	path := filepath.Join(t.TempDir(), "db")
	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
	scope := &wikiCheckpointScope{7, "kb", "doc", "pairs"}
	chunks := wikiCatalogTestChunks(9)
	plan, err := buildWikiCitationPlan(wikiCatalogTestCandidates(), nil, chunks, "Chinese")
	require.NoError(t, err)
	require.Len(t, plan.Shards, 2)
	require.Len(t, plan.Sources, 9)
	var calls, active, peak atomic.Int32
	var mu sync.Mutex
	seen := map[string]int{}
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls.Add(1)
		concurrent := active.Add(1)
		defer active.Add(-1)
		for old := peak.Load(); concurrent > old; old = peak.Load() {
			if peak.CompareAndSwap(old, concurrent) {
				break
			}
		}
		prompt := messages[0].Content
		if utf8.RuneCountInString(prompt) > wikiCitationRequestRunes {
			return nil, errors.New("unbounded request")
		}
		slug := "entity/beta"
		if strings.Contains(prompt, `name: "Alpha"`) {
			slug = "entity/alpha"
		}
		for i := range chunks {
			if strings.Contains(prompt, fmt.Sprintf("SOURCE_%02d", i)) {
				mu.Lock()
				seen[fmt.Sprintf("%s/%d", slug, i)]++
				mu.Unlock()
			}
		}
		if slug == "entity/alpha" && !strings.Contains(prompt, "SOURCE_08") {
			return &types.ChatResponse{Content: `{"citations":{}}`}, nil
		}
		return &types.ChatResponse{Content: fmt.Sprintf(`{"citations":{%q:["c000"]}}`, slug)}, nil
	}}
	got, _, completed, err := svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", scope)
	require.ErrorIs(t, err, errWikiCheckpointYield)
	require.Nil(t, got)
	require.Equal(t, 16, completed)
	require.Equal(t, int32(16), calls.Load())
	var progress wikiCitationProgress
	_, err = svc.loadWikiCheckpoint(context.Background(), scope, "citation_progress", plan.Hash, &progress)
	require.NoError(t, err)
	require.False(t, progress.StageComplete)
	require.Equal(t, 18, progress.TotalPairs)
	sqlDB, err := svc.checkpointDB.DB()
	require.NoError(t, err)
	require.NoError(t, sqlDB.Close())
	svc = &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
	got, _, completed, err = svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", scope)
	require.NoError(t, err)
	require.Equal(t, 18, completed)
	require.Equal(t, []string{chunks[8].ID}, got["entity/alpha"], "an early shard must inspect the last source batch")
	require.Len(t, got["entity/beta"], 9)
	require.Equal(t, int32(18), calls.Load())
	require.LessOrEqual(t, peak.Load(), int32(maxCitationBatchConcurrency))
	for key, count := range seen {
		require.Equal(t, 1, count, key)
	}
	require.Len(t, seen, 18)
	_, err = svc.loadWikiCheckpoint(context.Background(), scope, "citation_progress", plan.Hash, &progress)
	require.NoError(t, err)
	require.True(t, progress.StageComplete)
	_, _, _, err = svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", scope)
	require.NoError(t, err)
	require.Equal(t, int32(18), calls.Load())
	_, _, _, err = (&wikiIngestService{}).classifyWikiCitationPlan(context.Background(), model, plan, chunks, "Chinese", nil)
	require.ErrorContains(t, err, "需要持久断点")
	require.Equal(t, int32(18), calls.Load())
}

func TestWikiCandidateCatalogFailedPairKeepsSuccessfulSiblings(t *testing.T) {
	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
	scope := &wikiCheckpointScope{7, "kb", "doc", "failure"}
	chunks := wikiCatalogTestChunks(3)
	plan, err := buildWikiCitationPlan(wikiCatalogTestCandidates(), nil, chunks, "English")
	require.NoError(t, err)
	var failed atomic.Bool
	var mu sync.Mutex
	calls := map[string]int{}
	model := &wikiScriptedChat{id: t.Name(), call: func(ctx context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		prompt := messages[0].Content
		slug := "entity/beta"
		if strings.Contains(prompt, `name: "Alpha"`) {
			slug = "entity/alpha"
		}
		index := 0
		for i := range chunks {
			if strings.Contains(prompt, fmt.Sprintf("SOURCE_%02d", i)) {
				index = i
			}
		}
		key := fmt.Sprintf("%s/%d", slug, index)
		mu.Lock()
		calls[key]++
		mu.Unlock()
		if slug == "entity/alpha" && index == 1 && failed.CompareAndSwap(false, true) {
			deadline := time.Now().Add(3 * time.Second)
			for time.Now().Before(deadline) {
				var count int64
				if err := svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "citation_pair").Count(&count).Error; err != nil {
					return nil, err
				}
				if count > 0 {
					return nil, errors.New("invalid request")
				}
				select {
				case <-ctx.Done():
					return nil, ctx.Err()
				case <-time.After(time.Millisecond):
				}
			}
			return nil, errors.New("sibling checkpoint did not persist")
		}
		return &types.ChatResponse{Content: fmt.Sprintf(`{"citations":{%q:["c000"]}}`, slug)}, nil
	}}
	got, _, completed, err := svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "English", scope)
	require.Error(t, err)
	require.Nil(t, got)
	require.GreaterOrEqual(t, completed, 1)
	require.Less(t, completed, 6)
	var progress wikiCitationProgress
	_, loadErr := svc.loadWikiCheckpoint(context.Background(), scope, "citation_progress", plan.Hash, &progress)
	require.NoError(t, loadErr)
	require.False(t, progress.StageComplete)
	got, _, completed, err = svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "English", scope)
	require.NoError(t, err)
	require.Equal(t, 6, completed)
	require.Len(t, got["entity/alpha"], 3)
	for key, count := range calls {
		want := 1
		if key == "entity/alpha/1" {
			want = 2
		}
		require.Equal(t, want, count, key)
	}
}

func TestWikiCandidateCatalogPairIsolationAndEvidenceUnion(t *testing.T) {
	for _, corrupt := range []bool{false, true} {
		t.Run(fmt.Sprintf("corrupt-%t", corrupt), func(t *testing.T) {
			svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
			scope := &wikiCheckpointScope{7, "kb", "doc", "isolation"}
			entities := wikiCatalogTestCandidates()
			chunks := []*types.Chunk{{ID: "first", Content: "Alpha Beta first", ChunkIndex: 0}, {ID: "second", Content: "Alpha Beta second", ChunkIndex: 1}}
			plan, err := buildWikiCitationPlan(entities, nil, chunks, "English")
			require.NoError(t, err)
			require.Len(t, plan.Shards, 2)
			require.Len(t, plan.Sources, 1)
			for i, shard := range plan.Shards {
				result := citationBatchResult{Citations: map[string][]string{"entity/alpha": {"first"}}}
				if i == 1 && !corrupt {
					result.NewSlugs = []newSlugFromCitation{{Type: "entity", Name: "Alpha", Slug: "entity/alpha", SourceChunks: []string{"second"}}}
					result.Citations["entity/beta"] = []string{"second"}
				}
				key := wikiCheckpointHash([]any{wikiCandidateCatalogVersion, plan.Hash, shard.Hash, plan.Sources[0].Hash, "English"})
				require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "citation_pair", key, result))
			}
			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
				return nil, errors.New("cached validation must not invoke model repair")
			}}
			citations, declarations, completed, err := svc.classifyWikiCitationPlan(context.Background(), model, plan, chunks, "English", scope)
			if corrupt {
				require.ErrorContains(t, err, "undeclared slug")
				require.Nil(t, citations)
				require.Less(t, completed, 2)
				return
			}
			require.NoError(t, err)
			require.Equal(t, 2, completed)
			got, _, _ := mergeCitationsIntoItems(entities, nil, citations, declarations)
			require.Equal(t, []string{"first", "second"}, findBySlug(got, "entity/alpha").SourceChunks)
			require.Equal(t, []string{"second"}, findBySlug(got, "entity/beta").SourceChunks)
		})
	}
}

func TestWikiCandidateCatalogCompleteRequestGuard(t *testing.T) {
	var calls atomic.Int32
	inner := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		calls.Add(1)
		return &types.ChatResponse{}, nil
	}}
	model := wikiCitationBoundedChat{inner}
	_, err := model.Chat(context.Background(), []chat.Message{{Content: strings.Repeat("x", wikiCitationRequestRunes-1)}, {Content: "repair instruction"}}, nil)
	require.ErrorContains(t, err, "complete-request budget")
	require.Zero(t, calls.Load())
	require.Equal(t, inner.GetModelID(), model.GetModelID())
	_, err = model.Chat(context.Background(), []chat.Message{{Content: "bounded request"}}, nil)
	require.NoError(t, err)
	require.Equal(t, int32(1), calls.Load())
}