seriatim/schema/output.go

package schema

import (
	"bytes"
	"embed"
	"encoding/json"
	"fmt"
	"sync"

	"github.com/santhosh-tekuri/jsonschema/v6"
)

//go:embed *.schema.json
var schemaFS embed.FS

const (
	fullOutputSchemaPath         = "full-output.schema.json"
	intermediateOutputSchemaPath = "intermediate-output.schema.json"
	minimalOutputSchemaPath      = "minimal-output.schema.json"
)

var (
	compiledSchemas = make(map[string]*jsonschema.Schema)
	compileErrs     = make(map[string]error)
	compileMu       sync.Mutex
)

// Transcript is seriatim's full public JSON output contract.
type Transcript struct {
	Metadata      Metadata       `json:"metadata"`
	Segments      []Segment      `json:"segments"`
	OverlapGroups []OverlapGroup `json:"overlap_groups"`
}

// IntermediateTranscript is seriatim's intermediate public JSON output contract.
type IntermediateTranscript struct {
	Metadata IntermediateMetadata  `json:"metadata"`
	Segments []IntermediateSegment `json:"segments"`
}

// MinimalTranscript is seriatim's compact public JSON output contract.
type MinimalTranscript struct {
	Metadata MinimalMetadata  `json:"metadata"`
	Segments []MinimalSegment `json:"segments"`
}

// Metadata records the pipeline configuration that produced an artifact.
type Metadata struct {
	Application           string   `json:"application"`
	Version               string   `json:"version"`
	InputReader           string   `json:"input_reader"`
	InputFiles            []string `json:"input_files"`
	PreprocessingModules  []string `json:"preprocessing_modules"`
	PostprocessingModules []string `json:"postprocessing_modules"`
	OutputModules         []string `json:"output_modules"`
}

// IntermediateMetadata records intermediate artifact identity.
type IntermediateMetadata struct {
	Application  string `json:"application"`
	Version      string `json:"version"`
	OutputSchema string `json:"output_schema"`
}

// MinimalMetadata records minimal artifact identity.
type MinimalMetadata struct {
	Application  string `json:"application"`
	Version      string `json:"version"`
	OutputSchema string `json:"output_schema"`
}

// Segment is the public transcript segment shape.
type Segment struct {
	ID                 int      `json:"id"`
	Source             string   `json:"source"`
	SourceSegmentIndex *int     `json:"source_segment_index,omitempty"`
	SourceRef          string   `json:"source_ref,omitempty"`
	DerivedFrom        []string `json:"derived_from,omitempty"`
	Speaker            string   `json:"speaker"`
	Start              float64  `json:"start"`
	End                float64  `json:"end"`
	Text               string   `json:"text"`
	Categories         []string `json:"categories,omitempty"`
	OverlapGroupID     int      `json:"overlap_group_id,omitempty"`
}

// IntermediateSegment is the compact public transcript segment shape with
// categories.
type IntermediateSegment struct {
	ID         int      `json:"id"`
	Start      float64  `json:"start"`
	End        float64  `json:"end"`
	Speaker    string   `json:"speaker"`
	Text       string   `json:"text"`
	Categories []string `json:"categories,omitempty"`
}

// MinimalSegment is the compact public transcript segment shape.
type MinimalSegment struct {
	ID      int     `json:"id"`
	Start   float64 `json:"start"`
	End     float64 `json:"end"`
	Speaker string  `json:"speaker"`
	Text    string  `json:"text"`
}

// OverlapGroup describes a detected overlapping speech region.
type OverlapGroup struct {
	ID         int      `json:"id"`
	Start      float64  `json:"start"`
	End        float64  `json:"end"`
	Segments   []string `json:"segments"`
	Speakers   []string `json:"speakers"`
	Class      string   `json:"class"`
	Resolution string   `json:"resolution"`
}

// ValidateTranscript validates a full transcript against the public JSON
// schema and seriatim-specific semantic rules.
func ValidateTranscript(transcript Transcript) error {
	if err := validateSemantics(transcript); err != nil {
		return err
	}

	data, err := json.Marshal(transcript)
	if err != nil {
		return fmt.Errorf("marshal transcript for schema validation: %w", err)
	}
	return ValidateJSON(data)
}

// ValidateIntermediateTranscript validates the intermediate transcript against
// the intermediate JSON schema and seriatim-specific semantic rules.
func ValidateIntermediateTranscript(transcript IntermediateTranscript) error {
	if err := validateIntermediateSemantics(transcript); err != nil {
		return err
	}

	data, err := json.Marshal(transcript)
	if err != nil {
		return fmt.Errorf("marshal intermediate transcript for schema validation: %w", err)
	}
	return ValidateIntermediateJSON(data)
}

