package validators

import (
	"encoding/json"
	"reflect"
	"slices"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/audita/v3/internal/core/chunking"
	"gitea.maximumdirect.net/eric/audita/v3/internal/core/glossary"
)

// Wrapping the same estimator exercises budgeting through exact serialization.
type serializedBudgetEstimator struct{ *chunking.SimpleTokenEstimator }

func TestComposedValidationBudgetMatchesSerializedPayload(t *testing.T) {
	items := make([]LLMValidationItem, 12)
	for i := range items {
		items[i] = LLMValidationItem{
			CorrectionIndex:      100 + i*7,
			OriginalText:         "one\u00a0two!",
			CorrectedText:        "one two",
			OriginalSegmentText:  "she said \"one\u00a0two!\"\n",
			CorrectedSegmentText: "she said <one & two> — there.",
			Categories:           []string{"say, then write", "repeat"},
		}
	}
	items[0].NeighboringSegments = []NeighboringSegment{
		{Position: "preceding", Text: "spaces\tand escapes\\", Categories: []string{"speaker line"}},
		{Position: "following", Text: "‘quoted’\u2028words"},
	}
	items[0].GlossaryEntries = []glossary.Entry{
		{ID: "name", Name: "Name?", Aliases: []string{"N. M.", "Name's"}},
		{ID: "two", Name: "One Two", Description: "some\u2003meaning…"},
	}
	items[1].CorrectedSegmentText = strings.Repeat("longer candidate. ", 12)

	simple := chunking.NewSimpleTokenEstimator()
	serialized := serializedBudgetEstimator{simple}
	full, err := estimateBatchTokens(serialized, items)
	if err != nil {
		t.Fatal(err)
	}
	for budget := 1; budget <= full; budget++ {
		got, rejected, err := prepareValidationBatches(slices.Clone(items), budget, simple)
		want, wantRejected, serializedErr := prepareValidationBatches(slices.Clone(items), budget, serialized)
		if err != nil || serializedErr != nil || !reflect.DeepEqual(got, want) || !reflect.DeepEqual(rejected, wantRejected) {
			t.Fatalf("budget %d: composed=%+v rejected=%+v err=%v; serialized=%+v rejected=%+v err=%v",
				budget, got, rejected, err, want, wantRejected, serializedErr)
		}
	}

	batches, err := ChunkLLMValidationItems(items, full, simple)
	if err != nil || len(batches) != 1 || batches[0].TokenCount != full {
		t.Fatalf("full batch: %+v, %v", batches, err)
	}
	variables, err := BuildValidationVariables(batches[0].Items, "")
	if err != nil {
		t.Fatal(err)
	}
	var payload []LLMValidationItem
	if err := json.Unmarshal([]byte(variables["PayloadJSON"]), &payload); err != nil {
		t.Fatal(err)
	}
	if len(payload) != len(items) {
		t.Fatalf("sent %d candidates, want %d", len(payload), len(items))
	}
	for i := range payload {
		if payload[i].CorrectionIndex != i+1 || batches[0].Items[i].CorrectionIndex != items[i].CorrectionIndex {
			t.Fatalf("candidate %d lost its local or run reference", i)
		}
	}
}