notarius/internal/framework/pipeline/runner.go

package pipeline

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"mime"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"time"

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

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
	SessionID   string
	RunID       string
	StartedAt   time.Time
	LLMProfiles []artifacts.LLMProfileManifest
	Metadata    map[string]any
	Warnings    []contracts.Warning
}

type RunOutput struct {
	Manifest         artifacts.RunManifest       `json:"manifest"`
	NormalizeOutputs []contracts.NormalizeOutput `json:"normalize_outputs,omitempty"`
	Rejected         []contracts.RejectedOutput  `json:"rejected,omitempty"`
	Warnings         []contracts.Warning         `json:"warnings,omitempty"`
	OutputFiles      []contracts.OutputFile      `json:"-"`
}

func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err error) {
	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)
	defer func() {
		output.Manifest.LLMProfiles = mergeLLMProfileManifests(input.LLMProfiles, llmProfileManifests(input.LLMClient))
	}()
	output.Warnings = append(output.Warnings, cloneWarnings(input.Warnings)...)

	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)
	}
	attachModuleManifestMetadata(&output, "input", adapter)
	doc, err := adapter.Parse(ctx, contracts.ParseRequest{
		SourceID:   input.SourceID,
		Path:       input.Path,
		Raw:        input.RawInput,
		LLMProfile: input.Pipeline.Input.LLMProfile,
		Options:    cloneOptions(input.Pipeline.Input.Options),
		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)
	}
	sourceInput := sourceInputMaterial(input.Path, input.RawInput)
	sessionID := resolvedSessionID(input.SessionID, doc.ID)
	output.Manifest.Metadata = manifestMetadataWithSessionID(output.Manifest.Metadata, sessionID)
	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)
	}
	attachModuleManifestMetadata(&output, "chunker", chunker)
	chunkResult, err := chunker.Chunk(ctx, contracts.ChunkRequest{
		Source:      doc,
		SourceInput: sourceInput.Clone(),
		SessionID:   sessionID,
		References:  CloneReferenceSet(input.Pipeline.ChunkReferences.ReferenceSet),
		LLMClient:   input.LLMClient,
		LLMProfile:  input.Pipeline.Chunk.LLMProfile,
		Options:     cloneOptions(input.Pipeline.Chunk.Options),
		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())
	}
	canonicalChunks, err := validateAndCanonicalizeChunkResult(doc, chunkResult.Chunks)
	if err != nil {
		return failOutput(output), fmt.Errorf("validate chunks from chunker %q: %w", chunker.Key(), err)
	}

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

	if len(output.Rejected) > 0 {
		output.Manifest.ValidationStatus = "rejected"
	} else {
		output.Manifest.ValidationStatus = "approved"
	}
	output.Manifest.CompletedAt = timePtr(time.Now().UTC())

	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)
	}
	attachModuleManifestMetadata(&output, "output", encoder)
	encoded, err := encoder.Encode(ctx, contracts.OutputRequest{
		Manifest:         output.Manifest,
		NormalizeOutputs: cloneNormalizeOutputs(output.NormalizeOutputs),
		Rejected:         cloneRejectedOutputs(output.Rejected),
		Warnings:         output.Warnings,
		LLMProfile:       input.Pipeline.Output.LLMProfile,
		Options:          cloneOptions(input.Pipeline.Output.Options),
		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)
	}
	files, err := outputFilesFromResult(encoded)
	if err != nil {
		return failOutput(output), fmt.Errorf("validate output files from encoder %q: %w", encoder.Key(), err)
	}
	output.OutputFiles = files

	return output, nil
}

func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.SourceDocument, sourceInput contracts.LLMInputMaterial, sessionID string, chunks []contracts.SourceChunk, lane ResolvedArtifactLane, output *RunOutput) 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)
	}
	setLaneManifestMetadata(output, lane.ID, extractor, merger, normalizer)

	extractOutputs := make([]contracts.ExtractOutput, 0, len(chunks))
	for index := range chunks {
		chunk := chunks[index]
		result, err := extractor.Extract(ctx, contracts.ExtractionRequest{
			Source:      doc,
			Chunk:       &chunk,
			SourceInput: sourceInput.Clone(),
			SessionID:   sessionID,
			References:  CloneReferenceSet(lane.ExtractReferences.ReferenceSet),
			LLMClient:   input.LLMClient,
			LLMProfile:  lane.Extract.LLMProfile,
			Options:     cloneOptions(lane.Extract.Options),
			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)
		}
		extractOutput := result.Output
		extractOutput.LaneID = lane.ID
		extractOutput.ExtractorKey = extractor.Key()
		extractOutput.SourceID = doc.ID
		extractOutput.ChunkID = chunk.ID
		extractOutput.ChunkIndex = chunk.Index
		extractOutput.Payload.Warnings = append(extractOutput.Payload.Warnings, result.Warnings...)
		extractOutputs = append(extractOutputs, cloneExtractOutput(extractOutput))
	}

	mergeResult, err := merger.Merge(ctx, contracts.MergeRequest{
		Source:         doc,
		LaneID:         lane.ID,
		ExtractOutputs: cloneExtractOutputs(extractOutputs),
		LLMProfile:     lane.Merge.LLMProfile,
		Options:        cloneOptions(lane.Merge.Options),
		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,
		MergeOutput: cloneMergeOutput(mergeResult.Output),
		SourceInput: sourceInput.Clone(),
		SessionID:   sessionID,
		References:  CloneReferenceSet(lane.NormalizeReferences.ReferenceSet),
		LLMClient:   input.LLMClient,
		LLMProfile:  lane.Normalize.LLMProfile,
		Options:     cloneOptions(lane.Normalize.Options),
		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)
	}
	normalizeOutput := normalizeResult.Output
	normalizeOutput.LaneID = lane.ID
	normalizeOutput.NormalizerKey = normalizer.Key()
	normalizeOutput.SourceID = doc.ID
	normalizeOutput.Payload.Warnings = append(normalizeOutput.Payload.Warnings, normalizeResult.Warnings...)
	output.NormalizeOutputs = append(output.NormalizeOutputs, cloneNormalizeOutput(normalizeOutput))
	return nil
}

