notarius/internal/framework/chunkmap/codec.go

package chunkmap

import (
	"bytes"
	"embed"
	"encoding/json"
	"fmt"
	"io"
	"regexp"
	"strings"

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

//go:embed assets/schemas/source_chunk_map.v1.json
var schemaAssets embed.FS

var digestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)

// Codec owns strict serialization for the durable chunk-map contract.
type Codec struct{}

func New() *Codec { return &Codec{} }

func (c *Codec) Kind() contracts.ArtifactKind { return ArtifactKind }

func (c *Codec) Schema() contracts.ArtifactSchema {
	raw, err := c.schemaBytes()
	if err != nil {
		return contracts.ArtifactSchema{}
	}
	return contracts.ArtifactSchema{
		ID:         SchemaID,
		Name:       SchemaName,
		Version:    SchemaVersion,
		JSONSchema: raw,
	}
}

func (c *Codec) MediaType() string { return MediaType }

// Build proves that a durable value describes the exact accepted source plan
// and materialized chunk list supplied by the framework.
func Build(request BuildRequest) (ChunkMap, error) {
	if err := source.ValidateDocument(request.Source); err != nil {
		return ChunkMap{}, fmt.Errorf("validate source document: %w", err)
	}
	sourceDigest, err := source.DigestDocument(request.Source)
	if err != nil {
		return ChunkMap{}, fmt.Errorf("digest source document: %w", err)
	}
	if sourceDigest != request.Source.Digest {
		return ChunkMap{}, fmt.Errorf("source digest %q does not match source document digest %q", sourceDigest, request.Source.Digest)
	}
	if sourceDigest != request.Plan.SourceDigest {
		return ChunkMap{}, fmt.Errorf("source digest %q does not match chunk plan source digest %q", sourceDigest, request.Plan.SourceDigest)
	}
	plan, err := source.CanonicalizeChunkPlan(request.Plan)
	if err != nil {
		return ChunkMap{}, fmt.Errorf("canonicalize chunk plan: %w", err)
	}
	if err := source.ValidateChunkPlan(request.Source, plan); err != nil {
		return ChunkMap{}, fmt.Errorf("validate accepted chunk plan: %w", err)
	}
	planDigest, err := source.DigestChunkPlan(plan)
	if err != nil {
		return ChunkMap{}, fmt.Errorf("digest accepted chunk plan: %w", err)
	}
	expected, err := source.MaterializeChunkPlan(request.Source, plan)
	if err != nil {
		return ChunkMap{}, fmt.Errorf("materialize accepted chunk plan: %w", err)
	}
	if err := verifyMaterializedChunks(request.Chunks, expected); err != nil {
		return ChunkMap{}, err
	}

	value := ChunkMap{
		SourceID:         request.Source.ID,
		SourceDigest:     sourceDigest,
		PlanDigest:       planDigest,
		RequestedChunker: request.RequestedChunker,
		Producer:         request.Producer,
		PlanAnnotations:  source.CloneChunkAnnotations(plan.Annotations),
		Chunks:           make([]Chunk, len(expected)),
	}
	for index, chunk := range expected {
		value.Chunks[index] = Chunk{
			ID:          chunk.ID,
			Index:       chunk.Index,
			SourceRef:   chunk.Ref,
			UnitCount:   len(chunk.Units),
			Annotations: source.CloneChunkAnnotations(chunk.Annotations),
		}
	}
	canonical, err := canonicalize(value)
	if err != nil {
		return ChunkMap{}, fmt.Errorf("validate chunk map: %w", err)
	}
	return clone(canonical), nil
}

// Serialize builds and encodes the framework-owned serialized artifact.
func Serialize(request BuildRequest) (contracts.SerializedArtifact, error) {
	value, err := Build(request)
	if err != nil {
		return contracts.SerializedArtifact{}, err
	}
	codec := New()
	content, err := codec.Encode(value)
	if err != nil {
		return contracts.SerializedArtifact{}, err
	}
	return contracts.SerializedArtifact{
		Kind:      ArtifactKind,
		Schema:    codec.Schema(),
		MediaType: MediaType,
		Content:   content,
	}, nil
}

func (c *Codec) Encode(value ChunkMap) ([]byte, error) {
	if _, err := c.schemaBytes(); err != nil {
		return nil, err
	}
	canonical, err := canonicalize(value)
	if err != nil {
		return nil, fmt.Errorf("encode source chunk map: %w", err)
	}
	content, err := json.Marshal(canonical)
	if err != nil {
		return nil, fmt.Errorf("encode source chunk map: %w", err)
	}
	return content, nil
}

