package pipeline

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"reflect"
	"strings"
	"testing"

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

func TestWalkingSkeletonFixture(t *testing.T) {
	inputBytes := readTestFixture(t, "testdata/walking_skeleton_input.json")
	expectedBytes := readTestFixture(t, "testdata/walking_skeleton_output.json")
	llmClient := &walkingSkeletonLLMClient{}

	resolved, err := ResolvePipeline(walkingSkeletonProfile(), ResolveOptions{}, walkingSkeletonCatalog(t))
	if err != nil {
		t.Fatalf("ResolvePipeline() error = %v, want nil", err)
	}

	output, err := newPreparedRunner(t, walkingSkeletonRegistries(t)).Run(context.Background(), RunInput{
		pipeline:  resolved,
		SourceID:  "fixture-source",
		Path:      "walking_skeleton_input.json",
		RawInput:  inputBytes,
		llmClient: llmClient,
	})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if len(output.OutputFiles) != 1 {
		t.Fatalf("len(OutputFiles) = %d, want 1", len(output.OutputFiles))
	}
	if output.OutputFiles[0].ContentType != "application/json" {
		t.Fatalf("ContentType = %q, want application/json", output.OutputFiles[0].ContentType)
	}
	assertStructuralJSONEqual(t, output.OutputFiles[0].Bytes, expectedBytes)
	if llmClient.calls != 3 {
		t.Fatalf("LLM calls = %d, want extractor calls plus normalizer call", llmClient.calls)
	}
}

func TestWalkingSkeletonResolutionRejectsMissingCapability(t *testing.T) {
	catalog := walkingSkeletonCatalog(t)
	catalog.Extractors = NewExtractorRegistry()
	if err := catalog.Extractors.RegisterLegacyRawWithSpec(ModuleSpec{
		Key:      "fake/extract",
		Stage:    StageExtract,
		Requires: []string{"missing"},
		Provides: []string{"fake_artifacts"},
	}, func() (contracts.LegacyRawExtractor, error) {
		return walkingSkeletonExtractor{}, nil
	}); err != nil {
		t.Fatalf("RegisterWithSpec() error = %v, want nil", err)
	}

	_, err := ResolvePipeline(walkingSkeletonProfile(), ResolveOptions{}, catalog)
	if err == nil {
		t.Fatal("ResolvePipeline() error = nil, want error")
	}
	if !strings.Contains(err.Error(), "missing") {
		t.Fatalf("ResolvePipeline() error = %q, want missing capability", err.Error())
	}
}

func TestWalkingSkeletonResolutionRejectsUnknownOnlyLane(t *testing.T) {
	_, err := ResolvePipeline(walkingSkeletonProfile(), ResolveOptions{Only: []string{"missing"}}, walkingSkeletonCatalog(t))
	if err == nil {
		t.Fatal("ResolvePipeline() error = nil, want error")
	}
	if !strings.Contains(err.Error(), "missing") || !strings.Contains(err.Error(), "not declared") {
		t.Fatalf("ResolvePipeline() error = %q, want unknown lane error", err.Error())
	}
}

func walkingSkeletonProfile() PipelineProfile {
	return PipelineProfile{
		ID:     "walking-skeleton",
		Input:  Binding("fake/input"),
		Chunk:  Binding("fake/chunk"),
		Output: Binding("json"),
		Artifacts: map[string]ArtifactLaneProfile{
			"events": {Extract: Binding("fake/extract")},
		},
	}
}

