notarius/internal/modules/dnd/normalize/itemregistry/normalizer.go

// Package itemregistry normalizes merged D&D item-registry candidates conservatively.
package itemregistry

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"sort"
	"strconv"
	"strings"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/items/identity"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared/diagnostics"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared/entityreconcile"
)

const (
	Key                   = "dnd/item-registry"
	PromptID              = "dnd.item_registry.normalize"
	normalizationPolicy   = "dnd.item_registry.normalize.v2"
	semanticContextPolicy = "dnd.entity_reconcile.context.v1"
	semanticContextRadius = 2
	NormalizationPolicy   = normalizationPolicy

	ReasonCodeItemFieldsNormalized                = "item_fields_normalized"
	ReasonCodeItemIDRecomputed                    = "item_id_recomputed"
	ReasonCodeSourceReferencesNormalized          = "source_references_normalized"
	ReasonCodeDuplicateItemCollapsed              = "duplicate_item_collapsed"
	ReasonCodeItemSemanticProposalInvalid         = "item_semantic_proposal_invalid"
	ReasonCodeItemSemanticReconciliationExhausted = "item_semantic_reconciliation_exhausted"
	ReasonCodeItemNormalizationWarningsOmitted    = "item_normalization_warnings_omitted"
)

var requiredCapabilities = []string{"merged"}
var providedCapabilities = []string{"normalized"}

var _ contracts.Normalizer[dnd.ItemRegistry] = (*Normalizer)(nil)
var _ contracts.ManifestMetadataProvider = (*Normalizer)(nil)
var _ pipeline.CheckpointFingerprintProvider = (*Normalizer)(nil)

type Options struct{}

type Normalizer struct {
	llm               contracts.StructuredLLMClient
	promptSHA         string
	responseSchemaSHA string
}

func New(llmClient contracts.StructuredLLMClient, _ Options) (*Normalizer, error) {
	if llmClient == nil {
		return nil, normalizerErrorf("LLM client must not be nil")
	}
	promptSHA, err := promptAssetMetadata()
	if err != nil {
		return nil, normalizerErrorf("load prompt metadata: %w", err)
	}
	responseSchema, err := entityreconcile.LoadResponseSchema()
	if err != nil {
		return nil, normalizerErrorf("load response schema: %w", err)
	}
	return &Normalizer{llm: llmClient, promptSHA: promptSHA, responseSchemaSHA: responseSchema.SHA256}, nil
}

func (n *Normalizer) Key() string                               { return Key }
func (n *Normalizer) ReferenceSlots() []contracts.ReferenceSlot { return nil }

func (n *Normalizer) ManifestMetadata() map[string]any {
	if n == nil {
		return nil
	}
	return map[string]any{
		"prompt_id": PromptID, "prompt_version": entityreconcile.SchemaVersion, "prompt_sha256": n.promptSHA,
		"response_schema_key": string(entityreconcile.ResponseSchemaKey), "response_schema_id": entityreconcile.ResponseSchemaID,
		"response_schema_name": entityreconcile.ResponseSchemaName, "response_schema_version": entityreconcile.SchemaVersion,
		"response_schema_sha256": n.responseSchemaSHA, "identity_policy": identity.Policy,
		"normalization_policy": normalizationPolicy, "semantic_context_policy": semanticContextPolicy, "semantic_context_radius": semanticContextRadius,
	}
}

func (n *Normalizer) CheckpointFingerprints() []pipeline.CheckpointFingerprint {
	if n == nil {
		return nil
	}
	return []pipeline.CheckpointFingerprint{
		{Name: "prompt", Value: n.promptSHA}, {Name: "response_schema", Value: n.responseSchemaSHA},
		{Name: "identity_policy", Value: identity.Policy}, {Name: "normalization_policy", Value: normalizationPolicy},
		{Name: "semantic_context_policy", Value: fmt.Sprintf("%s:%d", semanticContextPolicy, semanticContextRadius)},
	}
}

