package pipeline

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

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

type capturedDebugRecorder struct {
	mu       sync.Mutex
	json     map[string][]byte
	bytes    map[string][]byte
	failPath string
}

func newCapturedDebugRecorder() *capturedDebugRecorder {
	return &capturedDebugRecorder{json: make(map[string][]byte), bytes: make(map[string][]byte)}
}

func (*capturedDebugRecorder) Enabled() bool { return true }

func (r *capturedDebugRecorder) WriteJSON(name string, payload any) error {
	if name == r.failPath {
		return errors.New("debug recorder failure")
	}
	data, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	r.mu.Lock()
	r.json[name] = data
	r.mu.Unlock()
	return nil
}

func (r *capturedDebugRecorder) WriteBytes(name string, data []byte) error {
	if name == r.failPath {
		return errors.New("debug recorder failure")
	}
	r.mu.Lock()
	r.bytes[name] = append([]byte(nil), data...)
	r.mu.Unlock()
	return nil
}

func (r *capturedDebugRecorder) has(name string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	_, jsonOK := r.json[name]
	_, bytesOK := r.bytes[name]
	return jsonOK || bytesOK
}

func (r *capturedDebugRecorder) names() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	names := make([]string, 0, len(r.json)+len(r.bytes))
	for name := range r.json {
		names = append(names, name)
	}
	for name := range r.bytes {
		names = append(names, name)
	}
	sort.Strings(names)
	return names
}

func (r *capturedDebugRecorder) envelope(t *testing.T, name string) debugTimedEnvelope {
	t.Helper()
	r.mu.Lock()
	data := append([]byte(nil), r.json[name]...)
	r.mu.Unlock()
	if len(data) == 0 {
		t.Fatalf("debug envelope %q was not written; names = %#v", name, r.names())
	}
	var envelope debugTimedEnvelope
	if err := json.Unmarshal(data, &envelope); err != nil {
		t.Fatalf("decode debug envelope %q: %v", name, err)
	}
	return envelope
}

type attemptDebugLLM struct{}

func (attemptDebugLLM) CompleteStructured(_ context.Context, request contracts.StructuredCompletionRequest, _ any) (contracts.StructuredCompletionResponse, error) {
	return contracts.StructuredCompletionResponse{
		Content:   json.RawMessage(`{"accepted":true}`),
		Provider:  "test",
		Model:     "test-model",
		ProfileID: request.ProfileID,
		Debug: &contracts.LLMDebugMaterial{
			Prompt: &contracts.LLMDebugPrompt{PromptID: request.PromptID, Messages: []contracts.LLMDebugMessage{{Role: "user", Content: "test"}}},
			Response: &contracts.LLMDebugResponse{
				Content:   `{"accepted":true}`,
				PromptID:  request.PromptID,
				ModelName: "test-model",
			},
		},
	}, nil
}

func preparedAttemptDebugPipeline(t *testing.T) *PreparedPipeline {
	return preparedAttemptDebugPipelineWithChunks(t, 1)
}

func preparedAttemptDebugPipelineWithChunks(t *testing.T, chunkCount int) *PreparedPipeline {
	t.Helper()
	prepared := preparedConcurrentPipeline(t, chunkCount)
	prepared.Steps[0].lanes = prepared.Steps[0].lanes[:1]
	prepared.resolved.Steps[0].ArtifactLanes = prepared.resolved.Steps[0].ArtifactLanes[:1]
	prepared.Steps[0].ArtifactLanes = prepared.Steps[0].ArtifactLanes[:1]
	prepared.Steps[0].lanes[0].mergeValidators = preparedValidatorChain{}
	prepared.Steps[0].lanes[0].normalizeValidators = preparedValidatorChain{}
	return prepared
}

