package pipeline

import (
	"context"
	"fmt"

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

type Registries struct {
	Inputs      *InputAdapterRegistry
	Chunkers    *ChunkerRegistry
	Extractors  *ExtractorRegistry
	Mergers     *MergerRegistry
	Normalizers *NormalizerRegistry
	Validators  *ValidatorRegistry
	Outputs     *OutputEncoderRegistry
}

type Runner struct {
	registries Registries
}

func New(registries Registries) *Runner {
	return &Runner{registries: registries}
}

type RunInput struct {
	Pipeline  ResolvedPipeline
	SourceID  string
	Path      string
	RawInput  []byte
	LLMClient contracts.StructuredLLMClient
	Metadata  map[string]any
}

type RunOutput struct {
	Manifest      artifacts.RunManifest        `json:"manifest"`
	Approved      []artifacts.Artifact         `json:"approved,omitempty"`
	Rejected      []artifacts.RejectedArtifact `json:"rejected,omitempty"`
	Warnings      []contracts.Warning          `json:"warnings,omitempty"`
	EncodedOutput []byte                       `json:"-"`
	ContentType   string                       `json:"content_type,omitempty"`
}

func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
	var output RunOutput
	if r == nil {
		return output, fmt.Errorf("runner must not be nil")
	}
	if err := validateRunInput(input); err != nil {
		return output, err
	}
	if err := r.validateRegistries(input.Pipeline); err != nil {
		return output, err
	}

	output.Manifest = manifestFromPipeline(input.Pipeline)

	adapter, err := r.registries.Inputs.Build(input.Pipeline.Input.Module)
	if err != nil {
		return failOutput(output), fmt.Errorf("build input adapter %q: %w", input.Pipeline.Input.Module, err)
	}
	doc, err := adapter.Parse(ctx, contracts.ParseRequest{
		SourceID: input.SourceID,
		Path:     input.Path,
		Raw:      input.RawInput,
		Metadata: input.Metadata,
	})
	if err != nil {
		return failOutput(output), fmt.Errorf("parse input with adapter %q: %w", adapter.Key(), err)
	}
	if err := source.ValidateDocument(doc); err != nil {
		return failOutput(output), fmt.Errorf("validate source document: %w", err)
	}
	output.Manifest.SourceDigests = []string{doc.Digest}

	chunker, err := r.registries.Chunkers.Build(input.Pipeline.Chunk.Module)
	if err != nil {
		return failOutput(output), fmt.Errorf("build chunker %q: %w", input.Pipeline.Chunk.Module, err)
	}
	chunkResult, err := chunker.Chunk(ctx, contracts.ChunkRequest{
		Source:   doc,
		Metadata: input.Metadata,
	})
	output.Warnings = append(output.Warnings, chunkResult.Warnings...)
	if err != nil {
		return failOutput(output), fmt.Errorf("chunk source with chunker %q: %w", chunker.Key(), err)
	}
	if len(chunkResult.Chunks) == 0 {
		return failOutput(output), fmt.Errorf("chunker %q returned no chunks", chunker.Key())
	}

	nextCandidateIndex := 0
	for _, lane := range input.Pipeline.ArtifactLanes {
		if err := r.runLane(ctx, input, doc, chunkResult.Chunks, lane, &output, &nextCandidateIndex); err != nil {
			return failOutput(output), err
		}
	}

	if len(output.Rejected) > 0 {
		output.Manifest.ValidationStatus = "rejected"
	} else {
		output.Manifest.ValidationStatus = "approved"
	}

	encoder, err := r.registries.Outputs.Build(input.Pipeline.Output.Module)
	if err != nil {
		return failOutput(output), fmt.Errorf("build output encoder %q: %w", input.Pipeline.Output.Module, err)
	}
	encoded, err := encoder.Encode(ctx, contracts.OutputRequest{
		Manifest: output.Manifest,
		Approved: output.Approved,
		Rejected: output.Rejected,
		Warnings: output.Warnings,
		Metadata: input.Metadata,
	})
	output.Warnings = append(output.Warnings, encoded.Warnings...)
	if err != nil {
		return failOutput(output), fmt.Errorf("encode output with encoder %q: %w", encoder.Key(), err)
	}
	output.EncodedOutput = encoded.Bytes
	output.ContentType = encoded.ContentType

	return output, nil
}