func (n *Normalizer) Normalize(ctx context.Context, req contracts.TypedNormalizeRequest[dnd.ItemRegistry]) (contracts.TypedNormalizeResult[dnd.ItemRegistry], error) {
	if n == nil {
		return contracts.TypedNormalizeResult[dnd.ItemRegistry]{}, normalizerErrorf("normalizer must not be nil")
	}
	if n.llm == nil {
		return contracts.TypedNormalizeResult[dnd.ItemRegistry]{}, normalizerErrorf("LLM client must not be nil")
	}
	if ctx == nil {
		return contracts.TypedNormalizeResult[dnd.ItemRegistry]{}, normalizerErrorf("context must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return contracts.TypedNormalizeResult[dnd.ItemRegistry]{}, normalizerErrorf("context error before normalize: %w", err)
	}

	order := shared.NewSourceRefOrder(req.Source)
	records, warnings := preprocessRecords(req.MergeOutput.Value, order)
	deterministic := recordList(records)
	materials, ready, err := entityreconcile.BuildContext(req.Source, reconciliationCandidates(records), semanticContextRadius)
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.ItemRegistry]{}, normalizerErrorf("build semantic context: %w", err)
	}
	if !ready {
		return contracts.TypedNormalizeResult[dnd.ItemRegistry]{Value: deterministic, Warnings: limitWarnings(warnings)}, nil
	}

	var response entityreconcile.ProposalResponse
	if _, err := n.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
		StageName: Key, PromptID: PromptID, PromptVersion: entityreconcile.SchemaVersion,
		ProfileID: req.LLMProfile, SessionID: req.SessionID,
		Inputs: contracts.LLMInputSet{"candidates": materials.Candidates, "transcript": materials.Transcript},
	}, &response); err != nil {
		if errors.Is(err, contracts.ErrInvalidStructuredOutput) {
			return n.invalidStructuredResult(deterministic, warnings), nil
		}
		return contracts.TypedNormalizeResult[dnd.ItemRegistry]{}, normalizerErrorf("complete structured output: %w", err)
	}
	assessment := materials.Assess(response)
	applied, semanticWarnings, rejectedGroups := applySafeGroups(records, reconciliationGroups(assessment, materials.CandidateKeys()), order)
	warnings = append(warnings, semanticWarnings...)
	if rejectedGroups > 0 {
		return currencyRetryResult(recordList(applied), warnings, rejectedGroups), nil
	}
	if assessment.DiscardedGroups() == 0 {
		return contracts.TypedNormalizeResult[dnd.ItemRegistry]{Value: recordList(applied), Warnings: limitWarnings(warnings)}, nil
	}
	return retryResult(recordList(applied), warnings, assessment), nil
}

func (n *Normalizer) invalidStructuredResult(value dnd.ItemRegistry, warnings []contracts.Warning) contracts.TypedNormalizeResult[dnd.ItemRegistry] {
	return contracts.TypedNormalizeResult[dnd.ItemRegistry]{Value: value, Warnings: limitWarningsForRetry(warnings), Retry: &contracts.NormalizeRetry{
		ReasonCode: ReasonCodeItemSemanticProposalInvalid, Message: "semantic proposal requires retry: invalid structured output",
		FallbackWarnings: []contracts.Warning{semanticFallbackWarning(-1)},
	}}
}

func retryResult(value dnd.ItemRegistry, warnings []contracts.Warning, assessment entityreconcile.Assessment) contracts.TypedNormalizeResult[dnd.ItemRegistry] {
	return contracts.TypedNormalizeResult[dnd.ItemRegistry]{Value: value, Warnings: limitWarningsForRetry(warnings), Retry: &contracts.NormalizeRetry{
		ReasonCode: ReasonCodeItemSemanticProposalInvalid, Message: diagnostics.Aggregate("semantic proposal requires retry", reconciliationIssues(assessment)),
		FallbackWarnings: []contracts.Warning{semanticFallbackWarning(assessment.DiscardedGroups())},
	}}
}

func currencyRetryResult(value dnd.ItemRegistry, warnings []contracts.Warning, rejectedGroups int) contracts.TypedNormalizeResult[dnd.ItemRegistry] {
	return contracts.TypedNormalizeResult[dnd.ItemRegistry]{Value: value, Warnings: limitWarningsForRetry(warnings), Retry: &contracts.NormalizeRetry{
		ReasonCode: ReasonCodeItemSemanticProposalInvalid, Message: "semantic proposal requires retry: currency may only be consolidated with aliases of one denomination",
		FallbackWarnings: []contracts.Warning{semanticFallbackWarning(rejectedGroups)},
	}}
}

