notarius/internal/framework/pipeline/profile.go

package pipeline

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sort"
	"strings"

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

const (
	DefaultChunkModule     = "generic"
	DefaultMergeModule     = "appendorder"
	DefaultNormalizeModule = "noop"
	DefaultOutputModule    = "json"
	DefaultLLMProfile      = "default"
)

type ModuleBinding struct {
	Module     string         `json:"module"`
	LLMProfile string         `json:"llm_profile,omitempty"`
	Options    map[string]any `json:"options,omitempty"`
}

type ArtifactLaneProfile struct {
	Extract    ModuleBinding     `json:"extract"`
	Merge      ModuleBinding     `json:"merge,omitempty"`
	Normalize  ModuleBinding     `json:"normalize,omitempty"`
	Validators []ModuleBinding   `json:"validators,omitempty"`
	References map[string]string `json:"references,omitempty"`
}

type PipelineProfile struct {
	ID         string                         `json:"id"`
	Input      ModuleBinding                  `json:"input"`
	Chunk      ModuleBinding                  `json:"chunk,omitempty"`
	Artifacts  map[string]ArtifactLaneProfile `json:"artifacts"`
	Output     ModuleBinding                  `json:"output,omitempty"`
	References map[string]string              `json:"references,omitempty"`
}

type ResolveOptions struct {
	Only               []string
	ReferenceOverrides []ReferenceBinding
	ReferenceUnbinds   []ReferenceUnbind
}

type ReferenceBinding struct {
	LaneID        string `json:"lane_id,omitempty"`
	SlotName      string `json:"slot_name"`
	Source        string `json:"source"`
	BindingSource string `json:"binding_source,omitempty"`
}

type ReferenceUnbind struct {
	LaneID   string `json:"lane_id"`
	SlotName string `json:"slot_name"`
}

type ResolvedArtifactLane struct {
	ID         string
	Extract    ModuleBinding
	Merge      ModuleBinding
	Normalize  ModuleBinding
	Validators []ModuleBinding
	References []ReferenceBinding `json:"references,omitempty"`
}

type ResolvedPipeline struct {
	ID            string
	Digest        string
	Input         ModuleBinding
	Chunk         ModuleBinding
	ArtifactLanes []ResolvedArtifactLane
	Output        ModuleBinding
}

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

func Binding(module string) ModuleBinding {
	return ModuleBinding{Module: strings.TrimSpace(module)}
}

func ResolvePipeline(profile PipelineProfile, options ResolveOptions, catalog ModuleCatalog) (ResolvedPipeline, error) {
	pipelineID := strings.TrimSpace(profile.ID)
	if pipelineID == "" {
		return ResolvedPipeline{}, fmt.Errorf("pipeline id must not be empty")
	}

	if len(profile.Artifacts) == 0 {
		return ResolvedPipeline{}, fmt.Errorf("pipeline %q must declare at least one artifact lane", pipelineID)
	}

	input := resolveBinding(profile.Input, "")
	if input.Module == "" {
		return ResolvedPipeline{}, fmt.Errorf("pipeline %q input module must not be empty", pipelineID)
	}
	inputModuleSpec, err := inputSpec(catalog, input.Module)
	if err != nil {
		return ResolvedPipeline{}, moduleLookupError(pipelineID, "", StageInput, input.Module, err)
	}

	capabilities := newCapabilitySet()
	if missing, ok := capabilities.missing(inputModuleSpec.Requires); ok {
		return ResolvedPipeline{}, capabilityError(pipelineID, "", StageInput, input.Module, missing)
	}
	capabilities.add(inputModuleSpec.Provides...)

	chunk := resolveBinding(profile.Chunk, DefaultChunkModule)
	chunkSpec, err := chunkerSpec(catalog, chunk.Module)
	if err != nil {
		return ResolvedPipeline{}, moduleLookupError(pipelineID, "", StageChunk, chunk.Module, err)
	}
	if missing, ok := capabilities.missing(chunkSpec.Requires); ok {
		return ResolvedPipeline{}, capabilityError(pipelineID, "", StageChunk, chunk.Module, missing)
	}
	capabilities.add(chunkSpec.Provides...)

	lanesByID, selectedLaneIDs, err := selectedArtifactLanes(pipelineID, profile.Artifacts, options)
	if err != nil {
		return ResolvedPipeline{}, err
	}
	if len(selectedLaneIDs) == 0 {
		return ResolvedPipeline{}, fmt.Errorf("pipeline %q must select at least one artifact lane", pipelineID)
	}

	resolved := ResolvedPipeline{
		ID:     pipelineID,
		Input:  input,
		Chunk:  chunk,
		Output: resolveBinding(profile.Output, DefaultOutputModule),
	}
	outputCapabilities := capabilities.clone()

	for _, laneID := range selectedLaneIDs {
		laneProfile := lanesByID[laneID]
		lane, laneCapabilities, err := resolveArtifactLane(pipelineID, laneID, laneProfile, profile.References, options, capabilities, catalog)
		if err != nil {
			return ResolvedPipeline{}, err
		}
		resolved.ArtifactLanes = append(resolved.ArtifactLanes, lane)
		outputCapabilities.addSet(laneCapabilities)
	}

	outputSpec, err := outputSpec(catalog, resolved.Output.Module)
	if err != nil {
		return ResolvedPipeline{}, moduleLookupError(pipelineID, "", StageOutput, resolved.Output.Module, err)
	}
	if missing, ok := outputCapabilities.missing(outputSpec.Requires); ok {
		return ResolvedPipeline{}, capabilityError(pipelineID, "", StageOutput, resolved.Output.Module, missing)
	}

	digest, err := resolvedPipelineDigest(resolved)
	if err != nil {
		return ResolvedPipeline{}, fmt.Errorf("pipeline %q digest: %w", pipelineID, err)
	}
	resolved.Digest = digest
	return resolved, nil
}

