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

// Package npcregistry normalizes merged D&D non-player character records.
package npcregistry

import (
	"context"
	"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/framework/semanticreconcile"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/npcs/identity"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared/diagnostics"
)

const (
	Key                 = "dnd/npc-registry"
	PromptID            = "dnd.npc_registry.normalize"
	PromptVersion       = "v1"
	normalizationPolicy = "dnd.npc_registry.normalize.v6"
	NormalizationPolicy = normalizationPolicy

	ReasonCodeNPCFieldsNormalized                = "npc_fields_normalized"
	ReasonCodeNPCIDRecomputed                    = "npc_id_recomputed"
	ReasonCodeSourceReferencesNormalized         = "source_references_normalized"
	ReasonCodeDuplicateNPCCollapsed              = "duplicate_npc_collapsed"
	ReasonCodeNPCSemanticProposalInvalid         = "npc_semantic_proposal_invalid"
	ReasonCodeNPCSemanticReconciliationExhausted = "npc_semantic_reconciliation_exhausted"
)

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

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

type Options struct{}

type Normalizer struct {
	engine *semanticreconcile.Engine
}

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)
	}
	engine, err := semanticreconcile.NewEngine(llmClient, semanticreconcile.PromptSpec{
		ID: PromptID, Version: PromptVersion, SHA256: promptSHA,
	}, semanticreconcile.DefaultLimits())
	if err != nil {
		return nil, normalizerErrorf("construct semantic reconciliation engine: %w", err)
	}
	return &Normalizer{engine: engine}, 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 || n.engine == nil {
		return nil
	}
	metadata := n.engine.ManifestMetadata()
	metadata["identity_policy"] = identity.Policy
	metadata["normalization_policy"] = normalizationPolicy
	return metadata
}

func (n *Normalizer) CheckpointFingerprints() []pipeline.CheckpointFingerprint {
	if n == nil || n.engine == nil {
		return nil
	}
	fingerprints := n.engine.CheckpointFingerprints()
	return append(fingerprints,
		pipeline.CheckpointFingerprint{Name: "identity_policy", Value: identity.Policy},
		pipeline.CheckpointFingerprint{Name: "normalization_policy", Value: normalizationPolicy},
	)
}

func (n *Normalizer) Normalize(ctx context.Context, req contracts.TypedNormalizeRequest[dnd.NPCRegistry]) (contracts.TypedNormalizeResult[dnd.NPCRegistry], error) {
	if n == nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("normalizer must not be nil")
	}
	if n.engine == nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("semantic reconciliation engine must not be nil")
	}
	if ctx == nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("context must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("context error before normalize: %w", err)
	}
	if req.Source == nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("source document must not be nil")
	}

	order := shared.NewSourceRefOrder(req.Source)
	records, findings := preprocessRecords(req.MergeOutput.Value, order)
	deterministic := recordList(records)
	if len(records) < 2 {
		return normalizationResult(deterministic, findings, nil)
	}

	candidates, envelopes, err := reconciliationInputs(records)
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("prepare semantic reconciliation inputs: %w", err)
	}
	reconciliation, err := n.engine.Reconcile(ctx, semanticreconcile.Request{
		StageName: Key, Source: req.Source, Candidates: candidates,
		ProfileID: req.LLMProfile, SessionID: req.SessionID, StructuredOutputRepairAttempts: req.StructuredOutputRepairAttempts, Correction: req.Correction,
	})
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("reconcile semantic duplicates: %w", err)
	}

	switch reconciliation.Disposition() {
	case semanticreconcile.SkippedInsufficientCandidates:
		return normalizationResult(deterministic, findings, nil)
	case semanticreconcile.SkippedLimitExceeded:
		return fallbackResult(deterministic, findings, semanticFallbackFinding(-1, "the candidate set exceeded the configured semantic-reconciliation limits"))
	case semanticreconcile.RetryableInvalidStructuredOutput:
		return n.invalidStructuredResult(deterministic, findings)
	case semanticreconcile.Complete, semanticreconcile.RetryableDiscardedProposalGroups:
	default:
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("unknown semantic reconciliation disposition %d", reconciliation.Disposition())
	}

	applied, semanticFindings, err := applyReconciliationPlan(reconciliation.Plan(), records, envelopes, order)
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("apply semantic reconciliation plan: %w", err)
	}
	findings = append(findings, semanticFindings...)
	if reconciliation.Disposition() == semanticreconcile.Complete {
		return normalizationResult(recordList(applied), findings, reconciliation.ModelCandidate())
	}
	return retryResult(recordList(applied), findings, reconciliation)
}

func (n *Normalizer) invalidStructuredResult(value dnd.NPCRegistry, findings []diagnostics.Finding) (contracts.TypedNormalizeResult[dnd.NPCRegistry], error) {
	return retryResultWithFallback(value, findings, nil, "", "semantic proposal requires retry: invalid structured output", semanticFallbackFinding(-1, "the final model response did not satisfy the structured-output contract"))
}

