package service

import (
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"strings"
	"testing"

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

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

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

func pageRepairInput() map[string]string {
	return map[string]string{
		"HasAdditions": "1", "HasRetractions": "1", "PageSlug": "concept/alpha", "PageTitle": "Alpha", "PageType": "concept",
		"ExistingContent":      "Existing fact is not valid unless condition A holds.",
		"SharedSourceContexts": "Source context and uncertainty.", "NewContent": "Source A states 42, except under condition B.",
		"DeletedContent": "Obsolete assertion.", "RemainingSourcesContent": "Preserved source B.",
		"AvailableSlugs": "ref-1 (Alpha)", "Language": "English", "CustomInstructions": "Keep attribution.", "InstructionScope": "wiki_content",
	}
}

func TestWikiPageOutputRepairRegeneratesWholePageWithIdenticalScope(t *testing.T) {
	svc := &wikiIngestService{}
	input := pageRepairInput()
	before := make(map[string]string, len(input))
	for key, value := range input {
		before[key] = value
	}
	var requests [][]chat.Message
	var prefixes []string
	model := &wikiPageRepairChat{id: t.Name(), call: func(ctx context.Context, messages []chat.Message, opts *chat.ChatOptions) (*types.ChatResponse, error) {
		requests = append(requests, append([]chat.Message(nil), messages...))
		purpose, prefix := types.LLMCallMetadataFromContext(ctx)
		require.Equal(t, "wiki_page_modify", purpose)
		prefixes = append(prefixes, prefix)
		require.Equal(t, wikiLLMMaxTokens, opts.MaxTokens)
		if len(requests) == 1 {
			return &types.ChatResponse{Content: "PARTIAL_OUTPUT_MUST_NOT_BE_REUSED", FinishReason: "length"}, nil
		}
		return &types.ChatResponse{Content: "SUMMARY: Complete revision\n# Alpha\nNot valid unless condition A holds; source A states 42 except under condition B.", FinishReason: "stop"}, nil
	}}
	got, err := svc.generateWikiPageModification(context.Background(), model, input)
	require.NoError(t, err)
	require.Contains(t, got, "Complete revision")
	require.Len(t, requests, 2)
	require.Equal(t, before, input)
	require.Equal(t, requests[0][0], requests[1][0], "base system prompt must remain unchanged")
	require.Equal(t, requests[0][1].Content+"\n\n"+wikiPageRepairInstruction, requests[1][1].Content)
	require.Equal(t, prefixes[0], prefixes[1], "page source prefix remains reusable")
	require.NotContains(t, requests[1][1].Content, "PARTIAL_OUTPUT_MUST_NOT_BE_REUSED")
}

func TestWikiPageOutputRepairSuccessNeedsNoSecondCall(t *testing.T) {
	calls := 0
	model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
		calls++
		return &types.ChatResponse{Content: "SUMMARY: Complete\nComplete body", FinishReason: "stop"}, nil
	}}
	got, err := (&wikiIngestService{}).generateWikiPageModification(context.Background(), model, pageRepairInput())
	require.NoError(t, err)
	require.NotEmpty(t, got)
	require.Equal(t, 1, calls)
}

func TestWikiPageOutputRepairStopsAfterOneFailedRegeneration(t *testing.T) {
	for _, response := range []*types.ChatResponse{
		{Content: "still partial", FinishReason: "length"},
		{Content: "still partial", FinishReason: "max_tokens"},
		{Content: "still partial", FinishReason: "max_output_tokens"},
		{Content: "still partial", FinishReason: types.FinishReasonIncomplete},
		{Content: " ", FinishReason: "stop"},
		{Content: "WIKI_PAGE_REPAIR_UNSAFE", FinishReason: "stop"},
		{Content: "SUMMARY: Summary without body", FinishReason: "stop"},
		{Content: "SUMMARY: \nBody without required summary", FinishReason: "stop"},
		{Content: "Missing required summary line", FinishReason: "stop"},
		{Content: "I cannot complete the requested page.", FinishReason: "stop"},
		{Content: strings.Repeat("x", 12001), FinishReason: "stop"},
	} {
		t.Run(fmt.Sprintf("%s-%d", response.FinishReason, len(response.Content)), 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: "partial", FinishReason: "length"}, nil
				}
				return response, nil
			}}
			got, err := (&wikiIngestService{}).generateWikiPageModification(context.Background(), model, pageRepairInput())
			require.Error(t, err)
			require.Empty(t, got)
			require.Equal(t, 2, calls)
		})
	}
}