func TestRunnerForwardsDetachedStructuredOutputRepairAttempts(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	lane := &prepared.Steps[0].lanes[0]
	producerAttempts := 2
	validatorAttempts := 0
	lane.resolved.Extract.StructuredOutputRepairAttempts = &producerAttempts
	if len(lane.extractValidators.validators) == 0 {
		t.Fatal("extract validators are empty")
	}
	lane.extractValidators.validators[0].resolved.Binding.StructuredOutputRepairAttempts = &validatorAttempts

	var observedProducer, observedValidator *int
	installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
		observedProducer = request.StructuredOutputRepairAttempts
		return erasedTypedResult{Value: codecNotes{Items: []string{"extract"}}}, nil
	})
	lane.extractValidators.validators[0].typedValidate = func(_ context.Context, _ any, target typedValidationTarget) (contracts.ValidationResult, error) {
		observedValidator = target.structuredOutputRepairAttempts
		return contracts.ValidationResult{Approved: true}, nil
	}

	if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if observedProducer == nil || *observedProducer != producerAttempts || observedProducer == lane.resolved.Extract.StructuredOutputRepairAttempts {
		t.Fatalf("producer repair attempts = %v, want detached value %d", observedProducer, producerAttempts)
	}
	if observedValidator == nil || *observedValidator != validatorAttempts || observedValidator == lane.extractValidators.validators[0].resolved.Binding.StructuredOutputRepairAttempts {
		t.Fatalf("validator repair attempts = %v, want detached value %d", observedValidator, validatorAttempts)
	}
}

func callAttemptDebugLLM(ctx context.Context, client contracts.StructuredLLMClient, name string) error {
	_, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: "unscoped-" + name, PromptID: name, ProfileID: "test"}, nil)
	return err
}

func TestRunnerWritesAttemptScopedMergeAndNormalizeDebug(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	debug := newCapturedDebugRecorder()
	client := WithDebugLLMRecording(attemptDebugLLM{}, debug)
	prepared.Steps[0].lanes[0].typed.merge = func(ctx context.Context, _ any, _ contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
		if err := callAttemptDebugLLM(ctx, client, "merge"); err != nil {
			return erasedTypedResult{}, err
		}
		return erasedTypedResult{Value: codecNotes{Items: []string{"merged"}}, Diagnostics: []contracts.ProducerDiagnostic{producerDiagnostic("merge", "merge diagnostic")}}, nil
	}
	prepared.Steps[0].lanes[0].typed.normalize = func(ctx context.Context, _ any, _ contracts.TypedNormalizeRequest[any]) (erasedTypedResult, error) {
		if err := callAttemptDebugLLM(ctx, client, "normalize"); err != nil {
			return erasedTypedResult{}, err
		}
		return erasedTypedResult{Value: codecNotes{Items: []string{"normalized"}}, Diagnostics: []contracts.ProducerDiagnostic{producerDiagnostic("normalize", "normalize diagnostic")}}, nil
	}

	output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if len(output.NormalizeOutputs) != 1 {
		t.Fatalf("normalize outputs = %d, want one", len(output.NormalizeOutputs))
	}

	wantPaths := []string{
		"merge/notes/input.json",
		"merge/notes/attempt-01.json",
		"merge/notes/attempt-01/prompt-0001.json",
		"merge/notes/attempt-01/response-0001.json",
		"merge/notes/attempt-01/response-content-0001.json",
		"merge/notes/output.json",
		"normalize/notes/input.json",
		"normalize/notes/attempt-01.json",
		"normalize/notes/attempt-01/prompt-0002.json",
		"normalize/notes/attempt-01/response-0002.json",
		"normalize/notes/attempt-01/response-content-0002.json",
		"normalize/notes/output.json",
	}
	for _, name := range wantPaths {
		if !debug.has(name) {
			t.Errorf("debug artifact %q missing; names = %#v", name, debug.names())
		}
	}
	for _, name := range debug.names() {
		if strings.HasPrefix(name, "unscoped-merge/") || strings.HasPrefix(name, "unscoped-normalize/") {
			t.Errorf("LLM call used fallback debug path %q", name)
		}
	}
	mergeEnvelope := debug.envelope(t, "merge/notes/attempt-01.json")
	normalizeEnvelope := debug.envelope(t, "normalize/notes/attempt-01.json")
	if len(mergeEnvelope.LLMCalls) != 1 || mergeEnvelope.LLMCalls[0].ResponsePath != "merge/notes/attempt-01/response-0001.json" {
		t.Fatalf("merge LLM calls = %#v, want attempt-scoped call", mergeEnvelope.LLMCalls)
	}
	if len(normalizeEnvelope.LLMCalls) != 1 || normalizeEnvelope.LLMCalls[0].ResponsePath != "normalize/notes/attempt-01/response-0002.json" {
		t.Fatalf("normalize LLM calls = %#v, want attempt-scoped call", normalizeEnvelope.LLMCalls)
	}
}

