notarius/internal/modules/generic/output/json/encoder.go

package json

import (
	"bytes"
	"context"
	stdjson "encoding/json"
	"fmt"
	"io"
	"mime"
	"regexp"
	"sort"
	"strings"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/chunkmap"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/evidencecontext"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
)

const Key = "json"

const contentTypeJSON = "application/json"

const chunkMapFileName = "chunk-map.json"

const evidenceContextFileName = "evidence-context.json"

var safeOutputFileChar = regexp.MustCompile(`[^A-Za-z0-9._-]`)

var _ contracts.OutputEncoder = (*Encoder)(nil)

type Options struct {
	IncludeChunkMap bool
	EvidenceContext pipeline.EvidenceContextPolicy
}

type Encoder struct {
	options Options
}

func New() *Encoder {
	return NewWithOptions(Options{})
}

func NewWithOptions(options Options) *Encoder {
	options.EvidenceContext.LaneIDs = append([]string(nil), options.EvidenceContext.LaneIDs...)
	return &Encoder{options: options}
}

func (e *Encoder) Key() string {
	return Key
}

func (e *Encoder) EvidenceContextPolicy() pipeline.EvidenceContextPolicy {
	if e == nil {
		return pipeline.EvidenceContextPolicy{}
	}
	policy := e.options.EvidenceContext
	policy.LaneIDs = append([]string(nil), policy.LaneIDs...)
	return policy
}

func (e *Encoder) Encode(ctx context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) {
	if e == nil {
		return contracts.OutputResult{}, encoderErrorf("encoder must not be nil")
	}
	if ctx == nil {
		return contracts.OutputResult{}, encoderErrorf("context must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return contracts.OutputResult{}, encoderErrorf("context error before encoding: %w", err)
	}

	files, err := logicalFiles(req, e.options)
	if err != nil {
		return contracts.OutputResult{}, err
	}
	return contracts.OutputResult{Files: files}, nil
}

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{
		Key:            Key,
		Stage:          pipeline.StageOutput,
		ExecutionClass: contracts.ExecutionClassDeterministic,
		Requires:       []string{"normalized"},
		Provides:       []string{"encoded"},
	}
}

func Register(registry *pipeline.OutputEncoderRegistry) error {
	return registry.RegisterBuilderWithProfileValidation(ModuleSpec(), validateOptions, validateProfileOptions, func(request pipeline.BuildRequest) (contracts.OutputEncoder, error) {
		options, err := DecodeOptions(request.Options)
		if err != nil {
			return nil, err
		}
		return &Encoder{options: options}, nil
	})
}

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, "include_chunk_map", "evidence_context"); err != nil {
		return Options{}, encoderErrorf("%w", err)
	}
	decoded := Options{}
	if value, ok := options["include_chunk_map"]; ok {
		enabled, ok := value.(bool)
		if !ok {
			return Options{}, encoderErrorf("option %q must be a boolean", "include_chunk_map")
		}
		decoded.IncludeChunkMap = enabled
	}
	if value, ok := options["evidence_context"]; ok {
		policy, err := decodeEvidenceContextPolicy(value)
		if err != nil {
			return Options{}, err
		}
		decoded.EvidenceContext = policy
	}
	return decoded, nil
}

func validateProfileOptions(context pipeline.OutputProfileOptionContext, options map[string]any) error {
	decoded, err := DecodeOptions(options)
	if err != nil || !decoded.EvidenceContext.Enabled {
		return err
	}
	configured := make(map[string]struct{}, len(context.LaneIDs))
	for _, laneID := range context.LaneIDs {
		configured[laneID] = struct{}{}
	}
	for _, laneID := range decoded.EvidenceContext.LaneIDs {
		if _, ok := configured[laneID]; !ok {
			return encoderErrorf("evidence_context lane %q is not configured", laneID)
		}
	}
	return nil
}