func TestWikiPageSummaryOnlyDoesNotOverwritePublishedPage(t *testing.T) {
	for _, repaired := range []string{"SUMMARY: Only a summary", "Missing required SUMMARY line", "I cannot complete the requested page.", "SUMMARY: \nMissing summary"} {
		t.Run(repaired, func(t *testing.T) {
			db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
			original := &types.WikiPage{ID: "page", TenantID: 7, KnowledgeBaseID: "kb", Slug: "concept/alpha", Title: "Alpha", PageType: types.WikiPageTypeConcept,
				Status: types.WikiPageStatusPublished, Content: "Original supported content", Summary: "Original summary", SourceRefs: types.StringArray{"old-doc"}}
			before := *original
			pages := &checkpointWikiPages{pages: map[string]*types.WikiPage{original.Slug: original}}
			svc := &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &wikiEvidenceChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}}
			batch := &WikiBatchContext{SlugTitleMany: func(context.Context, []string) map[string]string { return nil }}
			batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "unchanged-fingerprint"})
			calls := 0
			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
				calls++
				if calls == 1 {
					return &types.ChatResponse{Content: "SUMMARY: Incomplete first page", FinishReason: "stop"}, nil
				}
				return &types.ChatResponse{Content: repaired, FinishReason: "stop"}, nil
			}}
			update := SlugUpdate{KnowledgeID: "doc", SourceRef: "doc", Slug: original.Slug, Type: types.WikiPageTypeConcept, DocTitle: "Source", SourceChunks: []string{"source-chunk"}, Item: extractedItem{Name: "Alpha", Slug: original.Slug, Details: "New fact"}}
			changed, _, failed, err := svc.reduceSlugUpdates(context.Background(), model, "kb", original.Slug, []SlugUpdate{update}, 7, batch, nil)
			require.Error(t, err)
			require.False(t, changed)
			require.True(t, failed)
			require.Equal(t, 2, calls)
			require.Zero(t, pages.writes)
			require.Equal(t, before, *pages.pages[original.Slug])
			var checkpoints int64
			require.NoError(t, db.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "page").Count(&checkpoints).Error)
			require.Zero(t, checkpoints)
		})
	}
}

func TestWikiPageOutputRepairDoesNotMultiplyProviderFailures(t *testing.T) {
	for _, afterOutputFailure := range []bool{false, true} {
		for _, failure := range []string{"API request failed with status 401", "API request failed with status 402", "API request failed with status 403", "API request failed with status 429", "connection closed", "timeout"} {
			if failure == "timeout" && !afterOutputFailure {
				continue // The first request retains its existing transient transport policy.
			}
			t.Run(fmt.Sprintf("repair-%t-%s", afterOutputFailure, failure), func(t *testing.T) {
				ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(700))
				calls := 0
				model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
					calls++
					if afterOutputFailure && calls == 1 {
						return &types.ChatResponse{Content: "partial", FinishReason: "length"}, nil
					}
					return nil, errors.New(failure)
				}}
				key := wikiProviderCircuitKey(700, model.id)
				t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
				got, err := (&wikiIngestService{}).generateWikiPageModification(ctx, model, pageRepairInput())
				require.Error(t, err)
				require.Empty(t, got)
				expected := 1
				if afterOutputFailure {
					expected++
				}
				require.Equal(t, expected, calls)
			})
		}
	}
}

func TestWikiPageOutputRepairHonorsCancellation(t *testing.T) {
	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++
		cancel()
		return &types.ChatResponse{Content: "partial", FinishReason: "length"}, nil
	}}
	got, err := (&wikiIngestService{}).generateWikiPageModification(ctx, model, pageRepairInput())
	require.Error(t, err)
	require.Empty(t, got)
	require.Equal(t, 1, calls)
}

func TestWikiPageOutputRepairFailurePreservesPublishedPageAndCheckpoint(t *testing.T) {
	for _, repairsSuccessfully := range []bool{false, true} {
		t.Run(fmt.Sprint(repairsSuccessfully), func(t *testing.T) {
			db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
			original := &types.WikiPage{ID: "page", TenantID: 7, KnowledgeBaseID: "kb", Slug: "concept/alpha", Title: "Alpha", PageType: types.WikiPageTypeConcept,
				Status: types.WikiPageStatusPublished, Content: "Original supported content", Summary: "Original summary", SourceRefs: types.StringArray{"old-doc"}}
			before := *original
			pages := &checkpointWikiPages{pages: map[string]*types.WikiPage{original.Slug: original}}
			svc := &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &wikiEvidenceChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}}
			batch := &WikiBatchContext{SlugTitleMany: func(context.Context, []string) map[string]string { return nil }}
			batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "unchanged-fingerprint"})
			calls := 0
			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
				calls++
				if repairsSuccessfully && calls == 2 {
					return &types.ChatResponse{Content: "SUMMARY: Complete revised summary\nOriginal supported content. New fact has a condition.", FinishReason: "stop"}, nil
				}
				return &types.ChatResponse{Content: "partial", FinishReason: "length"}, nil
			}}
			update := SlugUpdate{KnowledgeID: "doc", SourceRef: "doc", Slug: original.Slug, Type: types.WikiPageTypeConcept, DocTitle: "Source", SourceChunks: []string{"source-chunk"}, Item: extractedItem{Name: "Alpha", Slug: original.Slug, Details: "New fact has a condition."}}
			changed, _, failed, err := svc.reduceSlugUpdates(context.Background(), model, "kb", original.Slug, []SlugUpdate{update}, 7, batch, nil)
			require.Equal(t, 2, calls)
			var checkpoints int64
			require.NoError(t, db.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "page").Count(&checkpoints).Error)
			if repairsSuccessfully {
				require.NoError(t, err)
				require.True(t, changed)
				require.False(t, failed)
				require.Equal(t, 1, pages.writes)
				require.Equal(t, int64(1), checkpoints)
				require.Contains(t, pages.pages[original.Slug].Content, "Original supported content")
				require.Contains(t, pages.pages[original.Slug].SourceRefs, "old-doc")
				require.Contains(t, pages.pages[original.Slug].SourceRefs, "doc")
			} else {
				require.Error(t, err)
				require.False(t, changed)
				require.True(t, failed)
				require.Zero(t, pages.writes)
				require.Zero(t, checkpoints)
				require.Equal(t, before, *pages.pages[original.Slug])
			}
		})
	}
}