notarius/internal/core/source/chunk_plan.go

package source

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"strings"
)

// CanonicalizeChunkAnnotations validates annotation namespaces and JSON values
// and returns an independently owned map whose values use canonical JSON bytes.
func CanonicalizeChunkAnnotations(annotations ChunkAnnotations) (ChunkAnnotations, error) {
	if len(annotations) == 0 {
		return nil, nil
	}
	canonical := make(ChunkAnnotations, len(annotations))
	for namespace, raw := range annotations {
		if strings.TrimSpace(namespace) == "" {
			return nil, fmt.Errorf("chunk annotation namespace must not be empty")
		}
		if strings.TrimSpace(namespace) != namespace {
			return nil, fmt.Errorf("chunk annotation namespace %q must not contain leading or trailing whitespace", namespace)
		}
		value, err := decodeAnnotation(raw)
		if err != nil {
			return nil, fmt.Errorf("chunk annotation %q: %w", namespace, err)
		}
		encoded, err := json.Marshal(value)
		if err != nil {
			return nil, fmt.Errorf("chunk annotation %q contains an unsupported value: %w", namespace, err)
		}
		canonical[namespace] = encoded
	}
	return canonical, nil
}

func decodeAnnotation(raw json.RawMessage) (any, error) {
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.UseNumber()
	var value any
	if err := decoder.Decode(&value); err != nil {
		return nil, fmt.Errorf("must contain valid JSON: %w", err)
	}
	var trailing any
	if err := decoder.Decode(&trailing); err != io.EOF {
		if err == nil {
			return nil, fmt.Errorf("must contain exactly one JSON value")
		}
		return nil, fmt.Errorf("must contain exactly one JSON value: %w", err)
	}
	return value, nil
}

// ValidateChunkAnnotations requires annotations to already contain canonical
// JSON. CanonicalizeChunkAnnotations can be used at producer boundaries.
func ValidateChunkAnnotations(annotations ChunkAnnotations) error {
	canonical, err := CanonicalizeChunkAnnotations(annotations)
	if err != nil {
		return err
	}
	for namespace, raw := range annotations {
		if !bytes.Equal(raw, canonical[namespace]) {
			return fmt.Errorf("chunk annotation %q must use canonical JSON", namespace)
		}
	}
	return nil
}

// CloneChunkAnnotations returns a deep clone, including every raw JSON value.
func CloneChunkAnnotations(annotations ChunkAnnotations) ChunkAnnotations {
	if len(annotations) == 0 {
		return nil
	}
	cloned := make(ChunkAnnotations, len(annotations))
	for namespace, raw := range annotations {
		cloned[namespace] = append(json.RawMessage(nil), raw...)
	}
	return cloned
}

// CloneChunkPlan returns a deep clone of a chunk plan.
func CloneChunkPlan(plan ChunkPlan) ChunkPlan {
	cloned := ChunkPlan{
		SourceDigest: plan.SourceDigest,
		Ranges:       make([]ChunkRange, len(plan.Ranges)),
		Annotations:  CloneChunkAnnotations(plan.Annotations),
	}
	for i, chunkRange := range plan.Ranges {
		cloned.Ranges[i] = ChunkRange{
			StartUnitID: chunkRange.StartUnitID,
			EndUnitID:   chunkRange.EndUnitID,
			Annotations: CloneChunkAnnotations(chunkRange.Annotations),
		}
	}
	return cloned
}

// CanonicalizeChunkPlan returns a deep clone with canonical annotation bytes.
func CanonicalizeChunkPlan(plan ChunkPlan) (ChunkPlan, error) {
	canonical := CloneChunkPlan(plan)
	annotations, err := CanonicalizeChunkAnnotations(plan.Annotations)
	if err != nil {
		return ChunkPlan{}, fmt.Errorf("chunk plan annotations: %w", err)
	}
	canonical.Annotations = annotations
	for i := range plan.Ranges {
		annotations, err := CanonicalizeChunkAnnotations(plan.Ranges[i].Annotations)
		if err != nil {
			return ChunkPlan{}, fmt.Errorf("chunk plan range[%d] annotations: %w", i, err)
		}
		canonical.Ranges[i].Annotations = annotations
	}
	return canonical, nil
}