func walkingSkeletonCatalog(t *testing.T) ModuleCatalog {
	t.Helper()

	catalog := ModuleCatalog{
		Inputs:          NewInputAdapterRegistry(),
		Chunkers:        NewChunkerRegistry(),
		ArtifactCodecs:  NewArtifactCodecRegistry(),
		Extractors:      NewExtractorRegistry(),
		Mergers:         NewMergerRegistry(),
		Normalizers:     NewNormalizerRegistry(),
		ValidatorChains: NewValidatorChainRegistry(),
		Outputs:         NewOutputEncoderRegistry(),
	}
	if err := catalog.Inputs.RegisterWithSpec(ModuleSpec{
		Key:      "fake/input",
		Stage:    StageInput,
		Provides: []string{"plain_text"},
	}, func() (contracts.InputAdapter, error) {
		return walkingSkeletonInput{}, nil
	}); err != nil {
		t.Fatalf("register fake input: %v", err)
	}
	if err := catalog.Chunkers.RegisterWithSpec(ModuleSpec{
		Key:      "fake/chunk",
		Stage:    StageChunk,
		Requires: []string{"plain_text"},
		Provides: []string{"chunks"},
	}, func() (contracts.Chunker, error) {
		return walkingSkeletonChunker{}, nil
	}); err != nil {
		t.Fatalf("register fake chunker: %v", err)
	}
	if err := catalog.Extractors.RegisterLegacyRawWithSpec(ModuleSpec{
		Key:      "fake/extract",
		Stage:    StageExtract,
		Requires: []string{"chunks"},
		Provides: []string{"fake_artifacts"},
	}, func() (contracts.LegacyRawExtractor, error) {
		return walkingSkeletonExtractor{}, nil
	}); err != nil {
		t.Fatalf("register fake extractor: %v", err)
	}
	if err := catalog.Mergers.RegisterLegacyRawWithSpec(ModuleSpec{
		Key:      DefaultMergeModule,
		Stage:    StageMerge,
		Requires: []string{"fake_artifacts"},
	}, func() (contracts.LegacyRawMerger, error) {
		return walkingSkeletonMerger{}, nil
	}); err != nil {
		t.Fatalf("register append-order merger: %v", err)
	}
	if err := catalog.Normalizers.RegisterLegacyRawWithSpec(ModuleSpec{
		Key:   DefaultNormalizeModule,
		Stage: StageNormalize,
	}, func() (contracts.LegacyRawNormalizer, error) {
		return walkingSkeletonNormalizer{}, nil
	}); err != nil {
		t.Fatalf("register no-op normalizer: %v", err)
	}
	if err := catalog.Outputs.RegisterWithSpec(ModuleSpec{
		Key:   "json",
		Stage: StageOutput,
	}, func() (contracts.OutputEncoder, error) {
		return walkingSkeletonOutput{}, nil
	}); err != nil {
		t.Fatalf("register fake output: %v", err)
	}
	return catalog
}

func walkingSkeletonRegistries(t *testing.T) Registries {
	t.Helper()

	catalog := walkingSkeletonCatalog(t)
	return Registries{
		Inputs:         catalog.Inputs,
		Chunkers:       catalog.Chunkers,
		ArtifactCodecs: catalog.ArtifactCodecs,
		Extractors:     catalog.Extractors,
		Mergers:        catalog.Mergers,
		Normalizers:    catalog.Normalizers,
		Outputs:        catalog.Outputs,
	}
}

type walkingSkeletonInput struct{}

func (input walkingSkeletonInput) Key() string {
	return "fake/input"
}

func (input walkingSkeletonInput) Parse(ctx context.Context, req contracts.ParseRequest) (*source.SourceDocument, error) {
	var fixture struct {
		ID    string `json:"id"`
		Units []struct {
			ID   int    `json:"id"`
			Text string `json:"text"`
		} `json:"units"`
	}
	if err := json.Unmarshal(req.Raw, &fixture); err != nil {
		return nil, err
	}

	units := make([]source.SourceUnit, 0, len(fixture.Units))
	for _, unit := range fixture.Units {
		units = append(units, source.SourceUnit{
			ID:   unit.ID,
			Kind: "unit",
			Text: unit.Text,
			Ref: source.SourceRef{
				SourceID:    fixture.ID,
				StartUnitID: unit.ID,
				EndUnitID:   unit.ID,
			},
		})
	}
	return &source.SourceDocument{
		ID:     fixture.ID,
		Kind:   "fixture",
		Format: "application/json",
		Digest: rawDigest(req.Raw),
		Units:  units,
	}, nil
}