func semanticFallbackWarning(discarded int) contracts.Warning {
	message := "semantic proposal could not be applied"
	if discarded >= 0 {
		message = fmt.Sprintf("%d proposal group(s) omitted after semantic proposal retry exhaustion", discarded)
	}
	return contracts.Warning{Scope: "items", ReasonCode: ReasonCodeItemSemanticReconciliationExhausted, Message: message}
}

func limitWarnings(warnings []contracts.Warning) []contracts.Warning {
	return diagnostics.LimitWarnings(warnings, "items", ReasonCodeItemNormalizationWarningsOmitted)
}

func limitWarningsForRetry(warnings []contracts.Warning) []contracts.Warning {
	if warnings == nil {
		return nil
	}
	if len(warnings) < diagnostics.MaxWarnings {
		return append([]contracts.Warning(nil), warnings...)
	}
	displayed := diagnostics.MaxWarnings - 2
	bounded := append([]contracts.Warning(nil), warnings[:displayed]...)
	return append(bounded, contracts.Warning{Scope: "items", ReasonCode: ReasonCodeItemNormalizationWarningsOmitted, Message: fmt.Sprintf("%d additional warning(s) omitted", len(warnings)-displayed)})
}

type normalizedRecord struct {
	item         dnd.Item
	inputIndexes []int
	earliest     int
}

func preprocessRecords(input dnd.ItemRegistry, order shared.SourceRefOrder) ([]normalizedRecord, []contracts.Warning) {
	if input.Items == nil {
		return nil, nil
	}
	records := make([]normalizedRecord, len(input.Items))
	warnings := make([]contracts.Warning, 0)
	for index, inputItem := range input.Items {
		item, fieldsChanged, refsChanged := normalizeRecord(inputItem, order)
		records[index] = normalizedRecord{item: item, inputIndexes: []int{index}, earliest: index}
		if fieldsChanged {
			warnings = append(warnings, contracts.Warning{Scope: itemScope(index), ReasonCode: ReasonCodeItemFieldsNormalized, Message: fmt.Sprintf("input index %d: item name normalized for %s", index, diagnostics.Quote(inputItem.Name))})
		}
		if refsChanged {
			warnings = append(warnings, contracts.Warning{Scope: itemScope(index), ReasonCode: ReasonCodeSourceReferencesNormalized, Message: fmt.Sprintf("input index %d: source references normalized (original count %d, final count %d)", index, len(inputItem.SourceRefs), len(item.SourceRefs))})
		}
		if inputItem.ID != item.ID {
			warnings = append(warnings, contracts.Warning{Scope: itemScope(index), ReasonCode: ReasonCodeItemIDRecomputed, Message: fmt.Sprintf("input index %d: item ID recomputed from %s", index, diagnostics.Quote(item.Name))})
		}
	}
	groups := comparisonNameGroups(records)
	output := make([]normalizedRecord, 0, len(groups))
	for _, members := range groups {
		retained := cloneRecord(records[members[0]])
		for _, member := range members[1:] {
			retained.item.SourceRefs = append(retained.item.SourceRefs, records[member].item.SourceRefs...)
			retained.inputIndexes = append(retained.inputIndexes, records[member].inputIndexes...)
		}
		retained.item.SourceRefs = order.Canonicalize(retained.item.SourceRefs)
		retained.item.ID = identity.DeriveID(retained.item.Name)
		retained.inputIndexes = sortedUniqueIndexes(retained.inputIndexes)
		output = append(output, retained)
		if len(members) > 1 {
			warnings = append(warnings, duplicateWarning(retained.earliest, memberInputIndexes(records, members[1:])))
		}
	}
	return output, warnings
}

func normalizeRecord(input dnd.Item, order shared.SourceRefOrder) (dnd.Item, bool, bool) {
	output := cloneItem(input)
	output.Name = identity.NormalizeDisplay(input.Name)
	output.SourceRefs = order.Canonicalize(input.SourceRefs)
	output.ID = identity.DeriveID(output.Name)
	return output, input.Name != output.Name, !reflect.DeepEqual(input.SourceRefs, output.SourceRefs)
}

