audita/internal/framework/validators/llm_validators_test.go

package validators

import (
	"context"
	"errors"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/audita/internal/core/chunking"
	"gitea.maximumdirect.net/eric/audita/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/internal/core/schema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
	"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
)

type fakeStructuredLLMClient struct {
	responses []LLMValidationResponse
	err       error
	calls     []StructuredCompletionRequest
}

type sleepingValidationClient struct {
	inFlight    int32
	maxInFlight int32
	entered     chan struct{}
	release     chan struct{}
}

type boundedScheduler struct {
	permits chan struct{}
}

type captureValidationDiagnosticsWriter struct {
	lastRequestMetadata any
}

func (w *captureValidationDiagnosticsWriter) WriteInteraction(stage string, requestMetadata any, requestPayload any, responsePayload any, errorPayload any) (InteractionArtifacts, error) {
	_ = stage
	_ = requestPayload
	_ = responsePayload
	_ = errorPayload
	w.lastRequestMetadata = requestMetadata
	return InteractionArtifacts{}, nil
}

func newBoundedScheduler(max int) *boundedScheduler {
	return &boundedScheduler{permits: make(chan struct{}, max)}
}

func (s *boundedScheduler) Run(ctx context.Context, fn func(context.Context) error) error {
	select {
	case s.permits <- struct{}{}:
	case <-ctx.Done():
		return ctx.Err()
	}
	defer func() { <-s.permits }()
	return fn(ctx)
}

func (c *sleepingValidationClient) CompleteStructured(ctx context.Context, req StructuredCompletionRequest, out any) (StructuredCompletionResponse, error) {
	_ = req
	current := atomic.AddInt32(&c.inFlight, 1)
	for {
		prior := atomic.LoadInt32(&c.maxInFlight)
		if current <= prior || atomic.CompareAndSwapInt32(&c.maxInFlight, prior, current) {
			break
		}
	}
	if c.entered != nil {
		c.entered <- struct{}{}
	}
	select {
	case <-c.release:
	case <-ctx.Done():
		atomic.AddInt32(&c.inFlight, -1)
		return StructuredCompletionResponse{}, ctx.Err()
	}
	atomic.AddInt32(&c.inFlight, -1)
	target := out.(*LLMValidationResponse)
	*target = LLMValidationResponse{
		Validations: []LLMValidationDecision{
			{CorrectionIndex: 0, Approved: true, Confidence: 0.9, Reason: "ok"},
		},
	}
	return StructuredCompletionResponse{}, nil
}

func (f *fakeStructuredLLMClient) CompleteStructured(ctx context.Context, req StructuredCompletionRequest, out any) (StructuredCompletionResponse, error) {
	_ = ctx
	f.calls = append(f.calls, req)
	if f.err != nil {
		return StructuredCompletionResponse{}, f.err
	}
	if len(f.responses) == 0 {
		return StructuredCompletionResponse{}, errors.New("unexpected call")
	}
	resp := f.responses[0]
	f.responses = f.responses[1:]
	target, ok := out.(*LLMValidationResponse)
	if !ok {
		return StructuredCompletionResponse{}, errors.New("unexpected output type")
	}
	*target = resp
	return StructuredCompletionResponse{}, nil
}

func makeReq(candidates []proposals.EnrichedCorrectionProposal) Request {
	cfg := config.Default()
	cfg.ValidationMaxPromptTokens = 10000
	return Request{
		WorkingTranscript: &schema.Transcript{Segments: []schema.Segment{{ID: 1, Text: "There were gestures at the temple.", Categories: []string{"narration"}}}},
		CandidateProposal: candidates,
		ModuleKey:         "homophones",
		ModuleInstance:    "homophones",
		ReplacementPolicy: proposals.ReplacementPolicyRequireUnique,
		Config:            &cfg,
	}
}

func mk(index int, orig, corr string) proposals.EnrichedCorrectionProposal {
	return proposals.EnrichedCorrectionProposal{
		CorrectionProposal: proposals.CorrectionProposal{TargetSegmentID: 1, OriginalText: orig, CorrectedText: corr, Confidence: 0.9},
		ProposalMetadata:   proposals.ProposalMetadata{ProposalIndex: index, ModuleKey: "homophones", ModuleInstance: "homophones"},
	}
}