type walkingSkeletonChunker struct{}

func (chunker walkingSkeletonChunker) Key() string {
	return "fake/chunk"
}

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

func (chunker walkingSkeletonChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
	if len(req.Source.Units) < 3 {
		return contracts.ChunkResult{}, fmt.Errorf("fixture source must contain at least three units")
	}
	return contracts.ChunkResult{
		Chunks: []source.Chunk{
			{
				ID:        req.Source.ID + ":chunk:0",
				SourceID:  req.Source.ID,
				Index:     0,
				Ref:       source.SourceRef{SourceID: req.Source.ID, StartUnitID: req.Source.Units[0].ID, EndUnitID: req.Source.Units[1].ID},
				Content:   []byte(`{"units":[1,2]}`),
				MediaType: "application/json",
				Units:     append([]source.SourceUnit(nil), req.Source.Units[:2]...),
			},
			{
				ID:        req.Source.ID + ":chunk:1",
				SourceID:  req.Source.ID,
				Index:     1,
				Ref:       source.SourceRef{SourceID: req.Source.ID, StartUnitID: req.Source.Units[2].ID, EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID},
				Content:   []byte(`{"units":[3]}`),
				MediaType: "application/json",
				Units:     append([]source.SourceUnit(nil), req.Source.Units[2:]...),
			},
		},
	}, nil
}

type walkingSkeletonExtractor struct{}

func (extractor walkingSkeletonExtractor) Key() string {
	return "fake/extract"
}

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

func (extractor walkingSkeletonExtractor) Extract(ctx context.Context, req contracts.ExtractionRequest) (contracts.ExtractionResult, error) {
	var response struct {
		Call int `json:"call"`
	}
	if _, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
		StageName:     "fake/extract",
		PromptID:      "fake.event",
		PromptVersion: "v1",
	}, &response); err != nil {
		return contracts.ExtractionResult{}, err
	}

	payload, err := json.Marshal(map[string]any{
		"chunk_id": req.Chunk.ID,
		"llm_call": response.Call,
		"text":     chunkText(req.Chunk.Units),
	})
	if err != nil {
		return contracts.ExtractionResult{}, err
	}

	return contracts.ExtractionResult{
		Output: contracts.ExtractOutput{
			Schema: contracts.ResponseSchema{ID: "fake_event", Name: "fake_event", Version: "v1"},
			Payload: contracts.RawPayload{
				Content:   payload,
				MediaType: "application/json",
			},
		},
	}, nil
}

type walkingSkeletonLLMClient struct {
	calls int
}

func (client *walkingSkeletonLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	client.calls++
	if response, ok := out.(*struct {
		Call int `json:"call"`
	}); ok {
		response.Call = client.calls
	}
	content, err := json.Marshal(map[string]any{"call": client.calls})
	if err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	return contracts.StructuredCompletionResponse{
		Content: content,
	}, nil
}

type walkingSkeletonMerger struct{}

func (merger walkingSkeletonMerger) Key() string {
	return DefaultMergeModule
}

func (merger walkingSkeletonMerger) Merge(ctx context.Context, req contracts.MergeRequest) (contracts.MergeResult, error) {
	outputs := make([]json.RawMessage, 0, len(req.ExtractOutputs))
	for _, output := range req.ExtractOutputs {
		outputs = append(outputs, json.RawMessage(output.Payload.Content))
	}
	content, err := json.Marshal(map[string]any{"outputs": outputs})
	if err != nil {
		return contracts.MergeResult{}, err
	}
	return contracts.MergeResult{
		Output: contracts.MergeOutput{
			LaneID:   req.LaneID,
			SourceID: req.Source.ID,
			Schema:   contracts.ResponseSchema{ID: "fake_event", Name: "fake_event", Version: "v1"},
			Payload: contracts.RawPayload{
				Content:   content,
				MediaType: "application/json",
			},
		},
	}, nil
}

type walkingSkeletonNormalizer struct{}

func (normalizer walkingSkeletonNormalizer) Key() string {
	return DefaultNormalizeModule
}

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

