package pipeline

import (
	"bytes"
	"context"
	"testing"

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

type correctionObservingNotesExtractor struct {
	correction *contracts.SemanticCorrection
	candidate  *contracts.ModelCandidate
}

func (*correctionObservingNotesExtractor) Key() string { return "test/correction-observing-extract" }
func (*correctionObservingNotesExtractor) ReferenceSlots() []contracts.ReferenceSlot {
	return nil
}

func (e *correctionObservingNotesExtractor) Extract(_ context.Context, request contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[codecNotes], error) {
	e.correction = request.Correction
	candidate, err := contracts.NewModelCandidate([]byte(`{"items":["extract"]}`), contracts.CorrectionProtocolSingleResponseV1)
	if err != nil {
		return contracts.TypedExtractionResult[codecNotes]{}, err
	}
	e.candidate = candidate
	return contracts.TypedExtractionResult[codecNotes]{Value: codecNotes{Items: []string{"extract"}}, ModelCandidate: candidate}, nil
}

type correctionObservingNotesMerger struct {
	correction *contracts.SemanticCorrection
	candidate  *contracts.ModelCandidate
}

func (*correctionObservingNotesMerger) Key() string { return "test/correction-observing-merge" }

func (m *correctionObservingNotesMerger) Merge(_ context.Context, request contracts.TypedMergeRequest[codecNotes]) (contracts.TypedMergeResult[codecNotes], error) {
	m.correction = request.Correction
	candidate, err := contracts.NewModelCandidate([]byte(`{"items":["merge"]}`), contracts.CorrectionProtocolSingleResponseV1)
	if err != nil {
		return contracts.TypedMergeResult[codecNotes]{}, err
	}
	m.candidate = candidate
	return contracts.TypedMergeResult[codecNotes]{Value: codecNotes{Items: []string{"merge"}}, ModelCandidate: candidate}, nil
}

type correctionObservingNotesNormalizer struct {
	correction *contracts.SemanticCorrection
	candidate  *contracts.ModelCandidate
}

func (*correctionObservingNotesNormalizer) Key() string { return "test/correction-observing-normalize" }
func (*correctionObservingNotesNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
	return nil
}

func (n *correctionObservingNotesNormalizer) Normalize(_ context.Context, request contracts.TypedNormalizeRequest[codecNotes]) (contracts.TypedNormalizeResult[codecNotes], error) {
	n.correction = request.Correction
	candidate, err := contracts.NewModelCandidate([]byte(`{"items":["normalize"]}`), contracts.CorrectionProtocolSingleResponseV1)
	if err != nil {
		return contracts.TypedNormalizeResult[codecNotes]{}, err
	}
	n.candidate = candidate
	return contracts.TypedNormalizeResult[codecNotes]{Value: codecNotes{Items: []string{"normalize"}}, ModelCandidate: candidate}, nil
}

type repairObservingNotesMerger struct {
	attempts *int
}

func (*repairObservingNotesMerger) Key() string { return "test/repair-observing-merge" }

func (m *repairObservingNotesMerger) Merge(_ context.Context, request contracts.TypedMergeRequest[codecNotes]) (contracts.TypedMergeResult[codecNotes], error) {
	m.attempts = request.StructuredOutputRepairAttempts
	return contracts.TypedMergeResult[codecNotes]{Value: codecNotes{Items: []string{"merged"}}}, nil
}

type repairObservingNotesNormalizer struct {
	attempts *int
}

func (*repairObservingNotesNormalizer) Key() string { return "test/repair-observing-normalize" }
func (*repairObservingNotesNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
	return nil
}

func (n *repairObservingNotesNormalizer) Normalize(_ context.Context, request contracts.TypedNormalizeRequest[codecNotes]) (contracts.TypedNormalizeResult[codecNotes], error) {
	n.attempts = request.StructuredOutputRepairAttempts
	return contracts.TypedNormalizeResult[codecNotes]{Value: request.MergeOutput.Value}, nil
}

func TestTypedRegistryErasurePreservesStructuredOutputRepairAttempts(t *testing.T) {
	const artifactKind contracts.ArtifactKind = "test/notes"

	t.Run("merge", func(t *testing.T) {
		implementation := &repairObservingNotesMerger{}
		registry := NewMergerRegistry()
		if err := RegisterMerger(registry, ModuleSpec{Key: implementation.Key(), Stage: StageMerge, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Merger[codecNotes], error) {
			return implementation, nil
		}); err != nil {
			t.Fatalf("RegisterMerger() error = %v", err)
		}
		entry, ok := registry.typedEntry(implementation.Key(), artifactKind)
		if !ok {
			t.Fatal("typed merger entry missing")
		}
		built, err := entry.builder(BuildRequest{})
		if err != nil {
			t.Fatalf("builder() error = %v", err)
		}
		attempts := 2
		_, err = entry.merge(context.Background(), built, contracts.TypedMergeRequest[any]{
			ExtractOutputs:                 []contracts.ExtractArtifact[any]{{Value: codecNotes{Items: []string{"extracted"}}}},
			StructuredOutputRepairAttempts: &attempts,
		})
		if err != nil {
			t.Fatalf("merge() error = %v", err)
		}
		if implementation.attempts == nil || *implementation.attempts != attempts {
			t.Fatalf("repair attempts = %v, want %d", implementation.attempts, attempts)
		}
	})

	t.Run("normalize", func(t *testing.T) {
		implementation := &repairObservingNotesNormalizer{}
		registry := NewNormalizerRegistry()
		if err := RegisterNormalizer(registry, ModuleSpec{Key: implementation.Key(), Stage: StageNormalize, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Normalizer[codecNotes], error) {
			return implementation, nil
		}); err != nil {
			t.Fatalf("RegisterNormalizer() error = %v", err)
		}
		entry, ok := registry.typedEntry(implementation.Key(), artifactKind)
		if !ok {
			t.Fatal("typed normalizer entry missing")
		}
		built, err := entry.builder(BuildRequest{})
		if err != nil {
			t.Fatalf("builder() error = %v", err)
		}
		attempts := 3
		_, err = entry.normalize(context.Background(), built, contracts.TypedNormalizeRequest[any]{
			MergeOutput:                    contracts.MergeArtifact[any]{Value: codecNotes{Items: []string{"merged"}}},
			StructuredOutputRepairAttempts: &attempts,
		})
		if err != nil {
			t.Fatalf("normalize() error = %v", err)
		}
		if implementation.attempts == nil || *implementation.attempts != attempts {
			t.Fatalf("repair attempts = %v, want %d", implementation.attempts, attempts)
		}
	})
}

func TestTypedRegistryErasurePreservesCorrectionAndCandidateOwnership(t *testing.T) {
	const artifactKind contracts.ArtifactKind = "test/notes"

	t.Run("extract", func(t *testing.T) {
		correction := newTestCorrection(t)
		implementation := &correctionObservingNotesExtractor{}
		registry := NewExtractorRegistry()
		if err := RegisterExtractor(registry, ModuleSpec{Key: implementation.Key(), Stage: StageExtract, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Extractor[codecNotes], error) {
			return implementation, nil
		}); err != nil {
			t.Fatalf("RegisterExtractor() error = %v", err)
		}
		entry, ok := registry.typedEntry(implementation.Key())
		if !ok {
			t.Fatal("typed extractor entry missing")
		}
		built, err := entry.builder(BuildRequest{})
		if err != nil {
			t.Fatalf("builder() error = %v", err)
		}
		result, err := entry.extract(context.Background(), built, contracts.TypedExtractionRequest{Correction: correction})
		if err != nil {
			t.Fatalf("extract() error = %v", err)
		}
		assertCorrectionAndCandidateOwnership(t, correction, implementation.correction, implementation.candidate, result.ModelCandidate)
	})

	t.Run("merge", func(t *testing.T) {
		correction := newTestCorrection(t)
		implementation := &correctionObservingNotesMerger{}
		registry := NewMergerRegistry()
		if err := RegisterMerger(registry, ModuleSpec{Key: implementation.Key(), Stage: StageMerge, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Merger[codecNotes], error) {
			return implementation, nil
		}); err != nil {
			t.Fatalf("RegisterMerger() error = %v", err)
		}
		entry, ok := registry.typedEntry(implementation.Key(), artifactKind)
		if !ok {
			t.Fatal("typed merger entry missing")
		}
		built, err := entry.builder(BuildRequest{})
		if err != nil {
			t.Fatalf("builder() error = %v", err)
		}
		result, err := entry.merge(context.Background(), built, contracts.TypedMergeRequest[any]{Correction: correction, ExtractOutputs: []contracts.ExtractArtifact[any]{{Value: codecNotes{Items: []string{"extract"}}}}})
		if err != nil {
			t.Fatalf("merge() error = %v", err)
		}
		assertCorrectionAndCandidateOwnership(t, correction, implementation.correction, implementation.candidate, result.ModelCandidate)
	})

	t.Run("normalize", func(t *testing.T) {
		correction := newTestCorrection(t)
		implementation := &correctionObservingNotesNormalizer{}
		registry := NewNormalizerRegistry()
		if err := RegisterNormalizer(registry, ModuleSpec{Key: implementation.Key(), Stage: StageNormalize, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Normalizer[codecNotes], error) {
			return implementation, nil
		}); err != nil {
			t.Fatalf("RegisterNormalizer() error = %v", err)
		}
		entry, ok := registry.typedEntry(implementation.Key(), artifactKind)
		if !ok {
			t.Fatal("typed normalizer entry missing")
		}
		built, err := entry.builder(BuildRequest{})
		if err != nil {
			t.Fatalf("builder() error = %v", err)
		}
		result, err := entry.normalize(context.Background(), built, contracts.TypedNormalizeRequest[any]{Correction: correction, MergeOutput: contracts.MergeArtifact[any]{Value: codecNotes{Items: []string{"merge"}}}})
		if err != nil {
			t.Fatalf("normalize() error = %v", err)
		}
		assertCorrectionAndCandidateOwnership(t, correction, implementation.correction, implementation.candidate, result.ModelCandidate)
	})
}

func newTestCorrection(t *testing.T) *contracts.SemanticCorrection {
	t.Helper()
	correction, err := contracts.NewSemanticCorrection([]byte(`{"items":["original"]}`), "Correct the response.")
	if err != nil {
		t.Fatalf("NewSemanticCorrection() error = %v", err)
	}
	return correction
}

func assertCorrectionAndCandidateOwnership(t *testing.T, callerCorrection, observedCorrection *contracts.SemanticCorrection, producerCandidate, returnedCandidate *contracts.ModelCandidate) {
	t.Helper()
	if observedCorrection == nil || producerCandidate == nil || returnedCandidate == nil {
		t.Fatal("correction and candidates must be present")
	}
	callerCorrection.AssistantResponse[0] = '['
	if got := string(observedCorrection.AssistantResponse); got != `{"items":["original"]}` {
		t.Fatalf("observed correction = %q, want owned original content", got)
	}
	producerCandidate.Response[0] = '['
	if bytes.Equal(producerCandidate.Response, returnedCandidate.Response) {
		t.Fatalf("returned candidate aliases producer candidate: %q", returnedCandidate.Response)
	}
}