func TestChunkLLMValidationItemsOneSmallBatch(t *testing.T) {
	items := []LLMValidationItem{{CorrectionIndex: 0, OriginalText: "a", CorrectedText: "b", OriginalSegmentText: "x", CorrectedSegmentText: "y"}}
	batches, err := ChunkLLMValidationItems(items, 1000, chunking.NewSimpleTokenEstimator())
	if err != nil {
		t.Fatalf("Chunk error: %v", err)
	}
	if len(batches) != 1 {
		t.Fatalf("expected 1 batch, got %d", len(batches))
	}
}

func TestChunkLLMValidationItemsMultipleBatchesStableNoDropNoDup(t *testing.T) {
	items := []LLMValidationItem{
		{CorrectionIndex: 0, OriginalText: strings.Repeat("a", 20), CorrectedText: "b", OriginalSegmentText: "x", CorrectedSegmentText: "y"},
		{CorrectionIndex: 1, OriginalText: strings.Repeat("c", 20), CorrectedText: "d", OriginalSegmentText: "x", CorrectedSegmentText: "y"},
		{CorrectionIndex: 2, OriginalText: strings.Repeat("e", 20), CorrectedText: "f", OriginalSegmentText: "x", CorrectedSegmentText: "y"},
	}
	batches, err := ChunkLLMValidationItems(items, 40, chunking.NewSimpleTokenEstimator())
	if err != nil {
		t.Fatalf("Chunk error: %v", err)
	}
	if len(batches) < 2 {
		t.Fatalf("expected multiple batches, got %d", len(batches))
	}
	seen := make([]int, 0)
	for _, b := range batches {
		for _, it := range b.Items {
			seen = append(seen, it.CorrectionIndex)
		}
	}
	if len(seen) != 3 || seen[0] != 0 || seen[1] != 1 || seen[2] != 2 {
		t.Fatalf("unexpected ordering/drops/dups: %v", seen)
	}
}

func TestPromptBuildersContainRequiredContextAndInstructions(t *testing.T) {
	payload := []LLMValidationItem{{CorrectionIndex: 0, SegmentID: 1, OriginalText: "gestures", CorrectedText: "Jesters", OriginalSegmentText: "There were gestures", CorrectedSegmentText: "There were Jesters", Categories: []string{"narration"}}}
	tests := []struct {
		name    string
		build   func([]LLMValidationItem, string) ([]LLMMessage, error)
		mustHas []string
	}{
		{"spoken_form", BuildSpokenFormPlausibilityMessages, []string{"plausible spoken-form", "correction_index", "original_segment_text", "corrected_segment_text"}},
		{"meaning_reversal", BuildMeaningReversalMessages, []string{"meaning reversals", "correction_index", "original_segment_text", "corrected_segment_text"}},
		{"editorial", BuildEditorialMessages, []string{"acceptable editorial revision", "correction_index", "original_segment_text", "corrected_segment_text"}},
		{"grammar_review", BuildGrammarReviewMessages, []string{"acceptable editorial revision", "correction_index"}},
		{"spoken_word", BuildSpokenWordReviewMessages, []string{"acceptable editorial revision", "correction_index"}},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			msgs, err := tt.build(payload, "Discussion among party members in a dungeon.")
			if err != nil {
				t.Fatalf("build err: %v", err)
			}
			if len(msgs) != 2 {
				t.Fatalf("expected 2 messages, got %d", len(msgs))
			}
			combined := msgs[0].Content + "\n" + msgs[1].Content
			for _, needle := range tt.mustHas {
				if !strings.Contains(combined, needle) {
					t.Fatalf("expected prompt to contain %q", needle)
				}
			}
			for _, needle := range []string{
				"Transcript description (background context only):",
				"must not override the transcript content",
				"Do not invent corrections, facts, names, events, motivations, or speaker intent based on this description.",
			} {
				if !strings.Contains(combined, needle) {
					t.Fatalf("expected prompt to contain %q", needle)
				}
			}
		})
	}
}

func TestPromptBuildersOmitTranscriptDescriptionSectionWhenEmpty(t *testing.T) {
	payload := []LLMValidationItem{{CorrectionIndex: 0, SegmentID: 1, OriginalText: "gestures", CorrectedText: "Jesters", OriginalSegmentText: "There were gestures", CorrectedSegmentText: "There were Jesters"}}
	msgs, err := BuildSpokenFormPlausibilityMessages(payload, " ")
	if err != nil {
		t.Fatalf("build err: %v", err)
	}
	combined := msgs[0].Content + "\n" + msgs[1].Content
	if strings.Contains(combined, "Transcript description (background context only):") {
		t.Fatalf("did not expect empty transcript description section in prompt")
	}
}

