package service

import (
	"context"
	"encoding/json"
	"errors"
	"path/filepath"
	"strings"
	"testing"

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

func TestWikiCitationCandidateHandlesPreserveExactDurableIdentity(t *testing.T) {
	entities := []extractedItem{{Name: "Long name", Slug: "entity/a-very-long-uuid-9f2704c9-f7e7-4848-ad9c-44af056c4cf1", Description: "Do not edit entity/inside-description"}}
	concepts := []extractedItem{{Name: "Concept", Slug: "concept/a"}}
	xml := renderCandidateSlugsXML(entities, concepts)
	encoded, handles := encodeWikiCitationCandidates(xml)
	require.Contains(t, encoded, "- id: s001, type: entity")
	require.NotContains(t, encoded, "entity/inside-description", "generated descriptions are audit-only, not routing metadata")
	require.NotContains(t, encoded, entities[0].Slug)
	_, reversed := encodeWikiCitationCandidates(renderCandidateSlugsXML(nil, []extractedItem{concepts[0], entities[0]}))
	want, _ := handles.Handle(entities[0].Slug)
	got, _ := reversed.Handle(entities[0].Slug)
	require.Equal(t, want, got, "mapping must not depend on candidate order")
	batch := splitChunksIntoCitationBatches([]*types.Chunk{{ID: "source", Content: "facts"}})[0]
	parsed, err := parseWikiCitationBatch(`{"citations":{"s001":["c000"],"s000":["c000"]}}`, batch, wikiCandidateSlugSet(xml), handles)
	require.NoError(t, err)
	require.Equal(t, []string{"source"}, parsed.Citations[entities[0].Slug])
	require.Equal(t, []string{"source"}, parsed.Citations[concepts[0].Slug])
	for _, raw := range []string{
		`{"citations":{"s999":["c000"]}}`,
		`{"citations":{"c000":["c000"]}}`,
		`{"citations":{"s001":["s000"]}}`,
		`{"citations":{},"new_slugs":[{"type":"entity","slug":"s001","name":"wrong","source_chunks":["c000"]}]}`,
	} {
		_, err := parseWikiCitationBatch(raw, batch, wikiCandidateSlugSet(xml), handles)
		require.Error(t, err, raw)
	}
}

func TestWikiCitationUndeclaredSlugRepairPreservesEveryEdgeAndCheckpoint(t *testing.T) {
	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))}
	scope := &wikiCheckpointScope{7, "kb", "doc", "repair"}
	chunks := []*types.Chunk{{ID: "first", Content: "Alpha source; New fact one."}, {ID: "second", ChunkIndex: 1, Content: "New fact two."}}
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		if calls == 1 {
			require.Contains(t, messages[0].Content, "- id: s000")
			return &types.ChatResponse{Content: `{"citations":{"s000":["c000"],"concept/new":["c000","c001"]}}`}, nil
		}
		require.Equal(t, 2, calls, "repair must not become repeated reclassification")
		require.Contains(t, messages[0].Content, `"concept/new":["c000","c001"]`)
		return &types.ChatResponse{Content: `{"cannot_repair":false,"new_slugs":[{"type":"concept","slug":"concept/new","name":"New","description":"New source facts","details":"New fact one; New fact two","source_chunks":["c001","c000"],"evidence_quotes":[{"chunk":"c000","quote":"New fact one."},{"chunk":"c001","quote":"New fact two."}]}]}`}, nil
	}}
	citations, newSlugs, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
	require.NoError(t, err)
	require.Equal(t, 1, completed)
	require.Equal(t, 2, calls)
	require.Equal(t, []string{"first", "second"}, citations["concept/new"])
	require.Equal(t, []string{"first"}, citations["entity/alpha"])
	require.Len(t, newSlugs, 1)
	require.ElementsMatch(t, []string{"first", "second"}, newSlugs[0].SourceChunks)
	_, _, _, err = svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
	require.NoError(t, err)
	require.Equal(t, 2, calls, "validated repaired checkpoint must resume without another model call")
}