func retryResult(value dnd.NPCRegistry, findings []diagnostics.Finding, reconciliation semanticreconcile.Result) (contracts.TypedNormalizeResult[dnd.NPCRegistry], error) {
	issues := reconciliation.Issues()
	correctionDetails, err := reconciliation.CorrectionDetails()
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("build semantic correction details: %w", err)
	}
	correctionGuidance, err := reconciliation.CorrectionGuidance()
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("build semantic correction guidance: %w", err)
	}
	return retryResultWithFallback(value, findings, reconciliation.ModelCandidate(), correctionGuidance, diagnostics.Aggregate("semantic proposal requires retry", semanticreconcile.IssueDetails(issues)), semanticFallbackFinding(reconciliation.DiscardedGroupCount(), correctionDetails...))
}

func semanticFallbackFinding(discardedGroups int, details ...string) diagnostics.Finding {
	message := "semantic reconciliation could not be completed; all affected candidates were preserved separately"
	if discardedGroups >= 0 {
		message = fmt.Sprintf("%d proposal group(s) omitted from the final semantic proposal; all affected candidates were preserved separately", discardedGroups)
	}
	if len(details) > 0 {
		message = diagnostics.Aggregate(message, details)
	}
	return diagnostics.Finding{Scope: "npcs", ReasonCode: ReasonCodeNPCSemanticReconciliationExhausted, Message: message}
}

func normalizationResult(value dnd.NPCRegistry, findings []diagnostics.Finding, candidate *contracts.ModelCandidate) (contracts.TypedNormalizeResult[dnd.NPCRegistry], error) {
	diagnosticGroups, err := diagnostics.Collect(findings, contracts.DiagnosticDispositionObservation, contracts.DiagnosticCategoryNormalization)
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("collect normalization diagnostics: %w", err)
	}
	return contracts.TypedNormalizeResult[dnd.NPCRegistry]{Value: value, Diagnostics: diagnosticGroups, ModelCandidate: candidate}, nil
}

func fallbackResult(value dnd.NPCRegistry, findings []diagnostics.Finding, fallback diagnostics.Finding) (contracts.TypedNormalizeResult[dnd.NPCRegistry], error) {
	result, err := normalizationResult(value, findings, nil)
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, err
	}
	fallbackGroups, err := diagnostics.Collect([]diagnostics.Finding{fallback}, contracts.DiagnosticDispositionWarning, contracts.DiagnosticCategoryFallback)
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("collect fallback diagnostic: %w", err)
	}
	result.Diagnostics = append(result.Diagnostics, fallbackGroups...)
	return result, nil
}

func retryResultWithFallback(value dnd.NPCRegistry, findings []diagnostics.Finding, candidate *contracts.ModelCandidate, correctionGuidance, message string, fallback diagnostics.Finding) (contracts.TypedNormalizeResult[dnd.NPCRegistry], error) {
	result, err := normalizationResult(value, findings, candidate)
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, err
	}
	fallbackGroups, err := diagnostics.Collect([]diagnostics.Finding{fallback}, contracts.DiagnosticDispositionWarning, contracts.DiagnosticCategoryFallback)
	if err != nil {
		return contracts.TypedNormalizeResult[dnd.NPCRegistry]{}, normalizerErrorf("collect fallback diagnostic: %w", err)
	}
	result.Retry = &contracts.NormalizeRetry{ReasonCode: ReasonCodeNPCSemanticProposalInvalid, Message: message, CorrectionGuidance: correctionGuidance, FallbackDiagnostics: fallbackGroups}
	return result, nil
}

type normalizedRecord struct {
	npc          dnd.NPC
	inputIndexes []int
	earliest     int
}

