notarius/internal/framework/evidencecontext/codec.go

package evidencecontext

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

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"github.com/santhosh-tekuri/jsonschema/v6"
)

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

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

var (
	loadSchemaOnce sync.Once
	loadedSchema   []byte
	compiledSchema *jsonschema.Schema
	loadSchemaErr  error
)

// Codec owns strict serialization for the durable evidence-context 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 }

// 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 Document) ([]byte, error) {
	if _, err := c.schemaBytes(); err != nil {
		return nil, err
	}
	canonical, err := canonicalize(value)
	if err != nil {
		return nil, fmt.Errorf("encode evidence context: %w", err)
	}
	content, err := json.Marshal(canonical)
	if err != nil {
		return nil, fmt.Errorf("encode evidence context: %w", err)
	}
	if err := validateSchemaInstance(content); err != nil {
		return nil, fmt.Errorf("encode evidence context: %w", err)
	}
	return content, nil
}

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

func (c *Codec) schemaBytes() ([]byte, error) {
	loadSchemaOnce.Do(loadAndCompileSchema)
	if loadSchemaErr != nil {
		return nil, loadSchemaErr
	}
	return append([]byte(nil), loadedSchema...), nil
}

func loadAndCompileSchema() {
	raw, err := schemaAssets.ReadFile("assets/schemas/source_evidence_context.v1.json")
	if err != nil {
		loadSchemaErr = fmt.Errorf("read source evidence context schema: %w", err)
		return
	}
	var identity struct {
		ID       string   `json:"$id"`
		Title    string   `json:"title"`
		Type     string   `json:"type"`
		Required []string `json:"required"`
	}
	if err := json.Unmarshal(raw, &identity); err != nil {
		loadSchemaErr = fmt.Errorf("decode source evidence context schema: %w", err)
		return
	}
	if identity.ID != SchemaID || identity.Title != SchemaName || identity.Type != "object" || !hasRequiredFields(identity.Required) {
		loadSchemaErr = fmt.Errorf("source evidence context schema identity or required fields are invalid")
		return
	}
	schemaDocument, err := jsonschema.UnmarshalJSON(bytes.NewReader(raw))
	if err != nil {
		loadSchemaErr = fmt.Errorf("parse source evidence context schema: %w", err)
		return
	}
	compiler := jsonschema.NewCompiler()
	if err := compiler.AddResource("source-evidence-context-schema.json", schemaDocument); err != nil {
		loadSchemaErr = fmt.Errorf("load source evidence context schema: %w", err)
		return
	}
	compiled, err := compiler.Compile("source-evidence-context-schema.json")
	if err != nil {
		loadSchemaErr = fmt.Errorf("compile source evidence context schema: %w", err)
		return
	}
	loadedSchema = append([]byte(nil), raw...)
	compiledSchema = compiled
}

func validateSchemaInstance(content []byte) error {
	instance, err := jsonschema.UnmarshalJSON(bytes.NewReader(content))
	if err != nil {
		return fmt.Errorf("payload is not valid JSON: %w", err)
	}
	if err := compiledSchema.Validate(instance); err != nil {
		return fmt.Errorf("payload does not conform to source evidence context schema: %w", err)
	}
	return nil
}

func hasRequiredFields(required []string) bool {
	want := map[string]bool{"source_id": true, "source_digest": true, "window_units": true, "selected_lanes": true, "contexts": true}
	for _, field := range required {
		delete(want, field)
	}
	return len(want) == 0
}

func canonicalize(value Document) (Document, error) {
	owned, err := clone(value)
	if err != nil {
		return Document{}, err
	}
	value = owned
	if err := requireIdentity("source_id", value.SourceID); err != nil {
		return Document{}, err
	}
	if !digestPattern.MatchString(value.SourceDigest) {
		return Document{}, fmt.Errorf("source_digest must be a sha256 digest")
	}
	if value.WindowUnits < 0 {
		return Document{}, fmt.Errorf("window_units must not be negative")
	}
	if err := validateSelectedLanes(value.SelectedLanes); err != nil {
		return Document{}, err
	}
	if value.Contexts == nil {
		value.Contexts = make([]Context, 0)
	}
	selected := make(map[string]struct{}, len(value.SelectedLanes))
	for _, laneID := range value.SelectedLanes {
		selected[laneID] = struct{}{}
	}
	seenUnits := make(map[int]struct{})
	for contextIndex := range value.Contexts {
		context, err := canonicalizeContext(value.SourceID, selected, seenUnits, value.Contexts[contextIndex], contextIndex)
		if err != nil {
			return Document{}, err
		}
		value.Contexts[contextIndex] = context
	}
	return value, nil
}

func validateSelectedLanes(lanes []string) error {
	if len(lanes) == 0 {
		return fmt.Errorf("selected_lanes must not be empty")
	}
	for index, laneID := range lanes {
		if err := requireIdentity(fmt.Sprintf("selected_lanes[%d]", index), laneID); err != nil {
			return err
		}
		if index > 0 && lanes[index-1] >= laneID {
			return fmt.Errorf("selected_lanes must be unique and in lexical order")
		}
	}
	return nil
}