type validatorExecution struct {
	validator contracts.Validator
	binding   ModuleBinding
}

func (r *Runner) buildConfiguredValidators(lane ResolvedArtifactLane) ([]validatorExecution, error) {
	validators := make([]validatorExecution, 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, validatorExecution{
			validator: validator,
			binding:   binding,
		})
	}
	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(input RunInput) artifacts.RunManifest {
	startedAt := input.StartedAt
	if startedAt.IsZero() {
		startedAt = time.Now().UTC()
	}
	runID := strings.TrimSpace(input.RunID)
	if runID == "" {
		runID = fmt.Sprintf("run-%d", startedAt.UnixNano())
	}

	pipeline := input.Pipeline
	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)),
		RunID:          runID,
		StartedAt:      timePtr(startedAt),
		References:     ReferenceProvenance(pipeline),
		LLMProfiles:    cloneLLMProfiles(input.LLMProfiles),
	}
	// The runner does not currently maintain a cache or idempotency key. Reference
	// digests are recorded in manifest provenance and intentionally kept separate
	// from source_digests.

	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"
		output.Manifest.CompletedAt = timePtr(time.Now().UTC())
	}
	return output
}

func setLaneManifestMetadata(output *RunOutput, laneID string, modules ...any) {
	if output == nil {
		return
	}
	for i := range output.Manifest.ArtifactLanes {
		if output.Manifest.ArtifactLanes[i].ID != laneID {
			continue
		}

		metadata := make(map[string]any)
		for _, module := range modules {
			moduleMetadata, ok := moduleManifestMetadata(module)
			if !ok {
				continue
			}
			key := manifestMetadataKey(module)
			if key == "" {
				continue
			}
			metadata[key] = moduleMetadata
		}
		if len(metadata) > 0 {
			output.Manifest.ArtifactLanes[i].Metadata = metadata
		}
		return
	}
}

func attachModuleManifestMetadata(output *RunOutput, moduleKey string, module any) {
	if output == nil {
		return
	}
	moduleMetadata, ok := moduleManifestMetadata(module)
	if !ok {
		return
	}
	if output.Manifest.ModuleMetadata == nil {
		output.Manifest.ModuleMetadata = make(map[string]map[string]any)
	}
	output.Manifest.ModuleMetadata[moduleKey] = moduleMetadata
}

func manifestMetadataKey(module any) string {
	switch module.(type) {
	case contracts.Extractor:
		return "extractor"
	case contracts.Merger:
		return "merger"
	case contracts.Normalizer:
		return "normalizer"
	default:
		return ""
	}
}

func moduleManifestMetadata(module any) (map[string]any, bool) {
	provider, ok := module.(contracts.ManifestMetadataProvider)
	if !ok {
		return nil, false
	}

	moduleMetadata := cloneMetadata(provider.ManifestMetadata())
	if len(moduleMetadata) == 0 {
		return nil, false
	}
	return moduleMetadata, true
}

func outputFilesFromResult(result contracts.OutputResult) ([]contracts.OutputFile, error) {
	out := make([]contracts.OutputFile, 0, len(result.Files))
	for _, file := range result.Files {
		if err := validateOutputFileName(file.Name); err != nil {
			return nil, err
		}
		out = append(out, contracts.OutputFile{
			Name:        file.Name,
			ContentType: file.ContentType,
			Bytes:       append([]byte(nil), file.Bytes...),
		})
	}
	return out, nil
}

func validateOutputFileName(name string) error {
	if strings.TrimSpace(name) == "" {
		return fmt.Errorf("output file name must not be empty")
	}
	if strings.Contains(name, "\\") {
		return fmt.Errorf("output file name %q must use slash-separated relative paths", name)
	}
	if path.IsAbs(name) {
		return fmt.Errorf("output file name %q must be relative", name)
	}
	if strings.Contains(name, "..") {
		return fmt.Errorf("output file name %q must not contain ..", name)
	}
	cleaned := path.Clean(name)
	if cleaned == "." || cleaned != name {
		return fmt.Errorf("output file name %q must be clean", name)
	}
	return nil
}