func resolveArtifactLane(
	pipelineID string,
	laneID string,
	profile ArtifactLaneProfile,
	pipelineReferences map[string]string,
	options ResolveOptions,
	inherited capabilitySet,
	catalog ModuleCatalog,
) (ResolvedArtifactLane, capabilitySet, error) {
	lane := ResolvedArtifactLane{
		ID:         laneID,
		Extract:    resolveBinding(profile.Extract, ""),
		Merge:      resolveBinding(profile.Merge, DefaultMergeModule),
		Normalize:  resolveBinding(profile.Normalize, DefaultNormalizeModule),
		Validators: resolveBindings(profile.Validators, ""),
	}
	if lane.Extract.Module == "" {
		return ResolvedArtifactLane{}, nil, fmt.Errorf("pipeline %q lane %q extract module must not be empty", pipelineID, laneID)
	}

	capabilities := inherited.clone()

	extractSpec, err := extractorSpec(catalog, lane.Extract.Module)
	if err != nil {
		return ResolvedArtifactLane{}, nil, moduleLookupError(pipelineID, laneID, StageExtract, lane.Extract.Module, err)
	}
	if missing, ok := capabilities.missing(extractSpec.Requires); ok {
		return ResolvedArtifactLane{}, nil, capabilityError(pipelineID, laneID, StageExtract, lane.Extract.Module, missing)
	}
	references, err := resolveReferenceBindings(pipelineID, laneID, lane.Extract.Module, extractSpec.ReferenceSlots, pipelineReferences, profile.References, options)
	if err != nil {
		return ResolvedArtifactLane{}, nil, err
	}
	lane.References = references
	capabilities.add(extractSpec.Provides...)

	mergeSpec, err := mergerSpec(catalog, lane.Merge.Module)
	if err != nil {
		return ResolvedArtifactLane{}, nil, moduleLookupError(pipelineID, laneID, StageMerge, lane.Merge.Module, err)
	}
	if missing, ok := capabilities.missing(mergeSpec.Requires); ok {
		return ResolvedArtifactLane{}, nil, capabilityError(pipelineID, laneID, StageMerge, lane.Merge.Module, missing)
	}
	capabilities.add(mergeSpec.Provides...)

	normalizeSpec, err := normalizerSpec(catalog, lane.Normalize.Module)
	if err != nil {
		return ResolvedArtifactLane{}, nil, moduleLookupError(pipelineID, laneID, StageNormalize, lane.Normalize.Module, err)
	}
	if missing, ok := capabilities.missing(normalizeSpec.Requires); ok {
		return ResolvedArtifactLane{}, nil, capabilityError(pipelineID, laneID, StageNormalize, lane.Normalize.Module, missing)
	}
	capabilities.add(normalizeSpec.Provides...)

	for _, validator := range lane.Validators {
		validatorSpec, err := validatorSpec(catalog, validator.Module)
		if err != nil {
			return ResolvedArtifactLane{}, nil, moduleLookupError(pipelineID, laneID, StageValidate, validator.Module, err)
		}
		if missing, ok := capabilities.missing(validatorSpec.Requires); ok {
			return ResolvedArtifactLane{}, nil, capabilityError(pipelineID, laneID, StageValidate, validator.Module, missing)
		}
		capabilities.add(validatorSpec.Provides...)
	}

	return lane, capabilities, nil
}

