package service

import (
	"context"
	"encoding/xml"
	"errors"
	"fmt"
	"strings"
	"sync/atomic"
	"testing"

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

type wikiScriptedChat struct {
	templateCaptureChatModel
	id   string
	call func(context.Context, []chat.Message) (*types.ChatResponse, error)
}

func TestWikiCitationTruncationRetriesTwoSmallerBatchesWithLocalHandles(t *testing.T) {
	s := &wikiIngestService{}
	var inputs []string
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		prompt := messages[0].Content
		inputs = append(inputs, prompt)
		if len(inputs) == 1 {
			return &types.ChatResponse{Content: `{"citations":{}}`, FinishReason: "length"}, nil
		}
		if strings.Contains(prompt, `<c id="c001"`) {
			t.Fatal("retry did not rebuild the local handle table")
		}
		return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[{"type":"concept","slug":"concept/new","name":"New","source_chunks":["c000"]}]}`}, nil
	}}
	chunks := []*types.Chunk{{ID: "first", Content: strings.Repeat("first", 600), ChunkIndex: 0}, {ID: "second", Content: strings.Repeat("second", 500), ChunkIndex: 1}}
	citations, newSlugs, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
	if err != nil || completed != 1 || len(inputs) != 3 || strings.Join(citations["entity/alpha"], ",") != "first,second" {
		t.Fatalf("failed bounded citation recovery: completed=%d calls=%d citations=%v err=%v", completed, len(inputs), citations, err)
	}
	if len(newSlugs) != 2 || strings.Join(newSlugs[0].SourceChunks, ",") != "first" || strings.Join(newSlugs[1].SourceChunks, ",") != "second" {
		t.Fatalf("new slug provenance lost: %+v", newSlugs)
	}
	if !strings.Contains(inputs[1], chunks[0].Content) || strings.Contains(inputs[1], chunks[1].Content) || !strings.Contains(inputs[2], chunks[1].Content) || strings.Contains(inputs[2], chunks[0].Content) {
		t.Fatal("retry did not partition all source chunks")
	}
}

func TestWikiCitationOversizedChunkRetryPreservesEveryRuneAndSourceID(t *testing.T) {
	s := &wikiIngestService{}
	content := strings.Repeat("前", 1200) + "\n" + strings.Repeat("后", 1200)
	var seen []string
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		prompt := messages[0].Content
		start, end := strings.Index(prompt, "<chunks>\n"), strings.Index(prompt, "\n</chunks>")
		var rendered struct {
			Chunks []struct {
				Content string `xml:",chardata"`
			} `xml:"c"`
		}
		if start < 0 || end < 0 {
			t.Fatal("missing chunks block")
		}
		if err := xml.Unmarshal([]byte(prompt[start:end+len("\n</chunks>")]), &rendered); err != nil {
			t.Fatal(err)
		}
		if len(rendered.Chunks) != 1 {
			t.Fatalf("unexpected retry chunks: %+v", rendered)
		}
		seen = append(seen, strings.TrimSuffix(strings.TrimPrefix(rendered.Chunks[0].Content, "\n"), "\n"))
		if len(seen) == 1 {
			return &types.ChatResponse{Content: `{"citations":`, FinishReason: "max_tokens"}, nil
		}
		return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[]}`}, nil
	}}
	chunk := &types.Chunk{ID: "source", Content: content}
	citations, _, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{chunk}, "Chinese", &WikiBatchContext{})
	if err != nil || completed != 1 || len(seen) != 3 || seen[1]+seen[2] != content || chunk.Content != content || strings.Join(citations["entity/alpha"], ",") != "source" {
		t.Fatalf("singleton split lost source or duplicated citation: calls=%d result=%v err=%v", len(seen), citations, err)
	}
}