// ValidateChunkPlan validates a canonical plan against the current source.
// Ranges may contain gaps or overlap, but their start positions must increase.
func ValidateChunkPlan(doc *SourceDocument, plan ChunkPlan) error {
	if err := ValidateDocument(doc); err != nil {
		return fmt.Errorf("source document: %w", err)
	}
	if plan.SourceDigest != doc.Digest {
		return fmt.Errorf("chunk plan source_digest %q does not match source document digest %q", plan.SourceDigest, doc.Digest)
	}
	if len(plan.Ranges) == 0 {
		return fmt.Errorf("chunk plan ranges must not be empty")
	}
	if err := ValidateChunkAnnotations(plan.Annotations); err != nil {
		return fmt.Errorf("chunk plan annotations: %w", err)
	}
	previousStart := -1
	for i, chunkRange := range plan.Ranges {
		start, ok := UnitIndex(doc, chunkRange.StartUnitID)
		if !ok {
			return fmt.Errorf("chunk plan range[%d] start_unit_id %d was not found", i, chunkRange.StartUnitID)
		}
		end, ok := UnitIndex(doc, chunkRange.EndUnitID)
		if !ok {
			return fmt.Errorf("chunk plan range[%d] end_unit_id %d was not found", i, chunkRange.EndUnitID)
		}
		if start > end {
			return fmt.Errorf("chunk plan range[%d] start_unit_id %d appears after end_unit_id %d", i, chunkRange.StartUnitID, chunkRange.EndUnitID)
		}
		if start <= previousStart {
			return fmt.Errorf("chunk plan range[%d] start_unit_id %d does not appear after the previous range start", i, chunkRange.StartUnitID)
		}
		if err := ValidateChunkAnnotations(chunkRange.Annotations); err != nil {
			return fmt.Errorf("chunk plan range[%d] annotations: %w", i, err)
		}
		previousStart = start
	}
	return nil
}

// MaterializeChunkPlan deterministically expands a validated plan into chunks.
func MaterializeChunkPlan(doc *SourceDocument, plan ChunkPlan) ([]Chunk, error) {
	if err := ValidateChunkPlan(doc, plan); err != nil {
		return nil, err
	}
	chunks := make([]Chunk, 0, len(plan.Ranges))
	for index, chunkRange := range plan.Ranges {
		start, _ := UnitIndex(doc, chunkRange.StartUnitID)
		end, _ := UnitIndex(doc, chunkRange.EndUnitID)
		units := cloneSourceUnits(doc.Units[start : end+1])
		content, err := json.Marshal(struct {
			Units []SourceUnit `json:"units"`
		}{Units: units})
		if err != nil {
			return nil, fmt.Errorf("encode chunk plan range[%d]: %w", index, err)
		}
		chunks = append(chunks, Chunk{
			ID:        fmt.Sprintf("chunk-%06d", index+1),
			SourceID:  doc.ID,
			Index:     index,
			Ref:       SourceRef{SourceID: doc.ID, StartUnitID: chunkRange.StartUnitID, EndUnitID: chunkRange.EndUnitID},
			Content:   content,
			MediaType: "application/json",
			Units:     units,
			Metadata: map[string]any{
				"start_unit_id": chunkRange.StartUnitID,
				"end_unit_id":   chunkRange.EndUnitID,
				"unit_count":    len(units),
			},
			Annotations:     CloneChunkAnnotations(chunkRange.Annotations),
			PlanAnnotations: CloneChunkAnnotations(plan.Annotations),
		})
	}
	return chunks, nil
}

func cloneSourceUnits(units []SourceUnit) []SourceUnit {
	if len(units) == 0 {
		return nil
	}
	cloned := make([]SourceUnit, len(units))
	for i, unit := range units {
		cloned[i] = unit
		cloned[i].Metadata = cloneJSONMap(unit.Metadata)
	}
	return cloned
}

func cloneJSONMap(values map[string]any) map[string]any {
	if len(values) == 0 {
		return nil
	}
	cloned := make(map[string]any, len(values))
	for key, value := range values {
		cloned[key] = value
	}
	return cloned
}