package pipeline

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"testing"

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

type terminalChunker struct {
	key         string
	plan        source.ChunkPlan
	diagnostics []contracts.ProducerDiagnostic
	err         error
	calls       *int
}

func (c terminalChunker) Key() string { return c.key }

func (terminalChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil }

func (c terminalChunker) Plan(context.Context, contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
	if c.calls != nil {
		(*c.calls)++
	}
	return contracts.ChunkPlanResult{Plan: source.CloneChunkPlan(c.plan), Diagnostics: contracts.CloneProducerDiagnostics(c.diagnostics)}, c.err
}

type terminalChunkValidator struct {
	result contracts.ValidationResult
	err    error
}

type retryingLLMChunkValidator struct {
	requests []contracts.ChunkValidationRequest
	calls    int
}

type observingChunkValidator struct {
	request contracts.ChunkValidationRequest
}

type observingSerializedValidator struct {
	request contracts.SerializedValidationRequest
}

func (*observingChunkValidator) Name() string { return "observing/chunk-validator" }
func (*observingChunkValidator) ExecutionClass() contracts.ExecutionClass {
	return contracts.ExecutionClassDeterministic
}
func (v *observingChunkValidator) Validate(_ context.Context, request contracts.ChunkValidationRequest) (contracts.ValidationResult, error) {
	v.request = request
	return contracts.ValidationResult{Approved: true}, nil
}

func (*observingSerializedValidator) Name() string { return "observing/serialized-validator" }
func (*observingSerializedValidator) ExecutionClass() contracts.ExecutionClass {
	return contracts.ExecutionClassDeterministic
}
func (v *observingSerializedValidator) Validate(_ context.Context, request contracts.SerializedValidationRequest) (contracts.ValidationResult, error) {
	v.request = request
	return contracts.ValidationResult{Approved: true}, nil
}

func (terminalChunkValidator) Name() string { return "terminal/chunk-validator" }

func (terminalChunkValidator) ExecutionClass() contracts.ExecutionClass {
	return contracts.ExecutionClassDeterministic
}

func (v terminalChunkValidator) Validate(context.Context, contracts.ChunkValidationRequest) (contracts.ValidationResult, error) {
	return v.result, v.err
}

func (*retryingLLMChunkValidator) Name() string { return "retrying/llm-chunk-validator" }

func (*retryingLLMChunkValidator) ExecutionClass() contracts.ExecutionClass {
	return contracts.ExecutionClassLLMBacked
}

func (validator *retryingLLMChunkValidator) Validate(_ context.Context, request contracts.ChunkValidationRequest) (contracts.ValidationResult, error) {
	validator.calls++
	validator.requests = append(validator.requests, request)
	if validator.calls == 1 && len(request.Chunks) > 0 {
		request.Chunks[0].ID = "mutated"
	}
	if validator.calls == 1 {
		return contracts.ValidationResult{}, errors.New("temporary validator failure")
	}
	return contracts.ValidationResult{Approved: true}, nil
}

func assertAttemptEnvelopeSequence(t *testing.T, debug *capturedDebugRecorder, prefix string, attempts ...int) {
	t.Helper()
	marker := strings.TrimSuffix(prefix, "/") + "/attempt-"
	var got []int
	for _, name := range debug.names() {
		if !strings.HasPrefix(name, marker) || !strings.HasSuffix(name, ".json") {
			continue
		}
		remainder := strings.TrimSuffix(strings.TrimPrefix(name, marker), ".json")
		if strings.Contains(remainder, "/") {
			continue
		}
		value, err := strconv.Atoi(remainder)
		if err != nil {
			t.Fatalf("parse attempt index from %q: %v", name, err)
		}
		got = append(got, value)
	}
	sort.Ints(got)
	if !reflect.DeepEqual(got, attempts) {
		t.Fatalf("attempt envelopes under %q = %#v, want %#v; names = %#v", prefix, got, attempts, debug.names())
	}
}