func TestWikiCitationPartitionFailureCannotReturnPartialCoverage(t *testing.T) {
	for _, failure := range []struct{ name, output, reason string }{
		{"truncated", `{"citations":{}}`, "max_output_tokens"},
		{"malformed", `{"citations":`, "stop"},
		{"missing-schema", `{}`, "stop"},
		{"unknown-handle", `{"citations":{"entity/alpha":["c999"]}}`, "stop"},
	} {
		t.Run(failure.name, func(t *testing.T) {
			s := &wikiIngestService{}
			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":`, FinishReason: "length"}, nil
				}
				if calls == 2 {
					return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[]}`}, nil
				}
				return &types.ChatResponse{Content: failure.output, FinishReason: failure.reason}, nil
			}}
			chunks := []*types.Chunk{{ID: "first", Content: "first", ChunkIndex: 0}, {ID: "second", Content: "second", ChunkIndex: 1}}
			citations, slugs, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
			if err == nil || calls != 3 || completed != 0 || citations != nil || slugs != nil || !strings.Contains(err.Error(), "partition 2/2") {
				t.Fatalf("partial citation accepted or retry budget exceeded: calls=%d completed=%d err=%v", calls, completed, err)
			}
		})
	}
}

func TestWikiCitationDoesNotPartitionProviderFailures(t *testing.T) {
	for _, status := range []int{401, 402, 403, 429} {
		t.Run(fmt.Sprint(status), func(t *testing.T) {
			s := &wikiIngestService{}
			ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(77))
			var calls atomic.Int32
			model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
				calls.Add(1)
				return nil, fmt.Errorf("API request failed with status %d: private text", status)
			}}
			key := wikiProviderCircuitKey(77, model.id)
			t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
			chunks := []*types.Chunk{{ID: "first", Content: "first"}, {ID: "second", Content: "second"}}
			citations, _, completed, err := s.classifyChunkCitations(ctx, model, "entity/alpha", chunks, "English", &WikiBatchContext{})
			if err == nil || calls.Load() != 1 || citations != nil || completed != 0 || strings.Contains(err.Error(), "private text") {
				t.Fatalf("provider failure split or leaked: calls=%d err=%v", calls.Load(), err)
			}
		})
	}
}

func TestWikiCitationCancellationStopsBeforeRetryPartition(t *testing.T) {
	s := &wikiIngestService{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	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":`, FinishReason: "length"}, nil
		}
		cancel()
		return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[]}`}, nil
	}}
	citations, _, completed, err := s.classifyChunkCitations(ctx, model, "entity/alpha", []*types.Chunk{{ID: "first", Content: "first"}, {ID: "second", Content: "second"}}, "English", &WikiBatchContext{})
	if err == nil || !errors.Is(err, context.Canceled) || calls != 2 || completed != 0 || citations != nil {
		t.Fatalf("cancelled partition continued: calls=%d completed=%d err=%v", calls, completed, err)
	}
}

func TestWikiCitationTwelveBatchesCompleteOnlyAfterTruncatedBatchRecovers(t *testing.T) {
	s := &wikiIngestService{}
	var calls atomic.Int32
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls.Add(1)
		prompt := messages[0].Content
		if strings.Contains(prompt, "target-left") && strings.Contains(prompt, "target-right") {
			return &types.ChatResponse{Content: `{"citations":`, FinishReason: "length"}, nil
		}
		handles := `["c000"]`
		if strings.Contains(prompt, `<c id="c001"`) {
			handles = `["c000","c001"]`
		}
		return &types.ChatResponse{Content: `{"citations":{"entity/alpha":` + handles + `},"new_slugs":[]}`}, nil
	}}
	var chunks []*types.Chunk
	var expected []string
	for i := 0; i < 12; i++ {
		for j := 0; j < 2; j++ {
			id := fmt.Sprintf("chunk-%02d-%d", i, j)
			prefix := id
			if i == 9 {
				prefix = []string{"target-left", "target-right"}[j]
			}
			chunks = append(chunks, &types.Chunk{ID: id, Content: prefix + strings.Repeat("x", 6000-len(prefix)), ChunkIndex: 2*i + j})
			expected = append(expected, id)
		}
	}
	citations, _, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
	if err != nil || calls.Load() != 14 || completed != 12 || strings.Join(citations["entity/alpha"], ",") != strings.Join(expected, ",") {
		t.Fatalf("batch 10/12 recovery lost coverage: calls=%d completed=%d citations=%v err=%v", calls.Load(), completed, citations, err)
	}
}

func TestWikiCitationRepeatedTruncationStopsAtFirstFailedPartition(t *testing.T) {
	s := &wikiIngestService{}
	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":`, FinishReason: "length"}, nil
	}}
	chunks := []*types.Chunk{{ID: "first", Content: strings.Repeat("x", 3000)}, {ID: "second", Content: strings.Repeat("y", 3000)}}
	citations, _, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
	if err == nil || calls.Load() != 2 || completed != 0 || citations != nil || !strings.Contains(err.Error(), "partition 1/2") {
		t.Fatalf("persistent truncation exceeded retry bound: calls=%d completed=%d err=%v", calls.Load(), completed, err)
	}
}

func (m *wikiScriptedChat) GetModelID() string { return m.id }
func (m *wikiScriptedChat) Chat(ctx context.Context, messages []chat.Message, _ *chat.ChatOptions) (*types.ChatResponse, error) {
	return m.call(ctx, messages)
}