func comparisonNameGroups(records []normalizedRecord) [][]int {
	groups := make([][]int, 0, len(records))
	for index, record := range records {
		key := identity.ComparisonKey(record.item.Name)
		found := false
		for groupIndex, members := range groups {
			first := records[members[0]]
			if identity.ComparisonKey(first.item.Name) == key {
				groups[groupIndex] = append(groups[groupIndex], index)
				found = true
				break
			}
		}
		if !found {
			groups = append(groups, []int{index})
		}
	}
	return groups
}

func cloneItem(input dnd.Item) dnd.Item {
	input.SourceRefs = cloneSourceRefs(input.SourceRefs)
	return input
}
func cloneRecord(input normalizedRecord) normalizedRecord {
	input.item = cloneItem(input.item)
	input.inputIndexes = append([]int(nil), input.inputIndexes...)
	return input
}
func cloneSourceRefs(input []source.SourceRef) []source.SourceRef {
	return append([]source.SourceRef(nil), input...)
}

func sortedUniqueIndexes(indexes []int) []int {
	if len(indexes) == 0 {
		return nil
	}
	out := append([]int(nil), indexes...)
	sort.Ints(out)
	write := 1
	for _, index := range out[1:] {
		if index != out[write-1] {
			out[write] = index
			write++
		}
	}
	return out[:write]
}
func memberInputIndexes(records []normalizedRecord, members []int) []int {
	indexes := make([]int, 0, len(members))
	for _, member := range members {
		indexes = append(indexes, records[member].inputIndexes...)
	}
	return sortedUniqueIndexes(indexes)
}
func recordValues(records []normalizedRecord) []dnd.Item {
	if records == nil {
		return nil
	}
	values := make([]dnd.Item, len(records))
	for index, record := range records {
		values[index] = cloneItem(record.item)
	}
	return values
}
func recordList(records []normalizedRecord) dnd.ItemRegistry {
	if records == nil {
		return dnd.ItemRegistry{}
	}
	return dnd.ItemRegistry{Items: recordValues(records)}
}

func duplicateWarning(retainedIndex int, removed []int) contracts.Warning {
	const maxDisplayedIndices = 20
	displayed := removed
	if len(displayed) > maxDisplayedIndices {
		displayed = displayed[:maxDisplayedIndices]
	}
	indices := make([]string, len(displayed))
	for index, removedIndex := range displayed {
		indices[index] = strconv.Itoa(removedIndex)
	}
	message := fmt.Sprintf("retained input index %d; removed input indices [%s]", retainedIndex, strings.Join(indices, ", "))
	if omitted := len(removed) - len(displayed); omitted > 0 {
		message += fmt.Sprintf("; %d additional removed input indices omitted", omitted)
	}
	return contracts.Warning{Scope: itemScope(retainedIndex), ReasonCode: ReasonCodeDuplicateItemCollapsed, Message: message}
}
func itemScope(index int) string { return fmt.Sprintf("items[%d]", index) }

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{Key: Key, Stage: pipeline.StageNormalize, ExecutionClass: contracts.ExecutionClassLLMBacked, Requires: append([]string(nil), requiredCapabilities...), Provides: append([]string(nil), providedCapabilities...), ArtifactKind: dnd.ItemRegistryKind}
}
func Register(registry *pipeline.NormalizerRegistry) error {
	return pipeline.RegisterNormalizerBuilder(registry, ModuleSpec(), validateOptions, func(request pipeline.BuildRequest) (contracts.Normalizer[dnd.ItemRegistry], error) {
		options, err := DecodeOptions(request.Options)
		if err != nil {
			return nil, err
		}
		return New(request.Dependencies.LLM, options)
	})
}
func validateOptions(options map[string]any) error { _, err := DecodeOptions(options); return err }
func DecodeOptions(options map[string]any) (Options, error) {
	if err := pipeline.RejectUnknownOptions(options); err != nil {
		return Options{}, normalizerErrorf("%w", err)
	}
	return Options{}, nil
}
func normalizerErrorf(format string, args ...any) error {
	return fmt.Errorf("dnd item registry normalizer: "+format, args...)
}