// ValidateMinimalTranscript validates a minimal transcript against the minimal
// JSON schema and seriatim-specific semantic rules.
func ValidateMinimalTranscript(transcript MinimalTranscript) error {
	if err := validateMinimalSemantics(transcript); err != nil {
		return err
	}

	data, err := json.Marshal(transcript)
	if err != nil {
		return fmt.Errorf("marshal minimal transcript for schema validation: %w", err)
	}
	return ValidateMinimalJSON(data)
}

// ValidateJSON validates serialized output JSON against the public schema.
func ValidateJSON(data []byte) error {
	return validateJSONWithSchema(data, fullOutputSchemaPath)
}

// ValidateIntermediateJSON validates serialized intermediate output JSON
// against the intermediate public schema.
func ValidateIntermediateJSON(data []byte) error {
	return validateJSONWithSchema(data, intermediateOutputSchemaPath)
}

// ValidateMinimalJSON validates serialized minimal output JSON against the
// minimal public schema.
func ValidateMinimalJSON(data []byte) error {
	return validateJSONWithSchema(data, minimalOutputSchemaPath)
}

func validateJSONWithSchema(data []byte, schemaPath string) error {
	var value any
	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.UseNumber()
	if err := decoder.Decode(&value); err != nil {
		return fmt.Errorf("decode output JSON for schema validation: %w", err)
	}

	compiled, err := outputSchema(schemaPath)
	if err != nil {
		return err
	}
	if err := compiled.Validate(value); err != nil {
		return fmt.Errorf("output schema validation failed: %w", err)
	}
	return nil
}

func outputSchema(schemaPath string) (*jsonschema.Schema, error) {
	compileMu.Lock()
	defer compileMu.Unlock()

	if compiled, exists := compiledSchemas[schemaPath]; exists {
		return compiled, compileErrs[schemaPath]
	}

	data, err := schemaFS.ReadFile(schemaPath)
	if err != nil {
		compileErrs[schemaPath] = fmt.Errorf("read embedded output schema: %w", err)
		return nil, compileErrs[schemaPath]
	}
	var schemaDocument any
	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.UseNumber()
	if err := decoder.Decode(&schemaDocument); err != nil {
		compileErrs[schemaPath] = fmt.Errorf("decode embedded output schema: %w", err)
		return nil, compileErrs[schemaPath]
	}

	compiler := jsonschema.NewCompiler()
	if err := compiler.AddResource(schemaPath, schemaDocument); err != nil {
		compileErrs[schemaPath] = fmt.Errorf("load embedded output schema: %w", err)
		return nil, compileErrs[schemaPath]
	}
	compiled, err := compiler.Compile(schemaPath)
	if err != nil {
		compileErrs[schemaPath] = fmt.Errorf("compile embedded output schema: %w", err)
		return nil, compileErrs[schemaPath]
	}
	compiledSchemas[schemaPath] = compiled
	return compiled, nil
}

func validateSemantics(transcript Transcript) error {
	for index, segment := range transcript.Segments {
		wantID := index + 1
		if segment.ID != wantID {
			return fmt.Errorf("segment %d has id %d; want %d", index, segment.ID, wantID)
		}
		if segment.End < segment.Start {
			return fmt.Errorf("segment %d has end %.3f before start %.3f", index, segment.End, segment.Start)
		}
	}
	for index, group := range transcript.OverlapGroups {
		if group.End < group.Start {
			return fmt.Errorf("overlap_group %d has end %.3f before start %.3f", index, group.End, group.Start)
		}
	}
	return nil
}

func validateIntermediateSemantics(transcript IntermediateTranscript) error {
	for index, segment := range transcript.Segments {
		wantID := index + 1
		if segment.ID != wantID {
			return fmt.Errorf("segment %d has id %d; want %d", index, segment.ID, wantID)
		}
		if segment.End < segment.Start {
			return fmt.Errorf("segment %d has end %.3f before start %.3f", index, segment.End, segment.Start)
		}
	}
	return nil
}

func validateMinimalSemantics(transcript MinimalTranscript) error {
	for index, segment := range transcript.Segments {
		wantID := index + 1
		if segment.ID != wantID {
			return fmt.Errorf("segment %d has id %d; want %d", index, segment.ID, wantID)
		}
		if segment.End < segment.Start {
			return fmt.Errorf("segment %d has end %.3f before start %.3f", index, segment.End, segment.Start)
		}
	}
	return nil
}