func cloneMetadata(metadata map[string]any) map[string]any {
	if len(metadata) == 0 {
		return nil
	}
	out := make(map[string]any, len(metadata))
	for key, value := range metadata {
		out[key] = value
	}
	return out
}

func cloneLLMProfiles(profiles []artifacts.LLMProfileManifest) []artifacts.LLMProfileManifest {
	if len(profiles) == 0 {
		return nil
	}
	return append([]artifacts.LLMProfileManifest(nil), profiles...)
}

func llmProfileManifests(client contracts.StructuredLLMClient) []artifacts.LLMProfileManifest {
	provider, ok := client.(contracts.LLMProfileManifestProvider)
	if !ok {
		return nil
	}
	return provider.LLMProfileManifests()
}

func mergeLLMProfileManifests(sources ...[]artifacts.LLMProfileManifest) []artifacts.LLMProfileManifest {
	merged := make(map[string]artifacts.LLMProfileManifest)
	for _, source := range sources {
		for _, profile := range source {
			id := strings.TrimSpace(profile.ID)
			provider := strings.TrimSpace(profile.Provider)
			model := strings.TrimSpace(profile.Model)
			key := id + "\x00" + provider + "\x00" + model
			if _, exists := merged[key]; exists {
				continue
			}
			merged[key] = artifacts.LLMProfileManifest{
				ID:       id,
				Provider: provider,
				Model:    model,
			}
		}
	}
	if len(merged) == 0 {
		return nil
	}
	keys := make([]string, 0, len(merged))
	for key := range merged {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	out := make([]artifacts.LLMProfileManifest, 0, len(keys))
	for _, key := range keys {
		out = append(out, merged[key])
	}
	return out
}

func sourceInputMaterial(inputPath string, content []byte) contracts.LLMInputMaterial {
	return contracts.NewLLMInputMaterial(
		"source",
		sourceInputMediaType(inputPath),
		content,
		sourceInputDigest(content),
		sourceInputOriginURI(inputPath),
	)
}

func sourceInputMediaType(inputPath string) string {
	extension := strings.ToLower(filepath.Ext(strings.TrimSpace(inputPath)))
	if extension == ".json" {
		return "application/json"
	}
	mediaType := mime.TypeByExtension(extension)
	if strings.TrimSpace(mediaType) == "" {
		return unknownMediaType
	}
	return canonicalMediaType(mediaType)
}

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

func sourceInputOriginURI(inputPath string) string {
	if strings.TrimSpace(inputPath) == "" {
		return ""
	}
	return fileURI(inputPath)
}

func resolvedSessionID(explicit string, sourceDocumentID string) string {
	if trimmed := strings.TrimSpace(explicit); trimmed != "" {
		return trimmed
	}
	return strings.TrimSpace(sourceDocumentID)
}

func manifestMetadataWithSessionID(metadata map[string]any, sessionID string) map[string]any {
	out := cloneMetadata(metadata)
	if strings.TrimSpace(sessionID) == "" {
		return out
	}
	if out == nil {
		out = make(map[string]any)
	}
	out["session_id"] = sessionID
	return out
}

func cloneWarnings(warnings []contracts.Warning) []contracts.Warning {
	if len(warnings) == 0 {
		return nil
	}
	return append([]contracts.Warning(nil), warnings...)
}

func cloneRawPayload(payload contracts.RawPayload) contracts.RawPayload {
	return contracts.RawPayload{
		Content:   append([]byte(nil), payload.Content...),
		MediaType: payload.MediaType,
		Metadata:  cloneMetadata(payload.Metadata),
		Warnings:  cloneWarnings(payload.Warnings),
	}
}

func cloneExtractOutput(output contracts.ExtractOutput) contracts.ExtractOutput {
	output.Payload = cloneRawPayload(output.Payload)
	return output
}

func cloneExtractOutputs(outputs []contracts.ExtractOutput) []contracts.ExtractOutput {
	if len(outputs) == 0 {
		return nil
	}
	out := make([]contracts.ExtractOutput, 0, len(outputs))
	for _, output := range outputs {
		out = append(out, cloneExtractOutput(output))
	}
	return out
}

func cloneMergeOutput(output contracts.MergeOutput) contracts.MergeOutput {
	output.Payload = cloneRawPayload(output.Payload)
	return output
}

func cloneNormalizeOutput(output contracts.NormalizeOutput) contracts.NormalizeOutput {
	output.Payload = cloneRawPayload(output.Payload)
	return output
}

func cloneNormalizeOutputs(outputs []contracts.NormalizeOutput) []contracts.NormalizeOutput {
	if len(outputs) == 0 {
		return nil
	}
	out := make([]contracts.NormalizeOutput, 0, len(outputs))
	for _, output := range outputs {
		out = append(out, cloneNormalizeOutput(output))
	}
	return out
}

func cloneRejectedOutputs(rejected []contracts.RejectedOutput) []contracts.RejectedOutput {
	if len(rejected) == 0 {
		return nil
	}
	return append([]contracts.RejectedOutput(nil), rejected...)
}

func timePtr(t time.Time) *time.Time {
	return &t
}

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