func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.SourceDocument, chunks []contracts.SourceChunk, lane ResolvedArtifactLane, output *RunOutput, nextCandidateIndex *int) error {
	extractor, err := r.registries.Extractors.Build(lane.Extract.Module)
	if err != nil {
		return fmt.Errorf("build extractor %q for lane %q: %w", lane.Extract.Module, lane.ID, err)
	}
	merger, err := r.registries.Mergers.Build(lane.Merge.Module)
	if err != nil {
		return fmt.Errorf("build merger %q for lane %q: %w", lane.Merge.Module, lane.ID, err)
	}
	normalizer, err := r.registries.Normalizers.Build(lane.Normalize.Module)
	if err != nil {
		return fmt.Errorf("build normalizer %q for lane %q: %w", lane.Normalize.Module, lane.ID, err)
	}

	var validators []contracts.Validator
	if len(lane.Validators) > 0 {
		validators, err = r.buildConfiguredValidators(lane)
		if err != nil {
			return err
		}
	} else {
		validators = extractor.Validators()
	}

	chunkArtifacts := make([]contracts.ChunkArtifacts, 0, len(chunks))
	for index := range chunks {
		chunk := chunks[index]
		result, err := extractor.Extract(ctx, contracts.ExtractionRequest{
			Source:    doc,
			Chunk:     &chunk,
			LLMClient: input.LLMClient,
			Metadata:  input.Metadata,
		})
		output.Warnings = append(output.Warnings, result.Warnings...)
		if err != nil {
			return fmt.Errorf("extract lane %q chunk %q with extractor %q: %w", lane.ID, chunk.ID, extractor.Key(), err)
		}

		candidates, err := normalizeCandidates(extractor, result.Candidates, nextCandidateIndex)
		if err != nil {
			return err
		}
		chunkArtifacts = append(chunkArtifacts, contracts.ChunkArtifacts{
			Chunk:      chunk,
			Candidates: candidates,
		})
	}

	mergeResult, err := merger.Merge(ctx, contracts.MergeRequest{
		Source:         doc,
		LaneID:         lane.ID,
		ChunkArtifacts: chunkArtifacts,
		Metadata:       input.Metadata,
	})
	output.Warnings = append(output.Warnings, mergeResult.Warnings...)
	if err != nil {
		return fmt.Errorf("merge lane %q with merger %q: %w", lane.ID, merger.Key(), err)
	}

	normalizeResult, err := normalizer.Normalize(ctx, contracts.NormalizeRequest{
		Source:     doc,
		LaneID:     lane.ID,
		Candidates: mergeResult.Candidates,
		Metadata:   input.Metadata,
	})
	output.Warnings = append(output.Warnings, normalizeResult.Warnings...)
	if err != nil {
		return fmt.Errorf("normalize lane %q with normalizer %q: %w", lane.ID, normalizer.Key(), err)
	}

	approved, rejected, warnings, err := runValidators(ctx, extractor.Key(), validators, doc, normalizeResult.Candidates, input.Metadata)
	output.Warnings = append(output.Warnings, warnings...)
	output.Rejected = append(output.Rejected, rejected...)
	if err != nil {
		return err
	}

	for _, candidate := range approved {
		output.Approved = append(output.Approved, artifacts.ArtifactFromCandidate(candidate))
	}
	return nil
}

func (r *Runner) buildConfiguredValidators(lane ResolvedArtifactLane) ([]contracts.Validator, error) {
	validators := make([]contracts.Validator, 0, len(lane.Validators))
	for _, binding := range lane.Validators {
		validator, err := r.registries.Validators.Build(binding.Module)
		if err != nil {
			return nil, fmt.Errorf("build validator %q for lane %q: %w", binding.Module, lane.ID, err)
		}
		validators = append(validators, validator)
	}
	return validators, nil
}

