package pipeline

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

func TestRunnerCorrectsExtractCandidatesIndependentlyPerChunk(t *testing.T) {
	prepared := preparedConcurrentPipeline(t, 2)
	lane := &prepared.Steps[0].lanes[0]
	lane.resolved.Extract.Retries = 1

	var mu sync.Mutex
	corrections := make(map[int]*contracts.SemanticCorrection)
	installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		chunkIndex := request.Chunk.Index
		if request.Correction != nil {
			correction, err := contracts.CloneSemanticCorrection(request.Correction)
			if err != nil {
				return erasedTypedResult{}, err
			}
			mu.Lock()
			corrections[chunkIndex] = correction
			mu.Unlock()
			candidate, err := contracts.NewModelCandidate([]byte(fmt.Sprintf("corrected-response-%d", chunkIndex)), contracts.CorrectionProtocolSingleResponseV1)
			if err != nil {
				return erasedTypedResult{}, err
			}
			return erasedTypedResult{Value: codecNotes{Items: []string{fmt.Sprintf("corrected-%d", chunkIndex)}}, ModelCandidate: candidate}, nil
		}
		candidate, err := contracts.NewModelCandidate([]byte(fmt.Sprintf("initial-response-%d", chunkIndex)), contracts.CorrectionProtocolSingleResponseV1)
		if err != nil {
			return erasedTypedResult{}, err
		}
		return erasedTypedResult{Value: codecNotes{Items: []string{fmt.Sprintf("invalid-%d", chunkIndex)}}, ModelCandidate: candidate}, nil
	})
	lane.extractValidators.validators[0].typedValidate = func(_ context.Context, _ any, target typedValidationTarget) (contracts.ValidationResult, error) {
		value := target.value.(codecNotes)
		if value.Items[0] == fmt.Sprintf("invalid-%d", target.chunk.Index) {
			return contracts.ValidationResult{Approved: false, ReasonCode: "incorrect_extract", Message: "candidate needs correction", CorrectionGuidance: fmt.Sprintf("correct chunk %d", target.chunk.Index)}, nil
		}
		return contracts.ValidationResult{Approved: true}, nil
	}

	debug := newCapturedDebugRecorder()
	output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2, Debug: debug})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if len(output.Rejected) != 0 || len(output.NormalizeOutputs) != 2 {
		t.Fatalf("run output = %#v, want corrected accepted extracts", output)
	}
	mu.Lock()
	defer mu.Unlock()
	if len(corrections) != 2 {
		t.Fatalf("corrections = %#v, want one correction per chunk", corrections)
	}
	for index := 0; index < 2; index++ {
		correction := corrections[index]
		if correction == nil || string(correction.AssistantResponse) != fmt.Sprintf("initial-response-%d", index) || !strings.Contains(correction.UserGuidance, fmt.Sprintf("correct chunk %d", index)) || !strings.Contains(correction.UserGuidance, "complete corrected replacement") || strings.Contains(correction.UserGuidance, "incorrect_extract") || strings.Contains(correction.UserGuidance, "candidate needs correction") {
			t.Fatalf("chunk %d correction = %#v, want its exact initial response and semantic replacement guidance only", index, correction)
		}
	}
	terminal := string(debug.json["extract/notes/chunk-000001/terminal.json"])
	if !strings.Contains(terminal, `"kind":"semantic_correction"`) || !strings.Contains(terminal, `"producer_attempt_count":2`) {
		t.Fatalf("extract terminal debug = %s, want correction provenance", terminal)
	}
	if strings.Contains(terminal, "initial-response-0") || strings.Contains(terminal, "corrected-response-0") || strings.Contains(terminal, "correct chunk 0") {
		t.Fatalf("extract terminal debug leaked correction content: %s", terminal)
	}
	manifest, err := json.Marshal(output.Manifest)
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(string(manifest), "initial-response-0") || strings.Contains(string(manifest), "corrected-response-0") || strings.Contains(string(manifest), "correct chunk 0") {
		t.Fatalf("manifest leaked correction content: %s", manifest)
	}
}

func containsValidationSummary(summaries []artifacts.ValidationSummary, status string) bool {
	for _, summary := range summaries {
		if summary.Status == status {
			return true
		}
	}
	return false
}

func TestRunnerContinuesValidationIncompleteExtractWithoutCheckpoint(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	lane := &prepared.Steps[0].lanes[0]
	lane.resolved.ExtractValidationPolicy.ValidatorFailure = ValidatorFailureWarnContinue
	installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		return erasedTypedResult{Value: typedValueForLane(0, request.Chunk.Index)}, nil
	})
	lane.extractValidators.validators[0].typedValidate = func(context.Context, any, typedValidationTarget) (contracts.ValidationResult, error) {
		return contracts.ValidationResult{}, fmt.Errorf("validator unavailable")
	}
	recorder := &extractCaptureRecorder{CheckpointRecorder: NoopCheckpointRecorder()}
	debug := newCapturedDebugRecorder()

	output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Checkpoints: recorder, Debug: debug})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if len(output.NormalizeOutputs) != 1 || len(output.Rejected) != 0 {
		t.Fatalf("run output = %#v, want accepted incomplete extract", output)
	}
	if output.Manifest.ValidationStatus != "incomplete" || !containsValidationSummary(output.Manifest.ValidationSummaries, "incomplete") {
		t.Fatalf("manifest validation = %#v, want incomplete extract provenance", output.Manifest)
	}
	if len(output.Diagnostics.Groups) != 1 || output.Diagnostics.Groups[0].ReasonCode != "validator_execution_incomplete" {
		t.Fatalf("diagnostics = %#v, want validation-incomplete group", output.Diagnostics)
	}
	if len(recorder.checkpoint.Outputs) != 0 || len(recorder.checkpoint.Rejected) != 0 {
		t.Fatalf("extract checkpoint = %#v, want no persisted incomplete output", recorder.checkpoint)
	}
	encoded := string(debug.json["extract/notes/output.json"])
	if !strings.Contains(encoded, `"validation_incomplete_chunks":[0]`) {
		t.Fatalf("extract output debug = %s, want incomplete chunk marker", encoded)
	}
}