func decodeEvidenceContextPolicy(value any) (pipeline.EvidenceContextPolicy, error) {
	object, ok := value.(map[string]any)
	if !ok {
		return pipeline.EvidenceContextPolicy{}, encoderErrorf("option %q must be an object", "evidence_context")
	}
	if err := pipeline.RejectUnknownOptions(object, "enabled", "lanes", "window_units"); err != nil {
		return pipeline.EvidenceContextPolicy{}, encoderErrorf("evidence_context: %w", err)
	}
	enabledValue, ok := object["enabled"]
	if !ok {
		return pipeline.EvidenceContextPolicy{}, encoderErrorf("evidence_context option %q is required", "enabled")
	}
	enabled, ok := enabledValue.(bool)
	if !ok {
		return pipeline.EvidenceContextPolicy{}, encoderErrorf("evidence_context option %q must be a boolean", "enabled")
	}
	if !enabled {
		if _, ok := object["lanes"]; ok {
			return pipeline.EvidenceContextPolicy{}, encoderErrorf("evidence_context option %q is not allowed when disabled", "lanes")
		}
		if _, ok := object["window_units"]; ok {
			return pipeline.EvidenceContextPolicy{}, encoderErrorf("evidence_context option %q is not allowed when disabled", "window_units")
		}
		return pipeline.EvidenceContextPolicy{}, nil
	}
	rawLanes, ok := object["lanes"]
	if !ok {
		return pipeline.EvidenceContextPolicy{}, encoderErrorf("evidence_context option %q is required when enabled", "lanes")
	}
	lanes, err := decodeEvidenceLaneIDs(rawLanes)
	if err != nil {
		return pipeline.EvidenceContextPolicy{}, err
	}
	windowUnits := 3
	if rawWindow, ok := object["window_units"]; ok {
		value, ok := rawWindow.(int)
		if !ok {
			return pipeline.EvidenceContextPolicy{}, encoderErrorf("evidence_context option %q must be an integer", "window_units")
		}
		if value < 0 {
			return pipeline.EvidenceContextPolicy{}, encoderErrorf("evidence_context option %q must not be negative", "window_units")
		}
		windowUnits = value
	}
	return pipeline.EvidenceContextPolicy{Enabled: true, WindowUnits: windowUnits, LaneIDs: lanes}, nil
}

func decodeEvidenceLaneIDs(value any) ([]string, error) {
	var raw []any
	switch typed := value.(type) {
	case []any:
		raw = typed
	case []string:
		raw = make([]any, len(typed))
		for i := range typed {
			raw[i] = typed[i]
		}
	default:
		return nil, encoderErrorf("evidence_context option %q must be an array", "lanes")
	}
	if len(raw) == 0 {
		return nil, encoderErrorf("evidence_context option %q must not be empty", "lanes")
	}
	seen := make(map[string]struct{}, len(raw))
	lanes := make([]string, 0, len(raw))
	for _, value := range raw {
		lane, ok := value.(string)
		if !ok {
			return nil, encoderErrorf("evidence_context lane values must be strings")
		}
		lane = strings.TrimSpace(lane)
		if lane == "" {
			return nil, encoderErrorf("evidence_context lane values must not be empty")
		}
		if _, ok := seen[lane]; ok {
			return nil, encoderErrorf("evidence_context lane %q is duplicated", lane)
		}
		seen[lane] = struct{}{}
		lanes = append(lanes, lane)
	}
	sort.Strings(lanes)
	return lanes, nil
}

type indexFile struct {
	ManifestFile    string            `json:"manifest_file"`
	OutputFiles     []outputFileIndex `json:"output_files"`
	RejectedFile    string            `json:"rejected_file"`
	WarningsFile    string            `json:"warnings_file"`
	ChunkMap        *artifactIndex    `json:"chunk_map,omitempty"`
	EvidenceContext *artifactIndex    `json:"evidence_context,omitempty"`
}

type artifactIndex struct {
	ArtifactKind  contracts.ArtifactKind `json:"artifact_kind"`
	File          string                 `json:"file"`
	MediaType     string                 `json:"media_type"`
	SchemaID      string                 `json:"schema_id"`
	SchemaName    string                 `json:"schema_name"`
	SchemaVersion string                 `json:"schema_version"`
}

type outputFileIndex struct {
	LaneID     string `json:"lane_id"`
	MediaType  string `json:"media_type,omitempty"`
	File       string `json:"file"`
	ModuleKey  string `json:"module_key,omitempty"`
	SchemaID   string `json:"schema_id,omitempty"`
	SchemaName string `json:"schema_name,omitempty"`
	SchemaVer  string `json:"schema_version,omitempty"`
}

type rejectedFile struct {
	Rejected []contracts.RejectedOutput `json:"rejected"`
}

type warningsFile struct {
	Warnings []contracts.Warning `json:"warnings"`
}