func TestWikiExtractionSplitsTruncatedOutputWithoutDroppingEitherHalf(t *testing.T) {
	s := &wikiIngestService{}
	var inputs []string
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		inputs = append(inputs, messages[0].Content)
		if len(inputs) == 1 {
			return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
		}
		if len(inputs) == 2 {
			return &types.ChatResponse{Content: `{"entities":[{"name":"Alpha","slug":"entity/alpha"}],"concepts":[]}`}, nil
		}
		return &types.ChatResponse{Content: `{"entities":[],"concepts":[{"name":"Beta","slug":"concept/beta"}]}`}, nil
	}}
	content := "Alpha " + strings.Repeat("a", 1194) + "\n" + strings.Repeat("b", 1200)
	result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": content})
	if err != nil {
		t.Fatal(err)
	}
	if len(inputs) != 3 || inputs[1]+inputs[2] != content || len(result.Entities) != 1 || len(result.Concepts) != 1 {
		t.Fatalf("partition recovery lost data: calls=%d entities=%d concepts=%d", len(inputs), len(result.Entities), len(result.Concepts))
	}
}

func TestWikiExtractionInvalidJSONIsBoundedAndDoesNotReturnPartialOutput(t *testing.T) {
	s := &wikiIngestService{}
	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: `{"entities":[`}, nil
	}}
	result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
	if err == nil || calls.Load() != 4 || len(result.Entities) != 0 {
		t.Fatalf("invalid output must fail after first failing half, calls=%d err=%v", calls.Load(), err)
	}
}

func TestWikiExtractionLargeTruncatedHalfGetsOneAdditionalSplit(t *testing.T) {
	s := &wikiIngestService{}
	var leaves []string
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		content := messages[0].Content
		if len([]rune(content)) > 8000 {
			return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
		}
		leaves = append(leaves, content)
		return &types.ChatResponse{Content: `{"entities":[{"name":"Valid","slug":"entity/valid"}],"concepts":[]}`, FinishReason: "stop"}, nil
	}}
	source := "Valid" + strings.Repeat("甲", 7995) + "Valid" + strings.Repeat("乙", 7995) + "Valid" + strings.Repeat("丙", 7995) + "Valid" + strings.Repeat("丁", 7995)
	result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": source})
	require.NoError(t, err)
	require.Equal(t, 7, calls)
	require.Equal(t, source, strings.Join(leaves, ""), "every source character must reach a successful leaf")
	require.Len(t, leaves, 4)
	require.Len(t, result.Entities, 1, "identical fresh candidates coalesce only after all leaves succeed")
}

func TestWikiExtractionAdditionalSplitIsBoundedAndCannotPublishPartialResult(t *testing.T) {
	for _, lastFailure := range []string{"length", "invalid"} {
		t.Run(lastFailure, func(t *testing.T) {
			s := &wikiIngestService{}
			calls := 0
			model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
				calls++
				if calls == 3 {
					return &types.ChatResponse{Content: `{"entities":[{"name":"x","slug":"entity/valid"}],"concepts":[]}`}, nil
				}
				if calls == 4 && lastFailure == "invalid" {
					return &types.ChatResponse{Content: `{"entities":[`}, nil
				}
				return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
			}}
			result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 32000)})
			require.Error(t, err)
			expected := 4
			if lastFailure == "invalid" {
				expected++
			}
			require.Equal(t, expected, calls)
			require.Empty(t, result.Entities, "a successful first leaf must not publish when its sibling fails")
		})
	}
}

func TestWikiExtractionNestedProviderFailureStopsFurtherPartitions(t *testing.T) {
	s := &wikiIngestService{}
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(77))
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
		calls++
		if calls < 3 {
			return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
		}
		return nil, errors.New("API request failed with status 429: private provider detail")
	}}
	key := wikiProviderCircuitKey(77, model.id)
	t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
	result, err := s.generateWikiExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 32000)})
	require.Error(t, err)
	require.Equal(t, 3, calls)
	require.Empty(t, result.Entities)
	require.NotContains(t, err.Error(), "private provider detail")
}

func TestWikiExtractionInvalidCandidatesUseOneWholeSourceCorrection(t *testing.T) {
	for _, invalid := range []string{
		`{"entities":[{"name":" ","slug":"entity/private"}],"concepts":[]}`,
		`{"entities":[{"name":"Private","slug":"concept/private"}],"concepts":[]}`,
		`{"entities":[],"concepts":[{"name":"Private","slug":"concept/ "}]}`,
	} {
		t.Run(invalid, func(t *testing.T) {
			s := &wikiIngestService{}
			var inputs []string
			model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
				inputs = append(inputs, messages[0].Content)
				if len(inputs) == 1 {
					return &types.ChatResponse{Content: invalid}, nil
				}
				return &types.ChatResponse{Content: `{"entities":[{"name":"Valid","slug":"entity/valid"}],"concepts":[]}`}, nil
			}}
			source := "Valid " + strings.Repeat("a", 1194) + "\n" + strings.Repeat("b", 1200)
			result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": source})
			if err != nil || len(inputs) != 2 || !strings.HasPrefix(inputs[1], source+"\n\n"+wikiCandidateRepairInstruction) || len(result.Entities) != 1 {
				t.Fatalf("invalid candidate repair lost source or exceeded budget: calls=%d err=%v", len(inputs), err)
			}
		})
	}
}