func TestRunnerRecordsDistinctRetryAttemptsAndPromotesAcceptedDiagnosticsOnly(t *testing.T) {
	for _, stage := range []ModuleStage{StageMerge, StageNormalize} {
		t.Run(string(stage), func(t *testing.T) {
			prepared := preparedAttemptDebugPipeline(t)
			debug := newCapturedDebugRecorder()
			client := WithDebugLLMRecording(attemptDebugLLM{}, debug)
			lane := &prepared.Steps[0].lanes[0]
			attempts := 0
			operation := func(ctx context.Context) (erasedTypedResult, error) {
				attempts++
				if err := callAttemptDebugLLM(ctx, client, string(stage)); err != nil {
					return erasedTypedResult{}, err
				}
				scope := "accepted"
				if attempts == 1 {
					scope = "discarded"
				}
				return erasedTypedResult{Value: codecNotes{Items: []string{scope}}, Diagnostics: []contracts.ProducerDiagnostic{producerDiagnostic(scope, scope)}, ModelCandidate: attemptCandidate(t, fmt.Sprintf(`{"items":["%s"]}`, scope))}, nil
			}
			validatorCalls := 0
			validator := preparedValidator{
				resolved: ResolvedValidator{Binding: Binding("retry-check"), Target: ValidatorTargetTyped, ArtifactKind: "test/notes"},
				typedValidate: func(context.Context, any, typedValidationTarget) (contracts.ValidationResult, error) {
					validatorCalls++
					return contracts.ValidationResult{Approved: validatorCalls > 1, ReasonCode: "retry", Message: "retry candidate", CorrectionGuidance: "return an acceptable candidate"}, nil
				},
			}
			switch stage {
			case StageMerge:
				lane.resolved.Merge.Retries = 1
				lane.mergeValidators.validators = []preparedValidator{validator}
				lane.typed.merge = func(ctx context.Context, _ any, _ contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
					return operation(ctx)
				}
			case StageNormalize:
				lane.resolved.Normalize.Retries = 1
				lane.normalizeValidators.validators = []preparedValidator{validator}
				lane.typed.normalize = func(ctx context.Context, _ any, _ contracts.TypedNormalizeRequest[any]) (erasedTypedResult, error) {
					return operation(ctx)
				}
			}

			output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
			if err != nil {
				t.Fatalf("Run() error = %v, want nil", err)
			}
			laneID := lane.resolved.ID
			firstPath := fmt.Sprintf("%s/%s/attempt-01.json", stage, laneID)
			secondPath := fmt.Sprintf("%s/%s/attempt-02.json", stage, laneID)
			assertAttemptEnvelopeSequence(t, debug, fmt.Sprintf("%s/%s", stage, laneID), 1, 2)
			first := debug.envelope(t, firstPath)
			second := debug.envelope(t, secondPath)
			if len(first.LLMCalls) != 1 || len(second.LLMCalls) != 1 || first.LLMCalls[0].CallID == second.LLMCalls[0].CallID {
				t.Fatalf("retry LLM calls = first %#v, second %#v; want distinct calls", first.LLMCalls, second.LLMCalls)
			}
			if !strings.Contains(string(debug.json[firstPath]), "rejection") {
				t.Fatalf("first attempt envelope = %s, want rejection", debug.json[firstPath])
			}
			if len(output.Diagnostics.Groups) != 1 || output.Diagnostics.Groups[0].Samples[0].Scope != "accepted" {
				t.Fatalf("promoted diagnostics = %#v, want accepted attempt only", output.Diagnostics)
			}
		})
	}
}

