notarius/internal/modules/dnd/locations/registry/registry.go

// Package registry resolves normalized location artifacts into immutable
// ID-grounding data for D&D extraction modules.
package registry

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"mime"
	"strings"
	"sync"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	locationcodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locations"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/identity"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared/diagnostics"
)

const (
	ReferenceSlot = "locations"
	MaxBytes      = 1048576
	emptyPrompt   = `{"locations":[]}`
)

// Registry is an immutable, validated location registry prepared for prompt
// grounding. All accessors return defensive copies.
type Registry struct {
	bound            bool
	list             dnd.LocationList
	canonical        []byte
	digest           string
	projectionDigest string
	promptInput      contracts.LLMInputMaterial
	lookupByID       map[string]int
}

// Resolver retains the construction-time registry and memoizes immutable
// operation-time registries. It never retains caller-owned reference bytes.
type Resolver struct {
	seeded *Registry

	mu       sync.Mutex
	cache    map[string]*Registry
	rawCache map[string]*Registry
}

// NewResolver validates a materialized construction-time location reference.
// An empty slot is permitted because a generated reference is supplied only at
// operation time.
func NewResolver(references contracts.ReferenceSet) (*Resolver, error) {
	seeded, err := Resolve(constructionReferences(references))
	if err != nil {
		return nil, err
	}
	return &Resolver{seeded: seeded, cache: make(map[string]*Registry), rawCache: make(map[string]*Registry)}, nil
}

func constructionReferences(references contracts.ReferenceSet) contracts.ReferenceSet {
	slot, ok := references.Slots[ReferenceSlot]
	if !ok || len(slot.Items) > 0 {
		return references
	}
	cloned := contracts.ReferenceSet{Slots: make(map[string]contracts.ResolvedReferenceSlot, len(references.Slots))}
	for name, value := range references.Slots {
		cloned.Slots[name] = value
	}
	delete(cloned.Slots, ReferenceSlot)
	return cloned
}

// Seeded returns the validated construction-time registry.
func (r *Resolver) Seeded() *Registry {
	if r == nil {
		return nil
	}
	return r.seeded
}