func (c *Codec) Decode(content []byte) (ChunkMap, error) {
	if _, err := c.schemaBytes(); err != nil {
		return ChunkMap{}, err
	}
	decoder := json.NewDecoder(bytes.NewReader(content))
	decoder.DisallowUnknownFields()
	var value ChunkMap
	if err := decoder.Decode(&value); err != nil {
		return ChunkMap{}, fmt.Errorf("decode source chunk map: %w", err)
	}
	var trailing any
	if err := decoder.Decode(&trailing); err != io.EOF {
		return ChunkMap{}, fmt.Errorf("decode source chunk map: multiple JSON values")
	}
	canonical, err := canonicalize(value)
	if err != nil {
		return ChunkMap{}, fmt.Errorf("decode source chunk map: %w", err)
	}
	return clone(canonical), nil
}

func (c *Codec) schemaBytes() ([]byte, error) {
	raw, err := schemaAssets.ReadFile("assets/schemas/source_chunk_map.v1.json")
	if err != nil {
		return nil, fmt.Errorf("read source chunk map schema: %w", err)
	}
	var schema struct {
		ID       string   `json:"$id"`
		Title    string   `json:"title"`
		Type     string   `json:"type"`
		Required []string `json:"required"`
	}
	if err := json.Unmarshal(raw, &schema); err != nil {
		return nil, fmt.Errorf("decode source chunk map schema: %w", err)
	}
	if schema.ID != SchemaID || schema.Title != SchemaName || schema.Type != "object" || !hasRequiredFields(schema.Required) {
		return nil, fmt.Errorf("source chunk map schema identity or required fields are invalid")
	}
	return append([]byte(nil), raw...), nil
}

func hasRequiredFields(required []string) bool {
	want := map[string]bool{
		"source_id": true, "source_digest": true, "plan_digest": true,
		"requested_chunker": true, "producer": true, "plan_annotations": true,
		"chunks": true,
	}
	for _, field := range required {
		delete(want, field)
	}
	return len(want) == 0
}

func canonicalize(value ChunkMap) (ChunkMap, error) {
	if err := requireIdentity("source_id", value.SourceID); err != nil {
		return ChunkMap{}, err
	}
	if err := requireDigest("source_digest", value.SourceDigest); err != nil {
		return ChunkMap{}, err
	}
	if err := requireDigest("plan_digest", value.PlanDigest); err != nil {
		return ChunkMap{}, err
	}
	if err := requireIdentity("requested_chunker", value.RequestedChunker); err != nil {
		return ChunkMap{}, err
	}
	if err := requireIdentity("producer.input_module", value.Producer.InputModule); err != nil {
		return ChunkMap{}, err
	}
	if err := requireIdentity("producer.chunk_module", value.Producer.ChunkModule); err != nil {
		return ChunkMap{}, err
	}
	if value.Producer.LLMProfile != "" {
		if err := requireIdentity("producer.llm_profile", value.Producer.LLMProfile); err != nil {
			return ChunkMap{}, err
		}
	}
	annotations, err := canonicalizeAnnotations("plan_annotations", value.PlanAnnotations)
	if err != nil {
		return ChunkMap{}, err
	}
	value.PlanAnnotations = annotations
	if len(value.Chunks) == 0 {
		return ChunkMap{}, fmt.Errorf("chunks must not be empty")
	}
	seenIDs := make(map[string]struct{}, len(value.Chunks))
	plan := source.ChunkPlan{SourceDigest: value.SourceDigest, Annotations: annotations, Ranges: make([]source.ChunkRange, len(value.Chunks))}
	for index := range value.Chunks {
		chunk := &value.Chunks[index]
		if err := requireIdentity(fmt.Sprintf("chunks[%d].id", index), chunk.ID); err != nil {
			return ChunkMap{}, err
		}
		if _, exists := seenIDs[chunk.ID]; exists {
			return ChunkMap{}, fmt.Errorf("chunks[%d].id %q is duplicated", index, chunk.ID)
		}
		seenIDs[chunk.ID] = struct{}{}
		if chunk.Index != index {
			return ChunkMap{}, fmt.Errorf("chunks[%d].index = %d, want %d", index, chunk.Index, index)
		}
		if chunk.SourceRef.SourceID != value.SourceID {
			return ChunkMap{}, fmt.Errorf("chunks[%d].source_ref.source_id %q does not match source_id %q", index, chunk.SourceRef.SourceID, value.SourceID)
		}
		if chunk.SourceRef.StartUnitID <= 0 || chunk.SourceRef.EndUnitID <= 0 {
			return ChunkMap{}, fmt.Errorf("chunks[%d].source_ref endpoints must be positive", index)
		}
		if chunk.UnitCount <= 0 {
			return ChunkMap{}, fmt.Errorf("chunks[%d].unit_count must be positive", index)
		}
		chunkAnnotations, err := canonicalizeAnnotations(fmt.Sprintf("chunks[%d].annotations", index), chunk.Annotations)
		if err != nil {
			return ChunkMap{}, err
		}
		chunk.Annotations = chunkAnnotations
		plan.Ranges[index] = source.ChunkRange{
			StartUnitID: chunk.SourceRef.StartUnitID,
			EndUnitID:   chunk.SourceRef.EndUnitID,
			Annotations: chunkAnnotations,
		}
	}
	planDigest, err := source.DigestChunkPlan(plan)
	if err != nil {
		return ChunkMap{}, fmt.Errorf("reconstruct plan digest: %w", err)
	}
	if planDigest != value.PlanDigest {
		return ChunkMap{}, fmt.Errorf("plan_digest %q does not match reconstructed plan digest %q", value.PlanDigest, planDigest)
	}
	return value, nil
}