func (r *Runner) validateRegistries(pipeline ResolvedPipeline) error {
	if r.registries.Inputs == nil {
		return fmt.Errorf("input registry must not be nil")
	}
	if r.registries.Chunkers == nil {
		return fmt.Errorf("chunker registry must not be nil")
	}
	if r.registries.Extractors == nil {
		return fmt.Errorf("extractor registry must not be nil")
	}
	if r.registries.Mergers == nil {
		return fmt.Errorf("merger registry must not be nil")
	}
	if r.registries.Normalizers == nil {
		return fmt.Errorf("normalizer registry must not be nil")
	}
	if r.registries.Outputs == nil {
		return fmt.Errorf("output encoder registry must not be nil")
	}
	if pipelineUsesConfiguredValidators(pipeline) && r.registries.Validators == nil {
		return fmt.Errorf("validator registry must not be nil")
	}
	return nil
}

func validateRunInput(input RunInput) error {
	if input.Pipeline.ID == "" {
		return fmt.Errorf("resolved pipeline id must not be empty")
	}
	if input.Pipeline.Digest == "" {
		return fmt.Errorf("resolved pipeline digest must not be empty")
	}
	if input.Pipeline.Input.Module == "" {
		return fmt.Errorf("resolved pipeline input module must not be empty")
	}
	if input.Pipeline.Chunk.Module == "" {
		return fmt.Errorf("resolved pipeline chunk module must not be empty")
	}
	if input.Pipeline.Output.Module == "" {
		return fmt.Errorf("resolved pipeline output module must not be empty")
	}
	if len(input.Pipeline.ArtifactLanes) == 0 {
		return fmt.Errorf("resolved pipeline artifact lanes must not be empty")
	}
	for _, lane := range input.Pipeline.ArtifactLanes {
		if lane.ID == "" {
			return fmt.Errorf("resolved pipeline artifact lane id must not be empty")
		}
		if lane.Extract.Module == "" {
			return fmt.Errorf("resolved pipeline lane %q extract module must not be empty", lane.ID)
		}
		if lane.Merge.Module == "" {
			return fmt.Errorf("resolved pipeline lane %q merge module must not be empty", lane.ID)
		}
		if lane.Normalize.Module == "" {
			return fmt.Errorf("resolved pipeline lane %q normalize module must not be empty", lane.ID)
		}
		for _, validator := range lane.Validators {
			if validator.Module == "" {
				return fmt.Errorf("resolved pipeline lane %q validator module must not be empty", lane.ID)
			}
		}
	}
	return nil
}

func manifestFromPipeline(pipeline ResolvedPipeline) artifacts.RunManifest {
	manifest := artifacts.RunManifest{
		PipelineID:     pipeline.ID,
		PipelineDigest: pipeline.Digest,
		InputModule:    pipeline.Input.Module,
		Chunker:        pipeline.Chunk.Module,
		OutputEncoder:  pipeline.Output.Module,
		ArtifactLanes:  make([]artifacts.ArtifactLaneManifest, 0, len(pipeline.ArtifactLanes)),
	}

	for _, lane := range pipeline.ArtifactLanes {
		laneManifest := artifacts.ArtifactLaneManifest{
			ID:         lane.ID,
			Extractor:  lane.Extract.Module,
			Merger:     lane.Merge.Module,
			Normalizer: lane.Normalize.Module,
		}
		for _, validator := range lane.Validators {
			laneManifest.Validators = append(laneManifest.Validators, validator.Module)
		}
		manifest.ArtifactLanes = append(manifest.ArtifactLanes, laneManifest)
	}
	return manifest
}

func failOutput(output RunOutput) RunOutput {
	if output.Manifest.PipelineID != "" {
		output.Manifest.ValidationStatus = "failed"
	}
	return output
}