func logicalFiles(req contracts.OutputRequest, options Options) ([]contracts.OutputFile, error) {
	outputs := cloneNormalizeOutputs(req.NormalizeOutputs)
	sort.SliceStable(outputs, func(i, j int) bool {
		return outputs[i].LaneID < outputs[j].LaneID
	})

	outputIndexes := make([]outputFileIndex, 0, len(outputs))
	files := make([]contracts.OutputFile, 0, len(outputs)+5)
	manifestFile, err := jsonFile("manifest.json", req.Manifest)
	if err != nil {
		return nil, err
	}
	files = append(files, manifestFile)

	usedOutputFiles := make(map[string]string, len(outputs))
	for _, output := range outputs {
		name, err := outputFileName(output.LaneID)
		if err != nil {
			return nil, err
		}
		if existingLane, ok := usedOutputFiles[name]; ok {
			return nil, encoderErrorf("lanes %q and %q produce duplicate output file %q", existingLane, output.LaneID, name)
		}
		usedOutputFiles[name] = output.LaneID
		outputIndexes = append(outputIndexes, outputFileIndex{
			LaneID:     output.LaneID,
			MediaType:  output.Artifact.MediaType,
			File:       name,
			ModuleKey:  output.NormalizerKey,
			SchemaID:   output.Artifact.Schema.ID,
			SchemaName: output.Artifact.Schema.Name,
			SchemaVer:  output.Artifact.Schema.Version,
		})
		file, err := serializedOutputFile(name, output.Artifact)
		if err != nil {
			return nil, err
		}
		files = append(files, file)
	}

	index := indexFile{
		ManifestFile: "manifest.json",
		OutputFiles:  outputIndexes,
		RejectedFile: "rejected.json",
		WarningsFile: "warnings.json",
	}
	if options.IncludeChunkMap && req.ChunkMap != nil {
		chunkMapOutput, chunkMapDescriptor, err := serializedChunkMapFile(*req.ChunkMap)
		if err != nil {
			return nil, err
		}
		files = append(files, chunkMapOutput)
		index.ChunkMap = &chunkMapDescriptor
	}
	if req.EvidenceContext != nil {
		evidenceOutput, evidenceDescriptor, err := serializedEvidenceContextFile(*req.EvidenceContext)
		if err != nil {
			return nil, err
		}
		files = append(files, evidenceOutput)
		index.EvidenceContext = &evidenceDescriptor
	}
	indexOutput, err := jsonFile("index.json", index)
	if err != nil {
		return nil, err
	}
	rejectedOutput, err := jsonFile("rejected.json", rejectedFile{Rejected: cloneRejected(req.Rejected)})
	if err != nil {
		return nil, err
	}
	warningsOutput, err := jsonFile("warnings.json", warningsFile{Warnings: cloneWarnings(req.Warnings)})
	if err != nil {
		return nil, err
	}
	files = append(files, indexOutput, rejectedOutput, warningsOutput)
	sort.Slice(files, func(i, j int) bool {
		return files[i].Name < files[j].Name
	})
	return files, nil
}

func serializedChunkMapFile(artifact contracts.SerializedArtifact) (contracts.OutputFile, artifactIndex, error) {
	if artifact.Kind != chunkmap.ArtifactKind {
		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("chunk map has unexpected artifact kind %q", artifact.Kind)
	}
	if artifact.Schema.ID != chunkmap.SchemaID || artifact.Schema.Name != chunkmap.SchemaName || artifact.Schema.Version != chunkmap.SchemaVersion {
		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("chunk map has unexpected schema identity")
	}
	if strings.TrimSpace(artifact.MediaType) != chunkmap.MediaType {
		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("chunk map has unsupported media type %q", artifact.MediaType)
	}
	if _, err := chunkmap.New().Decode(artifact.Content); err != nil {
		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("decode chunk map: %w", err)
	}
	file, err := serializedOutputFile(chunkMapFileName, artifact)
	if err != nil {
		return contracts.OutputFile{}, artifactIndex{}, err
	}
	return file, artifactIndex{
		ArtifactKind:  artifact.Kind,
		File:          chunkMapFileName,
		MediaType:     chunkmap.MediaType,
		SchemaID:      artifact.Schema.ID,
		SchemaName:    artifact.Schema.Name,
		SchemaVersion: artifact.Schema.Version,
	}, nil
}