func canonicalizeAnnotations(name string, annotations source.ChunkAnnotations) (source.ChunkAnnotations, error) {
	for namespace := range annotations {
		if strings.TrimSpace(namespace) == "" || namespace != strings.TrimSpace(namespace) {
			return nil, fmt.Errorf("%s namespace %q must be non-empty and trimmed", name, namespace)
		}
	}
	canonical, err := source.CanonicalizeChunkAnnotations(annotations)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", name, err)
	}
	if canonical == nil {
		canonical = source.ChunkAnnotations{}
	}
	return canonical, nil
}

func requireIdentity(name, value string) error {
	if strings.TrimSpace(value) == "" || value != strings.TrimSpace(value) {
		return fmt.Errorf("%s must be non-empty and trimmed", name)
	}
	return nil
}

func requireDigest(name, value string) error {
	if !digestPattern.MatchString(value) {
		return fmt.Errorf("%s must be a canonical sha256 digest", name)
	}
	return nil
}

func verifyMaterializedChunks(actual, expected []source.Chunk) error {
	if len(actual) != len(expected) {
		return fmt.Errorf("materialized chunks length = %d, want %d", len(actual), len(expected))
	}
	for index := range expected {
		got, want := actual[index], expected[index]
		if got.ID != want.ID || got.SourceID != want.SourceID || got.Index != want.Index || got.Ref != want.Ref {
			return fmt.Errorf("materialized chunk[%d] identity or source range differs from accepted plan", index)
		}
		if len(got.Units) != len(want.Units) || !sameUnits(got.Units, want.Units) {
			return fmt.Errorf("materialized chunk[%d] units differ from accepted source range", index)
		}
		if !sameAnnotations(got.PlanAnnotations, want.PlanAnnotations) || !sameAnnotations(got.Annotations, want.Annotations) {
			return fmt.Errorf("materialized chunk[%d] annotations differ from accepted plan", index)
		}
	}
	return nil
}

func sameUnits(left, right []source.SourceUnit) bool {
	leftJSON, leftErr := json.Marshal(left)
	rightJSON, rightErr := json.Marshal(right)
	return leftErr == nil && rightErr == nil && bytes.Equal(leftJSON, rightJSON)
}

func sameAnnotations(left, right source.ChunkAnnotations) bool {
	leftCanonical, leftErr := source.CanonicalizeChunkAnnotations(left)
	rightCanonical, rightErr := source.CanonicalizeChunkAnnotations(right)
	if leftErr != nil || rightErr != nil {
		return false
	}
	leftJSON, leftErr := json.Marshal(leftCanonical)
	rightJSON, rightErr := json.Marshal(rightCanonical)
	return leftErr == nil && rightErr == nil && bytes.Equal(leftJSON, rightJSON)
}

func clone(value ChunkMap) ChunkMap {
	value.PlanAnnotations = cloneAnnotations(value.PlanAnnotations)
	value.Chunks = append([]Chunk(nil), value.Chunks...)
	for index := range value.Chunks {
		value.Chunks[index].Annotations = cloneAnnotations(value.Chunks[index].Annotations)
	}
	return value
}

func cloneAnnotations(annotations source.ChunkAnnotations) source.ChunkAnnotations {
	cloned := source.CloneChunkAnnotations(annotations)
	if cloned == nil {
		return source.ChunkAnnotations{}
	}
	return cloned
}