func TestWikiCitationDeclarationRepairCannotRewriteOrDropEvidence(t *testing.T) {
	batch := splitChunksIntoCitationBatches([]*types.Chunk{{ID: "source", Content: "Fact one."}, {ID: "other", ChunkIndex: 1, Content: "Other fact."}})[0]
	original := citationBatchResult{Citations: map[string][]string{"entity/known": {"other"}, "concept/new": {"source"}}}
	missing := map[string][]string{"concept/new": {"c000"}}
	valid := `{"cannot_repair":false,"new_slugs":[{"type":"concept","slug":"concept/new","name":"New","description":"Fact","details":"Fact one.","source_chunks":["c000"],"evidence_quotes":[{"chunk":"c000","quote":"Fact one."}]}]}`
	for name, raw := range map[string]string{
		"missing declaration": `{"new_slugs":[]}`,
		"declined":            strings.Replace(valid, `"cannot_repair":false`, `"cannot_repair":true`, 1),
		"renamed":             strings.Replace(valid, "concept/new", "concept/other", 1),
		"wrong type":          strings.Replace(valid, `"type":"concept"`, `"type":"entity"`, 1),
		"drop source":         strings.Replace(valid, `"source_chunks":["c000"]`, `"source_chunks":[]`, 1),
		"changed source":      strings.Replace(valid, `"source_chunks":["c000"]`, `"source_chunks":["c001"]`, 1),
		"unknown source":      strings.Replace(valid, `"source_chunks":["c000"]`, `"source_chunks":["c999"]`, 1),
		"fake quote":          strings.Replace(valid, `"quote":"Fact one."`, `"quote":"Made up"`, 1),
		"missing quote":       strings.Replace(valid, `"evidence_quotes":[{"chunk":"c000","quote":"Fact one."}]`, `"evidence_quotes":[]`, 1),
		"rewrite citations":   strings.Replace(valid, `"cannot_repair":false`, `"cannot_repair":false,"citations":{}`, 1),
		"trailing object":     valid + `{}`,
	} {
		t.Run(name, func(t *testing.T) {
			got, err := applyWikiCitationDeclarations(raw, original, missing, batch, map[string]bool{"entity/known": true})
			require.Error(t, err)
			require.Nil(t, got.Citations)
			require.Equal(t, []string{"source"}, original.Citations["concept/new"])
			require.Empty(t, original.NewSlugs)
		})
	}
}

func TestWikiCitationUnknownSourceAndHandleNeverTriggerDeclarationRepair(t *testing.T) {
	for _, raw := range []string{
		`{"citations":{"concept/new":["c999"]}}`,
		`{"citations":{"s999":["c000"]}}`,
		`{"citations":{"concept/new":[]}}`,
	} {
		calls := 0
		model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
			calls++
			return &types.ChatResponse{Content: raw}, nil
		}}
		_, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "one", Content: "Source one"}, {ID: "two", Content: "Source two"}}, "English", &WikiBatchContext{})
		require.Error(t, err)
		require.Zero(t, completed)
		if strings.Contains(raw, "c999") {
			require.Equal(t, 2, calls, "unknown source gets only one independent full regeneration")
		} else {
			require.Equal(t, 1, calls)
		}
	}
}

func TestWikiCitationUnknownSourceRegenerationPreservesWholeInputAndValidEdges(t *testing.T) {
	chunks := []*types.Chunk{{ID: "first", Content: "Alpha fact one."}, {ID: "second", ChunkIndex: 1, Content: "Alpha fact two."}}
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		if calls == 1 {
			return &types.ChatResponse{Content: `{"citations":{"s000":["c000","c999"]}}`}, nil
		}
		prompt := messages[0].Content
		require.Contains(t, prompt, "complete_citation_regeneration")
		require.Contains(t, prompt, `The ONLY allowed source IDs are: ["c000","c001"]`)
		require.Contains(t, prompt, `{"s000":["c000"]}`)
		require.Contains(t, prompt, chunks[0].Content)
		require.Contains(t, prompt, chunks[1].Content)
		require.NotContains(t, prompt, "c999", "invalid source identifier must not be proposed as a mapping")
		return &types.ChatResponse{Content: `{"citations":{"s000":["c000","c001"]},"new_slugs":[]}`}, nil
	}}
	got, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
	require.NoError(t, err)
	require.Equal(t, 2, calls)
	require.Equal(t, 1, completed)
	require.Equal(t, []string{"first", "second"}, got["entity/alpha"])
}