func preparedTerminalDebugPipeline(t *testing.T) (*PreparedPipeline, source.ChunkPlan) {
	t.Helper()
	prepared := preparedAttemptDebugPipeline(t)
	prepared.resolved.ChunkValidationPolicy = ValidationPolicy{
		ProducerStructuralFailure: ProducerStructuralFailureFailRun,
		SemanticRejection:         SemanticRejectionRejectOutput,
		ValidatorFailure:          ValidatorFailureFailRun,
	}
	prepared.Steps[0].lanes[0].resolved.ExtractValidationPolicy = ValidationPolicy{
		ProducerStructuralFailure: ProducerStructuralFailureFailRun,
		SemanticRejection:         SemanticRejectionRejectOutput,
		ValidatorFailure:          ValidatorFailureFailRun,
	}
	prepared.Steps[0].lanes[0].resolved.MergeValidationPolicy = ValidationPolicy{
		ProducerStructuralFailure: ProducerStructuralFailureFailRun,
		SemanticRejection:         SemanticRejectionRejectOutput,
		ValidatorFailure:          ValidatorFailureFailRun,
	}
	prepared.Steps[0].lanes[0].resolved.NormalizeValidationPolicy = ValidationPolicy{
		ProducerStructuralFailure: ProducerStructuralFailureFailRun,
		SemanticRejection:         SemanticRejectionRejectOutput,
		ValidatorFailure:          ValidatorFailureFailRun,
	}
	chunker, ok := prepared.chunker.(*typedTestChunker)
	if !ok {
		t.Fatalf("prepared chunker = %T, want *typedTestChunker", prepared.chunker)
	}
	return prepared, source.CloneChunkPlan(chunker.plan)
}

func TestRunnerRecordsChunkTerminalOutcomes(t *testing.T) {
	tests := []struct {
		name          string
		moduleError   error
		validator     terminalChunkValidator
		wantError     string
		wantRejection bool
		attemptError  bool
	}{
		{name: "accepted", validator: terminalChunkValidator{result: contracts.ValidationResult{Approved: true}}},
		{name: "module error", moduleError: errors.New("chunk module failed"), wantError: "chunk module failed", attemptError: true},
		{name: "validator rejection", validator: terminalChunkValidator{result: contracts.ValidationResult{Approved: false, ReasonCode: "chunk_rejected", Message: "chunk rejected", CorrectionGuidance: "return an acceptable chunk plan"}}, wantRejection: true},
		{name: "validator error", validator: terminalChunkValidator{err: errors.New("chunk validator failed")}, wantError: "chunk validator failed"},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			prepared, plan := preparedTerminalDebugPipeline(t)
			prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan, diagnostics: []contracts.ProducerDiagnostic{producerDiagnostic("chunk", "chunk diagnostic")}, err: tc.moduleError}
			prepared.chunkValidators.validators = []preparedValidator{{resolved: ResolvedValidator{Binding: Binding("terminal/chunk-validator"), Target: ValidatorTargetChunk}, chunk: tc.validator}}
			debug := newCapturedDebugRecorder()

			output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
			assertAttemptEnvelopeSequence(t, debug, "chunk", 1)
			envelope := debug.envelope(t, "chunk/attempt-01.json")
			if tc.wantError != "" {
				if err == nil || !strings.Contains(err.Error(), tc.wantError) {
					t.Fatalf("Run() error = %v, attempt error = %q; want %q", err, envelope.Error, tc.wantError)
				}
				if tc.attemptError && !strings.Contains(envelope.Error, tc.wantError) {
					t.Fatalf("attempt error = %q, want %q", envelope.Error, tc.wantError)
				}
				if !tc.attemptError && envelope.Error != "" {
					t.Fatalf("settled validator failure attempt error = %q, want empty", envelope.Error)
				}
			} else if err != nil {
				t.Fatalf("Run() error = %v, want nil", err)
			}
			if tc.wantRejection {
				if envelope.Error != "" || len(output.Rejected) != 1 || output.Rejected[0].ReasonCode != "chunk_rejected" {
					t.Fatalf("chunk rejection = envelope %#v, outputs %#v", envelope, output.Rejected)
				}
			}
			if tc.name == "accepted" {
				encoded := string(debug.json["chunk/attempt-01.json"])
				if !strings.Contains(encoded, `"plan":`) || !strings.Contains(encoded, `"materialized_chunks":`) || strings.Contains(encoded, `"chunks":`) {
					t.Fatalf("chunk attempt debug = %s, want separate plan and materialized_chunks", encoded)
				}
			}
		})
	}
}