func canonicalizeContext(sourceID string, selected map[string]struct{}, seenUnits map[int]struct{}, value Context, contextIndex int) (Context, error) {
	prefix := fmt.Sprintf("contexts[%d]", contextIndex)
	if len(value.EvidenceRefs) == 0 {
		return Context{}, fmt.Errorf("%s.evidence_refs must not be empty", prefix)
	}
	if len(value.Units) == 0 {
		return Context{}, fmt.Errorf("%s.units must not be empty", prefix)
	}
	if err := validateRefIdentity(sourceID, value.ContextRef, prefix+".context_ref"); err != nil {
		return Context{}, err
	}
	positions := make(map[int]int, len(value.Units))
	for unitIndex := range value.Units {
		unit, err := cloneSourceUnit(value.Units[unitIndex])
		if err != nil {
			return Context{}, fmt.Errorf("%s.units[%d]: %w", prefix, unitIndex, err)
		}
		if unit.ID <= 0 || strings.TrimSpace(unit.Kind) == "" || strings.TrimSpace(unit.Text) == "" {
			return Context{}, fmt.Errorf("%s.units[%d] has invalid required fields", prefix, unitIndex)
		}
		if err := validateRefIdentity(sourceID, unit.Ref, fmt.Sprintf("%s.units[%d].ref", prefix, unitIndex)); err != nil {
			return Context{}, err
		}
		if unit.Ref.StartUnitID != unit.ID || unit.Ref.EndUnitID != unit.ID {
			return Context{}, fmt.Errorf("%s.units[%d].ref must identify unit id %d", prefix, unitIndex, unit.ID)
		}
		if _, exists := positions[unit.ID]; exists {
			return Context{}, fmt.Errorf("%s.units contains duplicate unit id %d", prefix, unit.ID)
		}
		if _, exists := seenUnits[unit.ID]; exists {
			return Context{}, fmt.Errorf("contexts contain duplicate unit id %d", unit.ID)
		}
		positions[unit.ID] = unitIndex
		seenUnits[unit.ID] = struct{}{}
		value.Units[unitIndex] = unit
	}
	if value.ContextRef.StartUnitID != value.Units[0].ID || value.ContextRef.EndUnitID != value.Units[len(value.Units)-1].ID {
		return Context{}, fmt.Errorf("%s.context_ref must identify the first and last units", prefix)
	}
	for evidenceIndex := range value.EvidenceRefs {
		evidence := value.EvidenceRefs[evidenceIndex]
		if _, ok := selected[evidence.LaneID]; !ok {
			return Context{}, fmt.Errorf("%s.evidence_refs[%d].lane_id is not selected", prefix, evidenceIndex)
		}
		if err := requireIdentity(fmt.Sprintf("%s.evidence_refs[%d].lane_id", prefix, evidenceIndex), evidence.LaneID); err != nil {
			return Context{}, err
		}
		if err := validateRefIdentity(sourceID, evidence.SourceRef, fmt.Sprintf("%s.evidence_refs[%d].source_ref", prefix, evidenceIndex)); err != nil {
			return Context{}, err
		}
		start, startOK := positions[evidence.SourceRef.StartUnitID]
		end, endOK := positions[evidence.SourceRef.EndUnitID]
		if !startOK || !endOK || start > end {
			return Context{}, fmt.Errorf("%s.evidence_refs[%d].source_ref is outside context units", prefix, evidenceIndex)
		}
		if evidenceIndex > 0 && !lessEvidenceRef(value.EvidenceRefs[evidenceIndex-1], evidence) {
			return Context{}, fmt.Errorf("%s.evidence_refs must be unique and in deterministic order", prefix)
		}
	}
	return value, nil
}

func validateRefIdentity(sourceID string, ref source.SourceRef, field string) error {
	if ref.SourceID != sourceID {
		return fmt.Errorf("%s.source_id does not match source_id", field)
	}
	if ref.StartUnitID <= 0 || ref.EndUnitID <= 0 {
		return fmt.Errorf("%s endpoints must be positive", field)
	}
	return nil
}

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

func clone(value Document) (Document, error) {
	value.SelectedLanes = append([]string(nil), value.SelectedLanes...)
	if value.Contexts == nil {
		value.Contexts = make([]Context, 0)
	} else {
		contexts := make([]Context, len(value.Contexts))
		for contextIndex, context := range value.Contexts {
			contexts[contextIndex].ContextRef = context.ContextRef
			contexts[contextIndex].EvidenceRefs = append([]EvidenceRef(nil), context.EvidenceRefs...)
			contexts[contextIndex].Units = make([]source.SourceUnit, len(context.Units))
			for unitIndex, unit := range context.Units {
				cloned, err := cloneSourceUnit(unit)
				if err != nil {
					return Document{}, fmt.Errorf("clone contexts[%d].units[%d]: %w", contextIndex, unitIndex, err)
				}
				contexts[contextIndex].Units[unitIndex] = cloned
			}
		}
		value.Contexts = contexts
	}
	return value, nil
}

func cloneSourceUnit(unit source.SourceUnit) (source.SourceUnit, error) {
	metadata, err := source.CloneMetadata(unit.Metadata)
	if err != nil {
		return source.SourceUnit{}, fmt.Errorf("clone metadata: %w", err)
	}
	unit.Metadata = metadata
	return unit, nil
}