package service

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

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

func TestWikiCandidateRepairClassifiesStrictOutputWithoutModelValues(t *testing.T) {
	for _, fixture := range []struct{ content, code, field string }{
		{`{"entities":[`, "invalid_json_syntax", ""},
		{`{"entities":"PRIVATE_VALUE","concepts":[]}`, "json_type_mismatch", ""},
		{`null`, "missing_entities_array", "entities"},
		{`{"concepts":[]}`, "missing_entities_array", "entities"},
		{`{"entities":[],"concepts":null}`, "missing_concepts_array", "concepts"},
		{`{"entities":[{"name":" ","slug":"entity/PRIVATE_VALUE"}],"concepts":[]}`, "candidate_name_empty", "entities.name"},
		{`{"entities":[],"concepts":[{"name":"PRIVATE_VALUE","slug":"concept/ "}]}`, "candidate_slug_invalid", "concepts.slug"},
		{`{"entities":[{"name":"PRIVATE_VALUE","slug":"concept/private"}],"concepts":[]}`, "candidate_slug_invalid", "entities.slug"},
		{`{"entities":[{"name":"PRIVATE_VALUE","slug":"entity/private","aliases":17}],"concepts":[]}`, "json_type_mismatch", ""},
	} {
		t.Run(fixture.code+fixture.field, func(t *testing.T) {
			result, validation := parseWikiCandidateOutput(fixture.content)
			require.NotNil(t, validation)
			require.Empty(t, result.Entities)
			require.Empty(t, result.Concepts)
			require.Equal(t, fixture.code, validation.Code)
			require.Equal(t, fixture.field, validation.Field)
			require.NotContains(t, validation.String(), "PRIVATE_VALUE")
			require.NotContains(t, wikiCandidateCorrectionInstruction(*validation), "PRIVATE_VALUE")
			if fixture.code == "invalid_json_syntax" || fixture.code == "json_type_mismatch" {
				require.Positive(t, validation.JSONOffset)
			}
		})
	}
}

func TestWikiCandidateRepairUsesExactSourceSharedContextAndOneFreshRequest(t *testing.T) {
	guard := &candidateRepairTestGuard{}
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()
	ctx = requestbudget.WithGuard(context.WithValue(ctx, types.TenantIDContextKey, uint64(71)), guard)
	deadline, _ := ctx.Deadline()
	source := "Valid\n" + strings.Repeat("original-source\n", 500)
	data := map[string]string{"Content": source, "Language": "English", "PreviousSlugs": "entity/known"}
	wantData := make(map[string]string)
	for key, value := range data {
		wantData[key] = value
	}
	calls := 0
	var first string
	model := &wikiScriptedChat{id: t.Name(), call: func(callCtx context.Context, messages []chat.Message) (*types.ChatResponse, error) {
		calls++
		require.Same(t, guard, requestbudget.FromContext(callCtx))
		actualDeadline, ok := callCtx.Deadline()
		require.True(t, ok)
		require.Equal(t, deadline, actualDeadline)
		require.NoError(t, callCtx.Err())
		require.Len(t, messages, 1)
		if calls == 1 {
			first = messages[0].Content
			require.Contains(t, first, source)
			return &types.ChatResponse{Content: `{"entities":"PRIVATE_REJECTED_OUTPUT","concepts":[]}`, FinishReason: "stop"}, nil
		}
		require.Equal(t, 2, calls)
		require.True(t, strings.HasPrefix(messages[0].Content, first+"\n\n"+wikiCandidateRepairInstruction))
		require.NotContains(t, messages[0].Content, "PRIVATE_REJECTED_OUTPUT")
		require.Contains(t, messages[0].Content, `"code":"json_type_mismatch"`)
		return &types.ChatResponse{Content: `{"entities":[{"name":"Valid","slug":"entity/valid"}],"concepts":[]}`, FinishReason: "stop"}, nil
	}}
	result, err := (&wikiIngestService{}).generateWikiExtraction(ctx, model, agent.WikiCandidateSlugPrompt, data)
	require.NoError(t, err)
	require.Equal(t, 2, calls)
	require.Len(t, result.Entities, 1)
	require.Equal(t, wantData, data)
	require.GreaterOrEqual(t, guard.checks, 4)
}