func serializedEvidenceContextFile(artifact contracts.SerializedArtifact) (contracts.OutputFile, artifactIndex, error) {
	codec := evidencecontext.New()
	expected := codec.Schema()
	if artifact.Kind != evidencecontext.ArtifactKind ||
		artifact.Schema.ID != expected.ID || artifact.Schema.Name != expected.Name || artifact.Schema.Version != expected.Version ||
		contracts.DigestArtifactSchema(artifact.Schema) != contracts.DigestArtifactSchema(expected) ||
		strings.TrimSpace(artifact.MediaType) != evidencecontext.MediaType {
		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("evidence context artifact is invalid")
	}
	if _, err := codec.Decode(artifact.Content); err != nil {
		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("evidence context artifact is invalid")
	}
	file, err := serializedOutputFile(evidenceContextFileName, artifact)
	if err != nil {
		return contracts.OutputFile{}, artifactIndex{}, encoderErrorf("evidence context artifact is invalid")
	}
	return file, artifactIndex{
		ArtifactKind:  evidencecontext.ArtifactKind,
		File:          evidenceContextFileName,
		MediaType:     evidencecontext.MediaType,
		SchemaID:      expected.ID,
		SchemaName:    expected.Name,
		SchemaVersion: expected.Version,
	}, nil
}

func serializedOutputFile(name string, artifact contracts.SerializedArtifact) (contracts.OutputFile, error) {
	content := append([]byte(nil), artifact.Content...)
	if len(content) == 0 {
		content = []byte("null")
	}
	mediaType := strings.TrimSpace(artifact.MediaType)
	if mediaType == "" {
		mediaType = "application/octet-stream"
	}
	if !isJSONMediaType(mediaType) {
		return contracts.OutputFile{}, encoderErrorf("normalized output %q has unsupported media type %q", name, mediaType)
	}
	decoder := stdjson.NewDecoder(bytes.NewReader(content))
	decoder.UseNumber()
	var decoded any
	if err := decoder.Decode(&decoded); err != nil {
		return contracts.OutputFile{}, encoderErrorf("normalized output %q contains invalid JSON: %w", name, err)
	}
	var trailing any
	if err := decoder.Decode(&trailing); err != io.EOF {
		return contracts.OutputFile{}, encoderErrorf("normalized output %q contains multiple JSON values", name)
	}
	pretty, err := marshalPretty(decoded)
	if err != nil {
		return contracts.OutputFile{}, err
	}
	return contracts.OutputFile{
		Name:        name,
		ContentType: mediaType,
		Bytes:       pretty,
	}, nil
}

func isJSONMediaType(mediaType string) bool {
	base, _, err := mime.ParseMediaType(strings.TrimSpace(mediaType))
	if err != nil {
		base = strings.TrimSpace(mediaType)
	}
	return strings.EqualFold(base, contentTypeJSON)
}

func jsonFile(name string, value any) (contracts.OutputFile, error) {
	data, err := marshalPretty(value)
	if err != nil {
		return contracts.OutputFile{}, encoderErrorf("encode %s: %w", name, err)
	}
	return contracts.OutputFile{
		Name:        name,
		ContentType: contentTypeJSON,
		Bytes:       data,
	}, nil
}

func marshalPretty(value any) ([]byte, error) {
	data, err := stdjson.MarshalIndent(value, "", "  ")
	if err != nil {
		return nil, err
	}
	return append(data, '\n'), nil
}

func outputFileName(laneID string) (string, error) {
	sanitized := safeOutputFileChar.ReplaceAllString(strings.TrimSpace(laneID), "_")
	for strings.Contains(sanitized, "..") {
		sanitized = strings.ReplaceAll(sanitized, "..", "__")
	}
	sanitized = strings.Trim(sanitized, "._")
	if sanitized == "" {
		return "", encoderErrorf("lane id %q cannot produce a safe file name", laneID)
	}
	return "lanes/" + sanitized + ".json", nil
}

func cloneNormalizeOutputs(outputs []contracts.SerializedOutput) []contracts.SerializedOutput {
	if len(outputs) == 0 {
		return nil
	}
	out := make([]contracts.SerializedOutput, 0, len(outputs))
	for _, output := range outputs {
		out = append(out, contracts.CloneSerializedOutput(output))
	}
	return out
}

func cloneRejected(rejected []contracts.RejectedOutput) []contracts.RejectedOutput {
	if len(rejected) == 0 {
		return []contracts.RejectedOutput{}
	}
	return append([]contracts.RejectedOutput(nil), rejected...)
}

func cloneWarnings(warnings []contracts.Warning) []contracts.Warning {
	if len(warnings) == 0 {
		return []contracts.Warning{}
	}
	return append([]contracts.Warning(nil), warnings...)
}

func encoderErrorf(format string, args ...any) error {
	return fmt.Errorf("json output encoder: "+format, args...)
}