func TestRunnerRetriesLLMValidatorsWithoutRegeneratingChunkCandidate(t *testing.T) {
	prepared, plan := preparedTerminalDebugPipeline(t)
	chunkCalls := 0
	prepared.resolved.Chunk.Retries = 0
	prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan, calls: &chunkCalls}
	validator := &retryingLLMChunkValidator{}
	prepared.chunkValidators.validators = []preparedValidator{{
		resolved: ResolvedValidator{Binding: ModuleBinding{Module: validator.Name(), Retries: 1}, ExecutionClass: contracts.ExecutionClassLLMBacked, Target: ValidatorTargetChunk},
		chunk:    validator,
	}}

	if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil {
		t.Fatalf("Run() error = %v", err)
	}
	if chunkCalls != 1 || validator.calls != 2 {
		t.Fatalf("producer calls = %d validator calls = %d, want 1 and 2", chunkCalls, validator.calls)
	}
	if len(validator.requests) != 2 || len(validator.requests[0].Chunks) == 0 || len(validator.requests[1].Chunks) == 0 || validator.requests[1].Chunks[0].ID == "mutated" {
		t.Fatalf("validator requests = %#v, want independently owned immutable candidates", validator.requests)
	}
}

func TestRunnerMaterializesAnnotatedPlanBeforeChunkValidation(t *testing.T) {
	prepared, plan := preparedTerminalDebugPipeline(t)
	plan.Annotations = source.ChunkAnnotations{"same": []byte(`{"plan":1}`)}
	plan.Ranges[0].Annotations = source.ChunkAnnotations{"same": []byte(`{"range":2}`)}
	prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan}
	validator := &observingChunkValidator{}
	prepared.chunkValidators.validators = []preparedValidator{{
		resolved: ResolvedValidator{Binding: Binding(validator.Name()), Target: ValidatorTargetChunk},
		chunk:    validator,
	}}

	if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if len(validator.request.Chunks) != 1 {
		t.Fatalf("validator chunks = %#v, want one", validator.request.Chunks)
	}
	chunk := validator.request.Chunks[0]
	if chunk.ID != "chunk-000001" || chunk.Index != 0 || chunk.MediaType != "application/json" {
		t.Fatalf("materialized chunk identity = %#v", chunk)
	}
	if string(chunk.Annotations["same"]) != `{"range":2}` || string(chunk.PlanAnnotations["same"]) != `{"plan":1}` {
		t.Fatalf("materialized annotations = %s / %s", chunk.Annotations["same"], chunk.PlanAnnotations["same"])
	}
	if chunk.Metadata["start_unit_id"] != 1 || chunk.Metadata["end_unit_id"] != 1 || chunk.Metadata["unit_count"] != 1 {
		t.Fatalf("materialized metadata = %#v", chunk.Metadata)
	}
}

func TestRunnerForwardsDetachedRepairAttemptsToChunkValidators(t *testing.T) {
	prepared, plan := preparedTerminalDebugPipeline(t)
	prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan}
	chunkValidator := &observingChunkValidator{}
	serializedValidator := &observingSerializedValidator{}
	chunkAttempts := 2
	serializedAttempts := 3
	chunkBinding := Binding(chunkValidator.Name())
	chunkBinding.StructuredOutputRepairAttempts = &chunkAttempts
	serializedBinding := Binding(serializedValidator.Name())
	serializedBinding.StructuredOutputRepairAttempts = &serializedAttempts
	prepared.chunkValidators.validators = []preparedValidator{
		{resolved: ResolvedValidator{Binding: chunkBinding, Target: ValidatorTargetChunk}, chunk: chunkValidator},
		{resolved: ResolvedValidator{Binding: serializedBinding, Target: ValidatorTargetSerialized}, serialized: serializedValidator},
	}

	if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if got := chunkValidator.request.StructuredOutputRepairAttempts; got == nil || *got != chunkAttempts || got == chunkBinding.StructuredOutputRepairAttempts {
		t.Fatalf("chunk validator repair attempts = %v, want detached value %d", got, chunkAttempts)
	}
	if got := serializedValidator.request.StructuredOutputRepairAttempts; got == nil || *got != serializedAttempts || got == serializedBinding.StructuredOutputRepairAttempts {
		t.Fatalf("serialized validator repair attempts = %v, want detached value %d", got, serializedAttempts)
	}
}

func TestRunnerRetriesMalformedPlanWithDebugEnabled(t *testing.T) {
	prepared, plan := preparedTerminalDebugPipeline(t)
	plan.Ranges[0].Annotations = source.ChunkAnnotations{"broken": []byte(`{"value":`)}
	prepared.resolved.Chunk.Retries = 1
	calls := 0
	prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan, calls: &calls}
	debug := newCapturedDebugRecorder()

	_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
	if err == nil || !strings.Contains(err.Error(), "valid JSON") {
		t.Fatalf("Run() error = %v, want malformed annotation error", err)
	}
	if calls != 2 {
		t.Fatalf("plan calls = %d, want two attempts", calls)
	}
	assertAttemptEnvelopeSequence(t, debug, "chunk", 1, 2)
}