func TestWikiCitationSourceRegenerationCannotDropKnownEvidenceOrCascadeRetries(t *testing.T) {
	for name, reply := range map[string]string{
		"dropped-known-edge":      `{"citations":{"s000":["c001"]}}`,
		"still-unknown":           `{"citations":{"s000":["c000","c999"]}}`,
		"missing-new-declaration": `{"citations":{"s000":["c000"],"concept/new":["c001"]}}`,
		"malformed":               `{"citations":`,
	} {
		t.Run(name, func(t *testing.T) {
			calls := 0
			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
				calls++
				if calls == 1 {
					return &types.ChatResponse{Content: `{"citations":{"s000":["c000","c999"]}}`}, nil
				}
				return &types.ChatResponse{Content: reply}, nil
			}}
			got, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "first", Content: "Alpha fact one."}, {ID: "second", Content: "Alpha fact two."}}, "English", &WikiBatchContext{})
			require.Error(t, err)
			require.Nil(t, got)
			require.Zero(t, completed)
			require.Equal(t, 2, calls, "no bisection, declaration repair or provider retry can be stacked")
		})
	}
}

func TestWikiCitationSourceRegenerationKeepsUndeclaredSlugEvidence(t *testing.T) {
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		if calls == 1 {
			return &types.ChatResponse{Content: `{"citations":{"s000":["c999"],"concept/new":["c000"]}}`}, nil
		}
		require.Contains(t, messages[0].Content, `"concept/new":["c000"]`)
		return &types.ChatResponse{Content: `{"citations":{"s000":["c000"]}}`}, nil
	}}
	got, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "first", Content: "Alpha source."}}, "English", &WikiBatchContext{})
	require.ErrorContains(t, err, "dropped previously valid evidence")
	require.Nil(t, got)
	require.Zero(t, completed)
	require.Equal(t, 2, calls)
}

func TestWikiCitationUnknownSourceCannotRepairInvalidCandidateOrType(t *testing.T) {
	for _, raw := range []string{
		`{"citations":{"s999":["c999"],"s000":["c999"]}}`,
		`{"citations":{"s000":["c999"]},"new_slugs":[{"type":"invalid","slug":"concept/new","name":"New","source_chunks":["c000"]}]}`,
	} {
		calls := 0
		model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
			calls++
			return &types.ChatResponse{Content: raw}, nil
		}}
		_, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "first", Content: "Alpha source."}}, "English", &WikiBatchContext{})
		require.Error(t, err)
		require.Zero(t, completed)
		require.Equal(t, 1, calls)
	}
}

func TestWikiCitationDeclarationRepairHasOneCallAndNoSplitFallback(t *testing.T) {
	for _, providerFailure := range []bool{false, true} {
		calls := 0
		model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
			calls++
			if calls == 1 {
				return &types.ChatResponse{Content: `{"citations":{"concept/new":["c000"]}}`}, nil
			}
			if providerFailure {
				return nil, errors.New("HTTP 429 rate limit")
			}
			return &types.ChatResponse{Content: `{"new_slugs":[]}`}, nil
		}}
		_, _, completed, err := (&wikiIngestService{}).classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "one", Content: "Source one"}, {ID: "two", Content: "Source two"}}, "English", &WikiBatchContext{})
		require.Error(t, err)
		require.Zero(t, completed)
		require.Equal(t, 2, calls)
	}
}

func TestWikiCitationPersistedUnknownSlugDoesNotCallModel(t *testing.T) {
	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))}
	scope := &wikiCheckpointScope{7, "kb", "doc", "corrupt"}
	chunks := []*types.Chunk{{ID: "source", Content: "Source text"}}
	key := wikiCheckpointHash([]any{"entity/alpha", "English", wikiCheckpointChunks(chunks)})
	require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "citation", key, citationBatchResult{Citations: map[string][]string{"concept/undeclared": {"source"}}}))
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		t.Fatal("persisted corruption must not spend a model request")
		return nil, nil
	}}
	_, _, _, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
	require.ErrorContains(t, err, "invalid persisted citation")
	// Diagnostics must not contain private candidate text or original documents.
	encoded, _ := json.Marshal(err.Error())
	require.NotContains(t, string(encoded), "undeclared\"")
}