func TestRunnerAttemptDebugRepresentsFailuresAndFinalRejection(t *testing.T) {
	tests := []struct {
		name         string
		configure    func(*PreparedPipeline)
		path         string
		wantError    string
		wantBody     string
		attemptError bool
	}{
		{
			name: "merge module error",
			configure: func(prepared *PreparedPipeline) {
				prepared.Steps[0].lanes[0].typed.merge = func(context.Context, any, contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
					return erasedTypedResult{}, errors.New("merge exploded")
				}
			},
			path:         "merge/notes/attempt-01.json",
			wantError:    "merge exploded",
			attemptError: true,
		},
		{
			name: "normalize validator error",
			configure: func(prepared *PreparedPipeline) {
				prepared.Steps[0].lanes[0].resolved.NormalizeValidationPolicy.ValidatorFailure = ValidatorFailureFailRun
				prepared.Steps[0].lanes[0].normalizeValidators.validators = []preparedValidator{{
					resolved: ResolvedValidator{Binding: Binding("error-check"), Target: ValidatorTargetTyped, ArtifactKind: "test/notes"},
					typedValidate: func(context.Context, any, typedValidationTarget) (contracts.ValidationResult, error) {
						return contracts.ValidationResult{}, errors.New("validator exploded")
					},
				}}
			},
			path:      "normalize/notes/attempt-01.json",
			wantError: "validator exploded",
			wantBody:  "output",
		},
		{
			name: "merge final rejection",
			configure: func(prepared *PreparedPipeline) {
				prepared.Steps[0].lanes[0].resolved.MergeValidationPolicy.SemanticRejection = SemanticRejectionRejectOutput
				prepared.Steps[0].lanes[0].mergeValidators.validators = []preparedValidator{{
					resolved: ResolvedValidator{Binding: Binding("reject-check"), Target: ValidatorTargetTyped, ArtifactKind: "test/notes"},
					typedValidate: func(context.Context, any, typedValidationTarget) (contracts.ValidationResult, error) {
						return contracts.ValidationResult{Approved: false, ReasonCode: "rejected", Message: "not accepted", CorrectionGuidance: "return an acceptable candidate"}, nil
					},
				}}
			},
			path:     "merge/notes/attempt-01.json",
			wantBody: "rejection",
		},
		{
			name: "normalize serialization error",
			configure: func(prepared *PreparedPipeline) {
				prepared.Steps[0].lanes[0].typed.normalize = func(context.Context, any, contracts.TypedNormalizeRequest[any]) (erasedTypedResult, error) {
					return erasedTypedResult{Value: "wrong artifact type"}, nil
				}
			},
			path:         "normalize/notes/attempt-01.json",
			wantError:    "serialize normalize candidate",
			attemptError: true,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			prepared := preparedAttemptDebugPipeline(t)
			debug := newCapturedDebugRecorder()
			tc.configure(prepared)
			output, runErr := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
			envelope := debug.envelope(t, tc.path)
			if tc.wantError != "" {
				if runErr == nil || !strings.Contains(runErr.Error(), tc.wantError) {
					t.Fatalf("run error = %v, envelope error = %q; want %q", runErr, envelope.Error, tc.wantError)
				}
				if tc.attemptError && !strings.Contains(envelope.Error, tc.wantError) {
					t.Fatalf("attempt envelope 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 runErr != nil {
				t.Fatalf("Run() error = %v, want nil rejection outcome", runErr)
			}
			if tc.wantBody != "" && !strings.Contains(string(debug.json[tc.path]), tc.wantBody) {
				t.Fatalf("attempt envelope = %s, want %q", debug.json[tc.path], tc.wantBody)
			}
			if tc.name == "merge final rejection" && len(output.Rejected) != 1 {
				t.Fatalf("rejected outputs = %#v, want one", output.Rejected)
			}
		})
	}
}

func TestRunnerKeepsValidatorLLMCallsOutOfModuleAttempt(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	debug := newCapturedDebugRecorder()
	client := WithDebugLLMRecording(attemptDebugLLM{}, debug)
	prepared.Steps[0].lanes[0].typed.merge = func(ctx context.Context, _ any, _ contracts.TypedMergeRequest[any]) (erasedTypedResult, error) {
		if err := callAttemptDebugLLM(ctx, client, "merge-module"); err != nil {
			return erasedTypedResult{}, err
		}
		return erasedTypedResult{Value: codecNotes{Items: []string{"merged"}}}, nil
	}
	prepared.Steps[0].lanes[0].mergeValidators.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, "merge-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)
	}
	moduleEnvelope := debug.envelope(t, "merge/notes/attempt-01.json")
	validatorPath := "validate/merge/notes/typed%2Fmerge/01-llm-check-attempt-01.json"
	validatorEnvelope := debug.envelope(t, validatorPath)
	if len(moduleEnvelope.LLMCalls) != 1 || !strings.Contains(moduleEnvelope.LLMCalls[0].ResponsePath, "merge/notes/attempt-01/") {
		t.Fatalf("module LLM calls = %#v, want module call only", moduleEnvelope.LLMCalls)
	}
	if len(validatorEnvelope.LLMCalls) != 1 || !strings.Contains(validatorEnvelope.LLMCalls[0].ResponsePath, "validate/merge/notes/") {
		t.Fatalf("validator LLM calls = %#v, want validator call only", validatorEnvelope.LLMCalls)
	}
	if moduleEnvelope.LLMCalls[0].CallID == validatorEnvelope.LLMCalls[0].CallID {
		t.Fatalf("module and validator envelopes reference the same call: %#v", moduleEnvelope.LLMCalls)
	}
}