type candidateRepairTestGuard struct {
	checks  int
	stopped bool
}

func (g *candidateRepairTestGuard) Key() string { return "candidate-repair-offline" }
func (g *candidateRepairTestGuard) Check(context.Context) error {
	g.checks++
	if g.stopped {
		return requestbudget.ErrStopped
	}
	return nil
}
func (g *candidateRepairTestGuard) Reserve(ctx context.Context, _, _ string) error {
	return g.Check(ctx)
}

func TestWikiCandidateRepairCannotBypassCancellationOrTransportBudget(t *testing.T) {
	for _, stop := range []string{"cancel", "budget"} {
		t.Run(stop, func(t *testing.T) {
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			guard := &candidateRepairTestGuard{}
			ctx = requestbudget.WithGuard(ctx, guard)
			calls := 0
			model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
				calls++
				if stop == "cancel" {
					cancel()
				} else {
					guard.stopped = true
				}
				return &types.ChatResponse{Content: `{"entities":[],"concepts":null}`, FinishReason: "stop"}, nil
			}}
			result, err := (&wikiIngestService{}).generateWikiExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
			require.Error(t, err)
			require.Equal(t, 1, calls)
			require.Empty(t, result.Entities)
			if stop == "cancel" {
				require.ErrorIs(t, err, context.Canceled)
			} else {
				require.ErrorIs(t, err, requestbudget.ErrStopped)
			}
		})
	}
}

func TestWikiCandidateRepairProviderFailuresNeverStartAnotherCorrection(t *testing.T) {
	for _, when := range []string{"initial", "correction"} {
		for _, code := range []string{"401", "403", "429", "503"} {
			t.Run(when+code, func(t *testing.T) {
				ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(72))
				calls := 0
				model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
					calls++
					if when == "correction" && calls == 1 {
						return &types.ChatResponse{Content: `{"concepts":[]}`, FinishReason: "stop"}, nil
					}
					return nil, errors.New("API request failed with status " + code + ": PRIVATE_PROVIDER_DETAIL")
				}}
				key := wikiProviderCircuitKey(72, model.id)
				t.Cleanup(func() { wikiProviderBlocks.Delete(key); wikiProviderSlots.Delete(key) })
				result, err := (&wikiIngestService{}).generateWikiExtraction(ctx, model, "{{.Content}}", map[string]string{"Content": strings.Repeat("x", 4000)})
				require.Error(t, err)
				want := 1
				if when == "correction" {
					want = 2
				} else if code == "503" {
					want = wikiLLMMaxAttempts
				}
				require.Equal(t, want, calls)
				require.Empty(t, result.Entities)
			})
		}
	}
}

func TestWikiCandidateRepairPartialFailureCannotAdvanceWindowCheckpoint(t *testing.T) {
	svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "db"))}
	scope := &wikiCheckpointScope{7, "kb", "doc", "candidate-correction"}
	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":"WINDOW_000","slug":"entity/valid"}],"concepts":[]}`, FinishReason: "stop"}, nil
		}
		return &types.ChatResponse{Content: `{"entities":[],"concepts":[{"name":"PRIVATE_VALUE","slug":"entity/private"}]}`, FinishReason: "stop"}, nil
	}}
	entities, concepts, items, err := svc.extractWikiFullCandidates(context.Background(), model, "kb", wikiFullReadTestSource(2), "English", nil, &WikiBatchContext{}, scope)
	require.ErrorContains(t, err, "partition 2/2")
	require.ErrorContains(t, err, "candidate_slug_invalid")
	require.NotContains(t, err.Error(), "PRIVATE_VALUE")
	require.Equal(t, 4, calls)
	require.Nil(t, entities)
	require.Nil(t, concepts)
	require.Nil(t, items)
	var count int64
	require.NoError(t, svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Where("stage IN ?", []string{"read_candidate", "read_progress", "candidates", "map", "page"}).Count(&count).Error)
	require.Zero(t, count, "partial successful leaf must not become a completed window or progress")
}