func TestLLMBackedValidatorApprovalAndRejection(t *testing.T) {
	client := &fakeStructuredLLMClient{responses: []LLMValidationResponse{{Validations: []LLMValidationDecision{
		{CorrectionIndex: 0, Approved: true, Confidence: 0.9, Reason: "ok"},
		{CorrectionIndex: 1, Approved: false, Confidence: 0.95, Reason: "bad"},
	}}}}
	v, err := NewLLMBackedValidator("spoken_form_plausibility_review", LLMValidatorTypeSpokenFormPlausibility, "test-model")
	if err != nil {
		t.Fatalf("new validator error: %v", err)
	}
	req := makeReq([]proposals.EnrichedCorrectionProposal{mk(0, "gestures", "Jesters"), mk(1, "gestures", "Lyra")})
	req.LLMClient = client
	res, err := v.Validate(context.Background(), req)
	if err != nil {
		t.Fatalf("validate err: %v", err)
	}
	if len(res.Decisions) != 2 || !res.Decisions[0].Approved || res.Decisions[1].Approved {
		t.Fatalf("unexpected decisions: %+v", res.Decisions)
	}
	if len(client.calls) != 1 {
		t.Fatalf("expected one LLM call, got %d", len(client.calls))
	}
	if client.calls[0].ResponseSchema == nil {
		t.Fatalf("expected response schema on structured validation request")
	}
	want := responseschema.MustLookup(responseschema.ValidatorDecisionSetKey)
	gotSchema := client.calls[0].ResponseSchema
	if gotSchema.ID != want.ID || gotSchema.Version != want.Version || gotSchema.Name != want.Name || gotSchema.SHA256 != want.SHA256 {
		t.Fatalf("unexpected validator response schema metadata: got=%+v want=%+v", *gotSchema, want)
	}
}

func TestLLMBackedValidatorMalformedOutputFails(t *testing.T) {
	client := &fakeStructuredLLMClient{err: errors.New("malformed structured output")}
	v, _ := NewLLMBackedValidator("spoken_form_plausibility_review", LLMValidatorTypeSpokenFormPlausibility, "test-model")
	req := makeReq([]proposals.EnrichedCorrectionProposal{mk(0, "gestures", "Jesters")})
	req.LLMClient = client
	_, err := v.Validate(context.Background(), req)
	if err == nil || !strings.Contains(err.Error(), "completion failed") {
		t.Fatalf("expected malformed output error, got %v", err)
	}
}

func TestLLMBackedValidatorMissingDecisionFails(t *testing.T) {
	client := &fakeStructuredLLMClient{responses: []LLMValidationResponse{{Validations: []LLMValidationDecision{}}}}
	v, _ := NewLLMBackedValidator("spoken_form_plausibility_review", LLMValidatorTypeSpokenFormPlausibility, "test-model")
	req := makeReq([]proposals.EnrichedCorrectionProposal{mk(0, "gestures", "Jesters")})
	req.LLMClient = client
	_, err := v.Validate(context.Background(), req)
	if err == nil || !strings.Contains(err.Error(), "response invalid") {
		t.Fatalf("expected missing decision error, got %v", err)
	}
}

func TestLLMBackedValidatorDuplicateDecisionFails(t *testing.T) {
	client := &fakeStructuredLLMClient{responses: []LLMValidationResponse{{Validations: []LLMValidationDecision{
		{CorrectionIndex: 0, Approved: true, Confidence: 0.9, Reason: "ok"},
		{CorrectionIndex: 0, Approved: false, Confidence: 0.9, Reason: "dup"},
	}}}}
	v, _ := NewLLMBackedValidator("spoken_form_plausibility_review", LLMValidatorTypeSpokenFormPlausibility, "test-model")
	req := makeReq([]proposals.EnrichedCorrectionProposal{mk(0, "gestures", "Jesters")})
	req.LLMClient = client
	_, err := v.Validate(context.Background(), req)
	if err == nil || !strings.Contains(err.Error(), "duplicate") {
		t.Fatalf("expected duplicate decision error, got %v", err)
	}
}