func TestWikiExtractionInvalidCandidateInSecondHalfCannotPublishFirstHalf(t *testing.T) {
	s := &wikiIngestService{}
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
		calls++
		if calls == 1 {
			return &types.ChatResponse{Content: `{"entities":[`, FinishReason: "length"}, nil
		}
		if calls == 2 {
			return &types.ChatResponse{Content: `{"entities":[{"name":"x","slug":"entity/valid"}],"concepts":[]}`}, nil
		}
		return &types.ChatResponse{Content: `{"entities":[],"concepts":[{"name":"private-source-name","slug":"entity/private-source-slug"}]}`}, nil
	}}
	result, err := s.generateWikiExtraction(context.Background(), model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
	if err == nil || calls != 4 || len(result.Entities) != 0 || len(result.Concepts) != 0 || !strings.Contains(err.Error(), "partition 2/2") || !strings.Contains(err.Error(), `"field":"concepts.slug","index":0`) || strings.Contains(err.Error(), "private-source") {
		t.Fatalf("invalid candidate must preserve scope and report only field path: calls=%d err=%v", calls, err)
	}
}

func TestWikiExtractionDoesNotSplitOrImmediatelyRetryRateLimit(t *testing.T) {
	s := &wikiIngestService{}
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(77))
	var calls atomic.Int32
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
		calls.Add(1)
		return nil, errors.New("API request failed with status 429: private text")
	}}
	key := wikiProviderCircuitKey(77, model.id)
	t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
	_, err := s.generateWikiExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
	if err == nil || calls.Load() != 1 || strings.Contains(err.Error(), "private text") {
		t.Fatalf("rate limit retry storm or raw provider leak: calls=%d err=%v", calls.Load(), err)
	}
}

func TestWikiRejectsTruncatedEvenSyntacticallyValidOutput(t *testing.T) {
	for _, reason := range []string{"length", "max_tokens", "max_output_tokens", "incomplete"} {
		for _, content := range []string{`{"entities":[],"concepts":[]}`, "", " "} {
			err := validateWikiResponse(&types.ChatResponse{Content: content, FinishReason: reason})
			var outputErr *wikiOutputError
			if !errors.As(err, &outputErr) || !outputErr.incomplete {
				t.Fatalf("lost incomplete response reason %s: %v", reason, err)
			}
		}
	}
	if err := validateWikiResponse(&types.ChatResponse{Content: " ", FinishReason: "stop"}); err == nil {
		t.Fatal("accepted empty output")
	}
}

func TestWikiCitationFailureCannotBecomePartialSuccess(t *testing.T) {
	s := &wikiIngestService{}
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		if strings.Contains(messages[0].Content, "BROKEN") {
			return &types.ChatResponse{Content: `{"citations":`}, nil
		}
		return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[]}`}, nil
	}}
	chunks := []*types.Chunk{
		{ID: "first", Content: strings.Repeat("a", 12000), ChunkIndex: 0},
		{ID: "second", Content: "BROKEN", ChunkIndex: 1},
	}
	citations, slugs, _, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
	if err == nil || citations != nil || slugs != nil {
		t.Fatalf("partial classification returned as success: %v", err)
	}
}

func TestWikiCitationValidBatchesTranslateHandlesAndIgnoreNilChunks(t *testing.T) {
	s := &wikiIngestService{}
	model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
		return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]},"new_slugs":[]}`}, nil
	}}
	chunks := []*types.Chunk{nil, {ID: "first", Content: strings.Repeat("a", 12000), ChunkIndex: 0}, {ID: "second", Content: "b", ChunkIndex: 1}}
	citations, _, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{})
	if err != nil || completed != 2 || strings.Join(citations["entity/alpha"], ",") != "first,second" {
		t.Fatalf("bad citation coverage: completed=%d result=%v err=%v", completed, citations, err)
	}
}

func TestWikiCitationRejectsMissingSchemaAndUnknownSourceHandles(t *testing.T) {
	for _, output := range []string{`null`, `{}`, `{"citations":{"entity/alpha":["c999"]}}`, `{"citations":{},"new_slugs":[{"slug":"concept/new","name":"New","source_chunks":["c999"]}]}`} {
		s := &wikiIngestService{}
		model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, _ []chat.Message) (*types.ChatResponse, error) {
			return &types.ChatResponse{Content: output}, nil
		}}
		citations, _, completed, err := s.classifyChunkCitations(context.Background(), model, "entity/alpha", []*types.Chunk{{ID: "first", Content: "Source text"}}, "English", &WikiBatchContext{})
		if err == nil || citations != nil || completed != 0 {
			t.Fatalf("accepted citations without valid source coverage: completed=%d err=%v", completed, err)
		}
	}
}