func preprocessRecords(input dnd.NPCRegistry, order shared.SourceRefOrder) ([]normalizedRecord, []diagnostics.Finding) {
	if input.NPCs == nil {
		return nil, nil
	}
	records := make([]normalizedRecord, len(input.NPCs))
	findings := make([]diagnostics.Finding, 0)
	for index, inputNPC := range input.NPCs {
		npc, fieldsChanged, referencesChanged := normalizeRecord(inputNPC, order)
		records[index] = normalizedRecord{npc: npc, inputIndexes: []int{index}, earliest: index}
		if fieldsChanged {
			findings = append(findings, diagnostics.Finding{Scope: npcScope(index), ReasonCode: ReasonCodeNPCFieldsNormalized, Message: fmt.Sprintf("input index %d: NPC name normalized for %s", index, diagnostics.Quote(inputNPC.Name))})
		}
		if referencesChanged {
			findings = append(findings, diagnostics.Finding{Scope: npcScope(index), ReasonCode: ReasonCodeSourceReferencesNormalized, Message: fmt.Sprintf("input index %d: source references normalized (original count %d, final count %d)", index, len(inputNPC.SourceRefs), len(npc.SourceRefs))})
		}
		if inputNPC.ID != npc.ID {
			findings = append(findings, diagnostics.Finding{Scope: npcScope(index), ReasonCode: ReasonCodeNPCIDRecomputed, Message: fmt.Sprintf("input index %d: NPC ID recomputed from %s", index, diagnostics.Quote(npc.Name))})
		}
	}

	groups := canonicalNameGroups(records)
	output := make([]normalizedRecord, 0, len(groups))
	for _, members := range groups {
		consolidated, referencesChanged := consolidate(records, members, order)
		output = append(output, consolidated)
		retainedIndex := consolidated.earliest
		if referencesChanged {
			findings = append(findings, diagnostics.Finding{Scope: npcScope(retainedIndex), ReasonCode: ReasonCodeSourceReferencesNormalized, Message: fmt.Sprintf("input index %d: source references normalized during identity consolidation (final count %d)", retainedIndex, len(consolidated.npc.SourceRefs))})
		}
		if len(members) > 1 {
			findings = append(findings, duplicateFinding(retainedIndex, memberInputIndexes(records, members[1:])))
		}
	}
	return output, findings
}

func normalizeRecord(input dnd.NPC, order shared.SourceRefOrder) (dnd.NPC, bool, bool) {
	output := cloneNPC(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 cloneNPC(input dnd.NPC) dnd.NPC {
	input.SourceRefs = cloneSourceRefs(input.SourceRefs)
	return input
}

func canonicalNameGroups(records []normalizedRecord) [][]int {
	groups := make([][]int, 0, len(records))
	ownerByKey := make(map[string]int, len(records))
	for index, record := range records {
		key := identity.ComparisonKey(record.npc.Name)
		if key != "" {
			if groupIndex, ok := ownerByKey[key]; ok {
				groups[groupIndex] = append(groups[groupIndex], index)
				continue
			}
			ownerByKey[key] = len(groups)
		}
		groups = append(groups, []int{index})
	}
	return groups
}

func consolidate(records []normalizedRecord, members []int, order shared.SourceRefOrder) (normalizedRecord, bool) {
	output := cloneRecord(records[members[0]])
	originalRefs := cloneSourceRefs(output.npc.SourceRefs)
	for _, member := range members[1:] {
		output.npc.SourceRefs = append(output.npc.SourceRefs, records[member].npc.SourceRefs...)
		output.inputIndexes = append(output.inputIndexes, records[member].inputIndexes...)
		if records[member].earliest < output.earliest {
			output.earliest = records[member].earliest
		}
	}
	output.inputIndexes = sortedUniqueIndexes(output.inputIndexes)
	output.npc.SourceRefs = order.Canonicalize(output.npc.SourceRefs)
	output.npc.ID = identity.DeriveID(output.npc.Name)
	return output, !reflect.DeepEqual(originalRefs, output.npc.SourceRefs)
}

func cloneRecord(input normalizedRecord) normalizedRecord {
	input.npc = cloneNPC(input.npc)
	input.inputIndexes = append([]int(nil), input.inputIndexes...)
	return input
}

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 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 recordValues(records []normalizedRecord) []dnd.NPC {
	if records == nil {
		return nil
	}
	values := make([]dnd.NPC, len(records))
	for index, record := range records {
		values[index] = cloneNPC(record.npc)
	}
	return values
}

func recordList(records []normalizedRecord) dnd.NPCRegistry {
	if records == nil {
		return dnd.NPCRegistry{}
	}
	return dnd.NPCRegistry{NPCs: recordValues(records)}
}

func cloneSourceRefs(input []source.SourceRef) []source.SourceRef {
	if input == nil {
		return nil
	}
	return append([]source.SourceRef(nil), input...)
}

func duplicateFinding(retainedIndex int, removed []int) diagnostics.Finding {
	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 diagnostics.Finding{Scope: npcScope(retainedIndex), ReasonCode: ReasonCodeDuplicateNPCCollapsed, Message: message}
}

func npcScope(index int) string { return fmt.Sprintf("npcs[%d]", index) }

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{Key: Key, Stage: pipeline.StageNormalize, ExecutionClass: contracts.ExecutionClassLLMBacked, CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1, Requires: append([]string(nil), requiredCapabilities...), Provides: append([]string(nil), providedCapabilities...), ArtifactKind: dnd.NPCRegistryKind}
}

func Register(registry *pipeline.NormalizerRegistry) error {
	return pipeline.RegisterNormalizerBuilder(registry, ModuleSpec(), validateOptions, func(request pipeline.BuildRequest) (contracts.Normalizer[dnd.NPCRegistry], 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 npc registry normalizer: "+format, args...)
}