func pipelineUsesConfiguredValidators(pipeline ResolvedPipeline) bool {
	for _, lane := range pipeline.ArtifactLanes {
		if len(lane.Validators) > 0 {
			return true
		}
	}
	return false
}

func normalizeCandidates(extractor contracts.Extractor, candidates []artifacts.ArtifactCandidate, nextIndex *int) ([]artifacts.ArtifactCandidate, error) {
	normalized := make([]artifacts.ArtifactCandidate, 0, len(candidates))
	for _, candidate := range candidates {
		candidate.Index = *nextIndex
		*nextIndex = *nextIndex + 1

		if candidate.ExtractorKey == "" {
			candidate.ExtractorKey = extractor.Key()
		} else if candidate.ExtractorKey != extractor.Key() {
			return nil, fmt.Errorf("candidate extractor_key %q does not match extractor %q", candidate.ExtractorKey, extractor.Key())
		}

		if candidate.ArtifactType == "" {
			candidate.ArtifactType = extractor.ArtifactType()
		} else if candidate.ArtifactType != extractor.ArtifactType() {
			return nil, fmt.Errorf("candidate artifact_type %q does not match extractor %q artifact type %q", candidate.ArtifactType, extractor.Key(), extractor.ArtifactType())
		}

		if candidate.SchemaVersion == "" {
			candidate.SchemaVersion = extractor.SchemaVersion()
		} else if candidate.SchemaVersion != extractor.SchemaVersion() {
			return nil, fmt.Errorf("candidate schema_version %q does not match extractor %q schema version %q", candidate.SchemaVersion, extractor.Key(), extractor.SchemaVersion())
		}

		normalized = append(normalized, candidate)
	}
	return normalized, nil
}

func runValidators(ctx context.Context, extractorKey string, validators []contracts.Validator, doc *source.SourceDocument, candidates []artifacts.ArtifactCandidate, metadata map[string]any) ([]artifacts.ArtifactCandidate, []artifacts.RejectedArtifact, []contracts.Warning, error) {
	eligible := candidates
	var rejected []artifacts.RejectedArtifact
	var warnings []contracts.Warning

	for validatorIndex, validator := range validators {
		if validator == nil {
			return nil, rejected, warnings, fmt.Errorf("extractor %q validator[%d] must not be nil", extractorKey, validatorIndex)
		}
		result, err := validator.Validate(ctx, contracts.ValidationRequest{
			Source:     doc,
			Candidates: eligible,
			Metadata:   metadata,
		})
		warnings = append(warnings, result.Warnings...)
		if err != nil {
			return nil, rejected, warnings, fmt.Errorf("validate extractor %q with validator %q: %w", extractorKey, validator.Name(), err)
		}
		if result.ValidatorName != validator.Name() {
			return nil, rejected, warnings, fmt.Errorf("validator %q returned result for %q", validator.Name(), result.ValidatorName)
		}
		if err := validate.EnforceDecisionCardinality(eligible, result.Decisions); err != nil {
			return nil, rejected, warnings, fmt.Errorf("validate extractor %q with validator %q: %w", extractorKey, validator.Name(), err)
		}

		decisions := make(map[int]contracts.ValidationDecision, len(result.Decisions))
		for _, decision := range result.Decisions {
			decisions[decision.CandidateIndex] = decision
		}

		nextEligible := make([]artifacts.ArtifactCandidate, 0, len(eligible))
		for _, candidate := range eligible {
			decision := decisions[candidate.Index]
			if decision.Approved {
				nextEligible = append(nextEligible, candidate)
				continue
			}
			rejected = append(rejected, artifacts.RejectedArtifact{
				Candidate:     candidate,
				ValidatorName: result.ValidatorName,
				ReasonCode:    decision.ReasonCode,
				Message:       decision.Message,
			})
		}
		eligible = nextEligible
	}

	return eligible, rejected, warnings, nil
}