type attemptReuseLoader struct {
	CheckpointLoader
	laneID    string
	merge     MergeCheckpoint
	normalize NormalizeCheckpoint
}

func (l attemptReuseLoader) Enabled() bool { return true }

func (l attemptReuseLoader) Merge(laneID, _ string, _ []CheckpointFingerprint) (MergeCheckpoint, CheckpointDecision) {
	if laneID == l.laneID {
		return l.merge, CheckpointDecision{Reused: true, Reason: "test reuse"}
	}
	return MergeCheckpoint{}, CheckpointDecision{Reason: "not found"}
}

func (l attemptReuseLoader) Normalize(laneID, _ string, _ []CheckpointFingerprint) (NormalizeCheckpoint, CheckpointDecision) {
	if laneID == l.laneID {
		return l.normalize, CheckpointDecision{Reused: true, Reason: "test reuse"}
	}
	return NormalizeCheckpoint{}, CheckpointDecision{Reason: "not found"}
}

func TestRunnerCheckpointReuseDoesNotSynthesizeModuleAttempts(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	lane := prepared.Steps[0].lanes[0]
	merge, err := checkpointArtifact(lane.typed.codec, lane.resolved.ID, lane.resolved.Merge.Module, "source", codecNotes{Items: []string{"merged"}})
	if err != nil {
		t.Fatalf("checkpointArtifact(merge): %v", err)
	}
	normalize, err := checkpointArtifact(lane.typed.codec, lane.resolved.ID, lane.resolved.Normalize.Module, "source", codecNotes{Items: []string{"normalized"}})
	if err != nil {
		t.Fatalf("checkpointArtifact(normalize): %v", err)
	}
	loader := attemptReuseLoader{
		CheckpointLoader: NoopCheckpointLoader(),
		laneID:           lane.resolved.ID,
		merge:            MergeCheckpoint{Output: merge},
		normalize:        NormalizeCheckpoint{Output: normalize},
	}
	debug := newCapturedDebugRecorder()
	if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug, Checkpoint: loader}); err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	for _, name := range debug.names() {
		if (strings.HasPrefix(name, "merge/notes/attempt-") || strings.HasPrefix(name, "normalize/notes/attempt-")) && name != "" {
			t.Fatalf("checkpoint reuse synthesized module attempt artifact %q", name)
		}
	}
	for _, name := range []string{"merge/notes/input.json", "merge/notes/output.json", "normalize/notes/input.json", "normalize/notes/output.json"} {
		if !debug.has(name) {
			t.Errorf("checkpoint reuse missing stage-level artifact %q", name)
		}
	}
}

func TestRunnerTreatsModuleAttemptDebugWriteFailureAsFrameworkError(t *testing.T) {
	prepared := preparedAttemptDebugPipeline(t)
	debug := newCapturedDebugRecorder()
	debug.failPath = "merge/notes/attempt-01.json"
	_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
	if err == nil || !strings.Contains(err.Error(), "write merge attempt debug artifact") || !strings.Contains(err.Error(), "debug recorder failure") {
		t.Fatalf("Run() error = %v, want merge attempt debug failure", err)
	}
}