func resolveReferenceBindings(
	pipelineID string,
	laneID string,
	extractorModule string,
	slots []contracts.ReferenceSlot,
	pipelineReferences map[string]string,
	laneReferences map[string]string,
	options ResolveOptions,
) ([]ReferenceBinding, error) {
	slotByName := make(map[string]contracts.ReferenceSlot, len(slots))
	for _, slot := range slots {
		slotByName[slot.Name] = slot
	}

	bindings := make(map[string]ReferenceBinding)
	addBinding := func(slotName, source, bindingSource string) error {
		slotName = strings.TrimSpace(slotName)
		source = strings.TrimSpace(source)
		if slotName == "" {
			return fmt.Errorf("pipeline %q lane %q reference slot name must not be empty", pipelineID, laneID)
		}
		if source == "" {
			return fmt.Errorf("pipeline %q lane %q reference slot %q source must not be empty", pipelineID, laneID, slotName)
		}
		if _, ok := slotByName[slotName]; !ok {
			return fmt.Errorf("pipeline %q lane %q reference slot %q is not declared by extractor %q", pipelineID, laneID, slotName, extractorModule)
		}
		bindings[slotName] = ReferenceBinding{
			LaneID:        laneID,
			SlotName:      slotName,
			Source:        source,
			BindingSource: bindingSource,
		}
		return nil
	}

	normalizedPipelineReferences, err := normalizedReferenceMap(pipelineReferences, fmt.Sprintf("pipeline %q reference slot", pipelineID))
	if err != nil {
		return nil, err
	}
	for _, slotName := range sortedStringMapKeys(normalizedPipelineReferences) {
		if _, ok := slotByName[slotName]; !ok {
			continue
		}
		if err := addBinding(slotName, normalizedPipelineReferences[slotName], contracts.ReferenceBindingSourceConfig); err != nil {
			return nil, err
		}
	}

	normalizedLaneReferences, err := normalizedReferenceMap(laneReferences, fmt.Sprintf("pipeline %q lane %q reference slot", pipelineID, laneID))
	if err != nil {
		return nil, err
	}
	for _, slotName := range sortedStringMapKeys(normalizedLaneReferences) {
		if err := addBinding(slotName, normalizedLaneReferences[slotName], contracts.ReferenceBindingSourceConfig); err != nil {
			return nil, err
		}
	}

	for _, override := range options.ReferenceOverrides {
		optionLaneID := strings.TrimSpace(override.LaneID)
		if optionLaneID == "" {
			return nil, fmt.Errorf("pipeline %q reference override lane id must not be empty", pipelineID)
		}
		if optionLaneID != laneID {
			continue
		}
		source := override.BindingSource
		if strings.TrimSpace(source) == "" {
			source = contracts.ReferenceBindingSourceCLI
		}
		if err := addBinding(override.SlotName, override.Source, strings.TrimSpace(source)); err != nil {
			return nil, err
		}
	}

	for _, unbind := range options.ReferenceUnbinds {
		optionLaneID := strings.TrimSpace(unbind.LaneID)
		if optionLaneID == "" {
			return nil, fmt.Errorf("pipeline %q reference unbind lane id must not be empty", pipelineID)
		}
		if optionLaneID != laneID {
			continue
		}
		slotName := strings.TrimSpace(unbind.SlotName)
		if slotName == "" {
			return nil, fmt.Errorf("pipeline %q lane %q reference unbind slot name must not be empty", pipelineID, laneID)
		}
		if _, ok := slotByName[slotName]; !ok {
			return nil, fmt.Errorf("pipeline %q lane %q reference slot %q is not declared", pipelineID, laneID, slotName)
		}
		delete(bindings, slotName)
	}

	for _, slot := range slots {
		if slot.Required {
			if _, ok := bindings[slot.Name]; !ok {
				return nil, fmt.Errorf("pipeline %q lane %q required reference slot %q is not bound", pipelineID, laneID, slot.Name)
			}
		}
	}

	keys := sortedReferenceBindingKeys(bindings)
	resolved := make([]ReferenceBinding, 0, len(keys))
	for _, slotName := range keys {
		resolved = append(resolved, bindings[slotName])
	}
	return resolved, nil
}