func TestLLMBackedValidatorUnknownProposalIndexFails(t *testing.T) {
	client := &fakeStructuredLLMClient{responses: []LLMValidationResponse{{Validations: []LLMValidationDecision{{CorrectionIndex: 99, Approved: true, Confidence: 0.9, Reason: "unknown"}}}}}
	v, _ := NewLLMBackedValidator("spoken_form_plausibility_review", LLMValidatorTypeSpokenFormPlausibility, "test-model")
	req := makeReq([]proposals.EnrichedCorrectionProposal{mk(0, "gestures", "Jesters")})
	req.LLMClient = client
	_, err := v.Validate(context.Background(), req)
	if err == nil || !strings.Contains(err.Error(), "unknown") {
		t.Fatalf("expected unknown index error, got %v", err)
	}
}

func TestLLMBackedValidatorRespectsSchedulerConcurrency(t *testing.T) {
	scheduler := newBoundedScheduler(2)
	client := &sleepingValidationClient{
		entered: make(chan struct{}, 16),
		release: make(chan struct{}),
	}
	v, err := NewLLMBackedValidator("spoken_form_plausibility_review", LLMValidatorTypeSpokenFormPlausibility, "test-model")
	if err != nil {
		t.Fatalf("new validator error: %v", err)
	}

	var wg sync.WaitGroup
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			req := makeReq([]proposals.EnrichedCorrectionProposal{mk(0, "gestures", "Jesters")})
			req.LLMClient = client
			req.Scheduler = scheduler
			if _, runErr := v.Validate(context.Background(), req); runErr != nil {
				t.Errorf("Validate error: %v", runErr)
			}
		}()
	}
	waitForValidationEntries(t, client.entered, 2)
	close(client.release)
	wg.Wait()

	if got := atomic.LoadInt32(&client.maxInFlight); got > 2 {
		t.Fatalf("expected scheduler cap <= 2, got %d", got)
	}
	if got := atomic.LoadInt32(&client.maxInFlight); got < 2 {
		t.Fatalf("expected observed concurrency of at least 2, got %d", got)
	}
}

func TestLLMBackedValidatorDiagnosticsIncludeSchemaMetadata(t *testing.T) {
	client := &fakeStructuredLLMClient{responses: []LLMValidationResponse{{Validations: []LLMValidationDecision{
		{CorrectionIndex: 0, Approved: true, Confidence: 0.9, Reason: "ok"},
	}}}}
	writer := &captureValidationDiagnosticsWriter{}
	v, err := NewLLMBackedValidator("spoken_form_plausibility_review", LLMValidatorTypeSpokenFormPlausibility, "test-model")
	if err != nil {
		t.Fatalf("new validator error: %v", err)
	}
	req := makeReq([]proposals.EnrichedCorrectionProposal{mk(0, "gestures", "Jesters")})
	req.LLMClient = client
	req.DiagnosticsWriter = writer
	_, err = v.Validate(context.Background(), req)
	if err != nil {
		t.Fatalf("validate error: %v", err)
	}
	metadata, ok := writer.lastRequestMetadata.(map[string]any)
	if !ok {
		t.Fatalf("expected metadata map, got %T", writer.lastRequestMetadata)
	}
	schemaMap, ok := metadata["response_schema"].(map[string]any)
	if !ok {
		t.Fatalf("expected response_schema map, got %T", metadata["response_schema"])
	}
	want := responseschema.MustLookup(responseschema.ValidatorDecisionSetKey)
	if schemaMap["id"] != want.ID || schemaMap["version"] != want.Version || schemaMap["name"] != want.Name || schemaMap["sha256"] != want.SHA256 {
		t.Fatalf("unexpected diagnostics schema metadata: got=%v want=%+v", schemaMap, want)
	}
}

func waitForValidationEntries(t *testing.T, entered <-chan struct{}, want int) {
	t.Helper()
	deadline := time.Now().Add(300 * time.Millisecond)
	got := 0
	for got < want && time.Now().Before(deadline) {
		select {
		case <-entered:
			got++
		default:
			runtime.Gosched()
		}
	}
	if got < want {
		t.Fatalf("timed out waiting for %d validator calls to enter (got %d)", want, got)
	}
}