// Resolve returns the generated operation-time registry when the locations
// slot is present, otherwise it returns the construction-time registry.
func (r *Resolver) Resolve(references contracts.ReferenceSet) (*Registry, error) {
	if r == nil {
		return Resolve(references)
	}
	if _, ok := references.Slots[ReferenceSlot]; !ok {
		return r.seeded, nil
	}

	slot := references.Slots[ReferenceSlot]
	rawKey := ""
	if len(slot.Items) == 1 {
		rawKey = strings.ToLower(strings.TrimSpace(slot.Items[0].MediaType)) + "\x00" + semanticDigest(slot.Items[0].Content)
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	if rawKey != "" {
		if cached, ok := r.rawCache[rawKey]; ok {
			return cached, nil
		}
	}
	resolved, err := Resolve(references)
	if err != nil {
		return nil, err
	}
	if sameRegistryIdentity(r.seeded, resolved) {
		if rawKey != "" {
			r.rawCache[rawKey] = r.seeded
		}
		return r.seeded, nil
	}
	if cached, ok := r.cache[resolved.Digest()]; ok {
		if rawKey != "" {
			r.rawCache[rawKey] = cached
		}
		return cached, nil
	}
	r.cache[resolved.Digest()] = resolved
	if rawKey != "" {
		r.rawCache[rawKey] = resolved
	}
	return resolved, nil
}

func sameRegistryIdentity(first, second *Registry) bool {
	if first == nil || second == nil {
		return first == second
	}
	return first.bound == second.bound && first.digest == second.digest
}

// Resolve validates an optional location-list reference. An absent reference
// uses the canonical empty projection and has no durable registry identity.
func Resolve(references contracts.ReferenceSet) (*Registry, error) {
	slot, ok := references.Slots[ReferenceSlot]
	if !ok {
		content := []byte(emptyPrompt)
		projectionDigest := semanticDigest(content)
		return &Registry{
			list:             dnd.LocationList{Locations: []dnd.Location{}},
			canonical:        append([]byte(nil), content...),
			projectionDigest: projectionDigest,
			promptInput:      contracts.NewLLMInputMaterial(ReferenceSlot, locationcodec.MediaType, content, projectionDigest, ""),
			lookupByID:       map[string]int{},
		}, nil
	}
	if len(slot.Items) != 1 {
		return nil, fmt.Errorf("reference slot %q must contain exactly one item", ReferenceSlot)
	}

	item := slot.Items[0]
	mediaType, _, err := mime.ParseMediaType(item.MediaType)
	if err != nil {
		return nil, fmt.Errorf("reference slot %q item media type is invalid", ReferenceSlot)
	}
	if !strings.EqualFold(mediaType, locationcodec.MediaType) {
		return nil, fmt.Errorf("reference slot %q item media type must be %s", ReferenceSlot, locationcodec.MediaType)
	}
	if len(item.Content) > MaxBytes {
		return nil, fmt.Errorf("reference slot %q item is %d bytes, limit %d", ReferenceSlot, len(item.Content), MaxBytes)
	}

	codec := locationcodec.New()
	value, err := codec.Decode(item.Content)
	if err != nil {
		return nil, fmt.Errorf("decode location registry: invalid approved location JSON")
	}
	if issues := identity.ValidateList(value); len(issues) > 0 {
		return nil, fmt.Errorf("%s", formatIdentityIssues(issues))
	}
	content, err := codec.Encode(value)
	if err != nil {
		return nil, fmt.Errorf("encode canonical location registry: approved location value could not be encoded")
	}

	list := cloneLocationList(value)
	lookupByID := make(map[string]int, len(list.Locations))
	for index, location := range list.Locations {
		lookupByID[location.ID] = index
	}
	projection, err := promptProjection(list)
	if err != nil {
		return nil, fmt.Errorf("encode location prompt projection: %w", err)
	}
	digest := semanticDigest(content)
	projectionDigest := semanticDigest(projection)
	return &Registry{
		bound:            true,
		list:             list,
		canonical:        append([]byte(nil), content...),
		digest:           digest,
		projectionDigest: projectionDigest,
		promptInput:      contracts.NewLLMInputMaterial(ReferenceSlot, locationcodec.MediaType, projection, projectionDigest, ""),
		lookupByID:       lookupByID,
	}, nil
}

// Bound reports whether a location reference was supplied and validated.
func (r *Registry) Bound() bool { return r != nil && r.bound }

// Locations returns a defensive copy of the validated location records.
func (r *Registry) Locations() []dnd.Location {
	if r == nil {
		return nil
	}
	return cloneLocations(r.list.Locations)
}

// List returns a defensive copy of the validated location list.
func (r *Registry) List() dnd.LocationList {
	if r == nil {
		return dnd.LocationList{}
	}
	return cloneLocationList(r.list)
}

// CanonicalBytes returns a defensive copy of the canonical durable JSON.
func (r *Registry) CanonicalBytes() []byte {
	if r == nil {
		return nil
	}
	return append([]byte(nil), r.canonical...)
}

// Digest returns the semantic digest of the canonical JSON, or an empty string
// when the registry is unbound.
func (r *Registry) Digest() string {
	if r == nil {
		return ""
	}
	return r.digest
}

// ProjectionDigest returns the digest of the exact source-free prompt
// projection, including for an unbound or empty registry.
func (r *Registry) ProjectionDigest() string {
	if r == nil {
		return ""
	}
	return r.projectionDigest
}

// Count returns the number of validated location records.
func (r *Registry) Count() int {
	if r == nil {
		return 0
	}
	return len(r.list.Locations)
}

// PromptInput returns the ordered ID-and-name projection without evidence or
// reference provenance.
func (r *Registry) PromptInput() contracts.LLMInputMaterial {
	if r == nil {
		return contracts.LLMInputMaterial{}
	}
	return r.promptInput.Clone()
}

// Lookup returns the canonical location for an exact durable location ID.
func (r *Registry) Lookup(id string) (dnd.Location, bool) {
	if r == nil {
		return dnd.Location{}, false
	}
	index, ok := r.lookupByID[id]
	if !ok {
		return dnd.Location{}, false
	}
	return cloneLocation(r.list.Locations[index]), true
}

// Matches reports whether id resolves to exactly the supplied canonical name.
func (r *Registry) Matches(id, name string) bool {
	location, ok := r.Lookup(id)
	return ok && location.Name == name
}

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

type projectedLocation struct {
	ID   string `json:"id"`
	Name string `json:"name"`
}

type projectedLocationList struct {
	Locations []projectedLocation `json:"locations"`
}

func promptProjection(list dnd.LocationList) ([]byte, error) {
	projection := projectedLocationList{Locations: make([]projectedLocation, len(list.Locations))}
	for index, location := range list.Locations {
		projection.Locations[index] = projectedLocation{ID: location.ID, Name: location.Name}
	}
	return json.Marshal(projection)
}

func formatIdentityIssues(issues []identity.Issue) string {
	parts := make([]string, len(issues))
	for index, issue := range issues {
		parts[index] = fmt.Sprintf("%s at record %d", issue.Code, issue.RecordIndex)
	}
	return diagnostics.Aggregate("validate location registry identity", parts)
}

func cloneLocationList(value dnd.LocationList) dnd.LocationList {
	return dnd.LocationList{Locations: cloneLocations(value.Locations)}
}

func cloneLocations(values []dnd.Location) []dnd.Location {
	if values == nil {
		return nil
	}
	cloned := make([]dnd.Location, len(values))
	for index, value := range values {
		cloned[index] = cloneLocation(value)
	}
	return cloned
}

func cloneLocation(value dnd.Location) dnd.Location {
	value.SourceRefs = append([]source.SourceRef(nil), value.SourceRefs...)
	return value
}