func normalizedReferenceMap(values map[string]string, keyName string) (map[string]string, error) {
	if len(values) == 0 {
		return nil, nil
	}
	out := make(map[string]string, len(values))
	for rawSlotName, rawSource := range values {
		slotName := strings.TrimSpace(rawSlotName)
		if slotName == "" {
			return nil, fmt.Errorf("%s must not be empty", keyName)
		}
		if _, ok := out[slotName]; ok {
			return nil, fmt.Errorf("%s %q is duplicated after trimming", keyName, slotName)
		}
		source := strings.TrimSpace(rawSource)
		if source == "" {
			return nil, fmt.Errorf("%s %q source must not be empty", keyName, slotName)
		}
		out[slotName] = source
	}
	return out, nil
}

func sortedStringMapKeys(values map[string]string) []string {
	if len(values) == 0 {
		return nil
	}
	keys := make([]string, 0, len(values))
	for key := range values {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	return keys
}

func sortedReferenceBindingKeys(values map[string]ReferenceBinding) []string {
	if len(values) == 0 {
		return nil
	}
	keys := make([]string, 0, len(values))
	for key := range values {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	return keys
}

func resolveBinding(binding ModuleBinding, defaultModule string) ModuleBinding {
	module := strings.TrimSpace(binding.Module)
	if module == "" {
		module = defaultModule
	}
	llmProfile := strings.TrimSpace(binding.LLMProfile)
	if llmProfile == "" {
		llmProfile = DefaultLLMProfile
	}
	return ModuleBinding{
		Module:     module,
		LLMProfile: llmProfile,
		Options:    cloneOptions(binding.Options),
	}
}

func resolveBindings(bindings []ModuleBinding, defaultModule string) []ModuleBinding {
	if len(bindings) == 0 {
		return nil
	}

	resolved := make([]ModuleBinding, 0, len(bindings))
	for _, binding := range bindings {
		resolvedBinding := resolveBinding(binding, defaultModule)
		resolved = append(resolved, resolvedBinding)
	}
	return resolved
}

func cloneOptions(options map[string]any) map[string]any {
	if len(options) == 0 {
		return nil
	}

	copied := make(map[string]any, len(options))
	for key, value := range options {
		copied[key] = value
	}
	return copied
}

func selectedArtifactLanes(pipelineID string, artifacts map[string]ArtifactLaneProfile, options ResolveOptions) (map[string]ArtifactLaneProfile, []string, error) {
	lanesByID := make(map[string]ArtifactLaneProfile, len(artifacts))
	for rawLaneID, lane := range artifacts {
		laneID := strings.TrimSpace(rawLaneID)
		if laneID == "" {
			return nil, nil, fmt.Errorf("pipeline %q artifact lane id must not be empty", pipelineID)
		}
		if _, ok := lanesByID[laneID]; ok {
			return nil, nil, fmt.Errorf("pipeline %q artifact lane %q is duplicated after trimming", pipelineID, laneID)
		}
		lanesByID[laneID] = lane
	}

	if len(options.Only) == 0 {
		keys := make([]string, 0, len(lanesByID))
		for laneID := range lanesByID {
			keys = append(keys, laneID)
		}
		sort.Strings(keys)
		return lanesByID, keys, nil
	}

	selected := make(map[string]struct{}, len(options.Only))
	for _, rawLaneID := range options.Only {
		laneID := strings.TrimSpace(rawLaneID)
		if laneID == "" {
			return nil, nil, fmt.Errorf("pipeline %q selected artifact lane id must not be empty", pipelineID)
		}
		if _, ok := lanesByID[laneID]; !ok {
			return nil, nil, fmt.Errorf("pipeline %q selected artifact lane %q is not declared", pipelineID, laneID)
		}
		selected[laneID] = struct{}{}
	}

	keys := make([]string, 0, len(selected))
	for laneID := range selected {
		keys = append(keys, laneID)
	}
	sort.Strings(keys)
	return lanesByID, keys, nil
}

func resolvedPipelineDigest(resolved ResolvedPipeline) (string, error) {
	withoutDigest := struct {
		ID            string
		Input         ModuleBinding
		Chunk         ModuleBinding
		ArtifactLanes []ResolvedArtifactLane
		Output        ModuleBinding
	}{
		ID:            resolved.ID,
		Input:         resolved.Input,
		Chunk:         resolved.Chunk,
		ArtifactLanes: resolved.ArtifactLanes,
		Output:        resolved.Output,
	}
	encoded, err := json.Marshal(withoutDigest)
	if err != nil {
		return "", err
	}

	sum := sha256.Sum256(encoded)
	return "sha256:" + hex.EncodeToString(sum[:]), nil
}

func inputSpec(catalog ModuleCatalog, key string) (ModuleSpec, error) {
	return registrySpec(catalog.Inputs, key)
}

func chunkerSpec(catalog ModuleCatalog, key string) (ModuleSpec, error) {
	return registrySpec(catalog.Chunkers, key)
}

func extractorSpec(catalog ModuleCatalog, key string) (ModuleSpec, error) {
	return registrySpec(catalog.Extractors, key)
}

func mergerSpec(catalog ModuleCatalog, key string) (ModuleSpec, error) {
	return registrySpec(catalog.Mergers, key)
}

func normalizerSpec(catalog ModuleCatalog, key string) (ModuleSpec, error) {
	return registrySpec(catalog.Normalizers, key)
}

func validatorSpec(catalog ModuleCatalog, key string) (ModuleSpec, error) {
	return registrySpec(catalog.Validators, key)
}

func outputSpec(catalog ModuleCatalog, key string) (ModuleSpec, error) {
	return registrySpec(catalog.Outputs, key)
}

type specRegistry interface {
	Spec(key string) (ModuleSpec, bool)
}

func registrySpec(registry specRegistry, key string) (ModuleSpec, error) {
	if registry == nil {
		return ModuleSpec{}, fmt.Errorf("module %q is not registered", key)
	}
	spec, ok := registry.Spec(key)
	if !ok {
		return ModuleSpec{}, fmt.Errorf("module %q is not registered", key)
	}
	return spec, nil
}

func moduleLookupError(pipelineID, laneID string, stage ModuleStage, module string, err error) error {
	if laneID != "" {
		return fmt.Errorf("pipeline %q lane %q %s module %q: %w", pipelineID, laneID, stage, module, err)
	}
	return fmt.Errorf("pipeline %q %s module %q: %w", pipelineID, stage, module, err)
}

func capabilityError(pipelineID, laneID string, stage ModuleStage, module, capability string) error {
	if laneID != "" {
		return fmt.Errorf("pipeline %q lane %q %s module %q requires missing capability %q", pipelineID, laneID, stage, module, capability)
	}
	return fmt.Errorf("pipeline %q %s module %q requires missing capability %q", pipelineID, stage, module, capability)
}

type capabilitySet map[string]struct{}

func newCapabilitySet() capabilitySet {
	return make(capabilitySet)
}

func (set capabilitySet) clone() capabilitySet {
	copied := make(capabilitySet, len(set))
	for capability := range set {
		copied[capability] = struct{}{}
	}
	return copied
}

func (set capabilitySet) add(values ...string) {
	for _, value := range values {
		set[value] = struct{}{}
	}
}

func (set capabilitySet) addSet(other capabilitySet) {
	for value := range other {
		set[value] = struct{}{}
	}
}

func (set capabilitySet) missing(required []string) (string, bool) {
	for _, capability := range required {
		if _, ok := set[capability]; !ok {
			return capability, true
		}
	}
	return "", false
}