func (normalizer walkingSkeletonNormalizer) Normalize(ctx context.Context, req contracts.NormalizeRequest) (contracts.NormalizeResult, error) {
	var response struct {
		Call int `json:"call"`
	}
	if _, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
		StageName:     "fake/normalize",
		PromptID:      "fake.normalize",
		PromptVersion: "v1",
	}, &response); err != nil {
		return contracts.NormalizeResult{}, err
	}
	return contracts.NormalizeResult{
		Output: contracts.NormalizeOutput{
			LaneID:   req.LaneID,
			SourceID: req.MergeOutput.SourceID,
			Schema:   req.MergeOutput.Schema,
			Payload:  req.MergeOutput.Payload,
		},
	}, nil
}

type walkingSkeletonOutput struct{}

func (output walkingSkeletonOutput) Key() string {
	return "json"
}

func (output walkingSkeletonOutput) Encode(ctx context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) {
	type rawOutput struct {
		LaneID        string                   `json:"lane_id"`
		NormalizerKey string                   `json:"normalizer_key"`
		SourceID      string                   `json:"source_id"`
		Schema        contracts.ResponseSchema `json:"schema"`
		MediaType     string                   `json:"media_type"`
		Content       json.RawMessage          `json:"content"`
	}
	rawOutputs := make([]rawOutput, 0, len(req.NormalizeOutputs))
	for _, output := range req.NormalizeOutputs {
		rawOutputs = append(rawOutputs, rawOutput{
			LaneID:        output.LaneID,
			NormalizerKey: output.NormalizerKey,
			SourceID:      output.SourceID,
			Schema:        contracts.ResponseSchema{ID: output.Artifact.Schema.ID, Name: output.Artifact.Schema.Name, Version: output.Artifact.Schema.Version},
			MediaType:     output.Artifact.MediaType,
			Content:       json.RawMessage(output.Artifact.Content),
		})
	}
	encoded, err := json.Marshal(struct {
		Manifest         artifacts.RunManifest `json:"manifest"`
		NormalizeOutputs []rawOutput           `json:"normalize_outputs"`
	}{
		Manifest: artifacts.RunManifest{
			PipelineID:       req.Manifest.PipelineID,
			PipelineDigest:   req.Manifest.PipelineDigest,
			ArtifactLanes:    req.Manifest.ArtifactLanes,
			ValidationStatus: req.Manifest.ValidationStatus,
		},
		NormalizeOutputs: rawOutputs,
	})
	if err != nil {
		return contracts.OutputResult{}, err
	}
	return contracts.OutputResult{
		Files: []contracts.OutputFile{
			{Name: "output.json", ContentType: "application/json", Bytes: encoded},
		},
	}, nil
}

func readTestFixture(t *testing.T, path string) []byte {
	t.Helper()

	bytes, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("read fixture %q: %v", path, err)
	}
	return bytes
}

func assertStructuralJSONEqual(t *testing.T, gotBytes, wantBytes []byte) {
	t.Helper()

	var got any
	if err := json.Unmarshal(gotBytes, &got); err != nil {
		t.Fatalf("unmarshal actual JSON: %v\n%s", err, gotBytes)
	}
	var want any
	if err := json.Unmarshal(wantBytes, &want); err != nil {
		t.Fatalf("unmarshal expected JSON: %v\n%s", err, wantBytes)
	}
	if !reflect.DeepEqual(got, want) {
		gotFormatted, _ := json.MarshalIndent(got, "", "  ")
		wantFormatted, _ := json.MarshalIndent(want, "", "  ")
		t.Fatalf("actual JSON:\n%s\nwant:\n%s", gotFormatted, wantFormatted)
	}
}

func chunkText(units []source.SourceUnit) string {
	parts := make([]string, 0, len(units))
	for _, unit := range units {
		parts = append(parts, unit.Text)
	}
	return strings.Join(parts, " ")
}

func rawDigest(raw []byte) string {
	sum := sha256.Sum256(raw)
	return "sha256:" + hex.EncodeToString(sum[:])
}