func TestRunnerRecordsExtractTerminalOutcomes(t *testing.T) {
	tests := []struct {
		name          string
		moduleError   error
		validatorErr  error
		reject        bool
		candidateFail bool
		finalFail     bool
		wantError     string
		attemptError  bool
	}{
		{name: "terminal rejection", reject: true},
		{name: "module error", moduleError: errors.New("extract module failed"), wantError: "extract module failed", attemptError: true},
		{name: "validator error", validatorErr: errors.New("extract validator failed"), wantError: "extract validator failed"},
		{name: "candidate codec error", candidateFail: true, wantError: "serialize extract candidate", attemptError: true},
		{name: "final codec error", finalFail: true, wantError: "serialize accepted extract output", attemptError: true},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			prepared := preparedAttemptDebugPipeline(t)
			prepared.Steps[0].lanes[0].resolved.ExtractValidationPolicy = ValidationPolicy{
				ProducerStructuralFailure: ProducerStructuralFailureFailRun,
				SemanticRejection:         SemanticRejectionRejectOutput,
				ValidatorFailure:          ValidatorFailureFailRun,
			}
			value := "extract-terminal"
			codec := &observedNotesCodec{}
			if tc.candidateFail {
				codec.candidateError = value
			}
			if tc.finalFail {
				codec.finalError = value
			}
			installObservedNotesCodec(t, prepared, codec)
			installExtractOperation(prepared, 0, func(context.Context, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
				if tc.moduleError != nil {
					return erasedTypedResult{}, tc.moduleError
				}
				return erasedTypedResult{Value: codecNotes{Items: []string{value}}, Diagnostics: []contracts.ProducerDiagnostic{producerDiagnostic("extract", "extract diagnostic")}}, nil
			})
			prepared.Steps[0].lanes[0].extractValidators.validators = []preparedValidator{{
				resolved: ResolvedValidator{Binding: Binding("terminal/extract-validator"), Target: ValidatorTargetTyped, ArtifactKind: "test/notes"},
				typedValidate: func(context.Context, any, typedValidationTarget) (contracts.ValidationResult, error) {
					return contracts.ValidationResult{Approved: !tc.reject, ReasonCode: "extract_rejected", Message: "extract rejected", CorrectionGuidance: "return an acceptable extract"}, tc.validatorErr
				},
			}}
			debug := newCapturedDebugRecorder()

			output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
			assertAttemptEnvelopeSequence(t, debug, "extract/notes/chunk-000001", 1)
			attemptPath := "extract/notes/chunk-000001/attempt-01.json"
			envelope := debug.envelope(t, attemptPath)
			if tc.wantError != "" {
				if err == nil || !strings.Contains(err.Error(), tc.wantError) {
					t.Fatalf("Run() error = %v, attempt error = %q; want %q", err, envelope.Error, tc.wantError)
				}
				if tc.attemptError && !strings.Contains(envelope.Error, tc.wantError) {
					t.Fatalf("attempt error = %q, want %q", envelope.Error, tc.wantError)
				}
				if !tc.attemptError && envelope.Error != "" {
					t.Fatalf("settled validator failure attempt error = %q, want empty", envelope.Error)
				}
			} else {
				if err != nil {
					t.Fatalf("Run() error = %v, want nil rejection", err)
				}
				if envelope.Error != "" || len(output.Rejected) != 1 || output.Rejected[0].ReasonCode != "extract_rejected" {
					t.Fatalf("extract rejection = envelope %#v, outputs %#v", envelope, output.Rejected)
				}
			}
		})
	}
}

func TestRunnerJoinsPrimaryAndAttemptWriteErrors(t *testing.T) {
	t.Run("chunk", func(t *testing.T) {
		prepared, plan := preparedTerminalDebugPipeline(t)
		primary := errors.New("chunk operation failed")
		prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan, err: primary}
		debug := newCapturedDebugRecorder()
		debug.failPath = "chunk/attempt-01.json"

		_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
		if err == nil || !errors.Is(err, primary) || !strings.Contains(err.Error(), "write chunk attempt debug artifact") || !strings.Contains(err.Error(), "debug recorder failure") {
			t.Fatalf("Run() error = %v, want joined operation and debug errors", err)
		}
	})

	t.Run("extract", func(t *testing.T) {
		prepared := preparedAttemptDebugPipeline(t)
		primary := errors.New("extract operation failed")
		installExtractOperation(prepared, 0, func(context.Context, contracts.TypedExtractionRequest) (erasedTypedResult, error) {
			return erasedTypedResult{}, primary
		})
		debug := newCapturedDebugRecorder()
		debug.failPath = "extract/notes/chunk-000001/attempt-01.json"

		_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
		if err == nil || !errors.Is(err, primary) || !strings.Contains(err.Error(), "write extract attempt debug artifact") || !strings.Contains(err.Error(), "debug recorder failure") {
			t.Fatalf("Run() error = %v, want joined operation and debug errors", err)
		}
	})
}

func TestRunnerDoesNotRetryAfterTerminalAttemptWriteFailure(t *testing.T) {
	prepared, plan := preparedTerminalDebugPipeline(t)
	prepared.resolved.Chunk.Retries = 1
	calls := 0
	prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan, calls: &calls}
	prepared.chunkValidators.validators = []preparedValidator{{
		resolved: ResolvedValidator{Binding: Binding("terminal/chunk-validator"), Target: ValidatorTargetChunk},
		chunk:    terminalChunkValidator{result: contracts.ValidationResult{Approved: true}},
	}}
	debug := newCapturedDebugRecorder()
	debug.failPath = "chunk/attempt-01.json"

	_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
	if err == nil || !strings.Contains(err.Error(), "write chunk attempt debug artifact") || !strings.Contains(err.Error(), "debug recorder failure") {
		t.Fatalf("Run() error = %v, want terminal attempt debug failure", err)
	}
	if calls != 1 {
		t.Fatalf("chunk calls = %d, want one attempt without retry", calls)
	}
	if debug.has("chunk/attempt-02.json") {
		t.Fatal("second chunk attempt envelope exists after non-retryable debug persistence failure")
	}
}

func TestRunnerKeepsExtractModuleAndValidatorLLMCallsIsolated(t *testing.T) {
	prepared := preparedAttemptDebugPipelineWithChunks(t, 2)
	debug := newCapturedDebugRecorder()
	client := WithDebugLLMRecording(attemptDebugLLM{}, debug)
	installExtractOperation(prepared, 0, func(ctx context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		if err := callAttemptDebugLLM(ctx, client, "extract-module"); err != nil {
			return erasedTypedResult{}, err
		}
		return erasedTypedResult{Value: typedValueForLane(0, request.Chunk.Index)}, nil
	})
	prepared.Steps[0].lanes[0].extractValidators.validators = []preparedValidator{{
		resolved: ResolvedValidator{Binding: Binding("llm-check"), Target: ValidatorTargetTyped, ArtifactKind: "test/notes"},
		typedValidate: func(ctx context.Context, _ any, _ typedValidationTarget) (contracts.ValidationResult, error) {
			if err := callAttemptDebugLLM(ctx, client, "extract-validator"); err != nil {
				return contracts.ValidationResult{}, err
			}
			return contracts.ValidationResult{Approved: true}, nil
		},
	}}

	if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug}); err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	validatorResponses := make(map[string]struct{})
	for chunkNumber := 1; chunkNumber <= 2; chunkNumber++ {
		chunkPath := fmt.Sprintf("chunk-%06d", chunkNumber)
		module := debug.envelope(t, "extract/notes/"+chunkPath+"/attempt-01.json")
		validator := debug.envelope(t, "validate/extract/notes/typed%2Fextract-notes/"+chunkPath+"/01-llm-check-attempt-01.json")
		if len(module.LLMCalls) != 1 || !strings.Contains(module.LLMCalls[0].ResponsePath, "extract/notes/"+chunkPath+"/attempt-01/") {
			t.Fatalf("module LLM calls for %s = %#v, want extract module call only", chunkPath, module.LLMCalls)
		}
		if len(validator.LLMCalls) != 1 || !strings.Contains(validator.LLMCalls[0].ResponsePath, "validate/extract/notes/typed%2Fextract-notes/"+chunkPath+"/") {
			t.Fatalf("validator LLM calls for %s = %#v, want chunk-scoped validator call only", chunkPath, validator.LLMCalls)
		}
		if module.LLMCalls[0].CallID == validator.LLMCalls[0].CallID {
			t.Fatalf("module and validator attempts for %s share LLM call %#v", chunkPath, module.LLMCalls)
		}
		validatorResponses[validator.LLMCalls[0].ResponsePath] = struct{}{}
	}
	if len(validatorResponses) != 2 {
		t.Fatalf("validator response paths = %#v, want one distinct path per chunk", validatorResponses)
	}
}
