seriatim/internal/cli/trim.go

package cli

import (
	"encoding/json"
	"fmt"
	"os"
	"sort"

	"github.com/spf13/cobra"

	"gitea.maximumdirect.net/eric/seriatim/internal/config"
	"gitea.maximumdirect.net/eric/seriatim/internal/report"
	triminternal "gitea.maximumdirect.net/eric/seriatim/internal/trim"
)

type trimAuditReport struct {
	Operation               string          `json:"operation"`
	InputFile               string          `json:"input_file"`
	OutputFile              string          `json:"output_file"`
	InputSchema             string          `json:"input_schema"`
	OutputSchema            string          `json:"output_schema"`
	Mode                    string          `json:"mode"`
	Selector                string          `json:"selector"`
	SelectedIDs             []int           `json:"selected_ids"`
	AllowEmpty              bool            `json:"allow_empty"`
	InputSegmentCount       int             `json:"input_segment_count"`
	RetainedSegmentCount    int             `json:"retained_segment_count"`
	RemovedSegmentCount     int             `json:"removed_segment_count"`
	RemovedInputIDs         []int           `json:"removed_input_ids"`
	OldToNewIDMapping       []trimIDMapping `json:"old_to_new_id_mapping"`
	OverlapGroupsRecomputed bool            `json:"overlap_groups_recomputed"`
}

type trimIDMapping struct {
	OldID int `json:"old_id"`
	NewID int `json:"new_id"`
}

func newTrimCommand() *cobra.Command {
	var opts config.TrimOptions

	cmd := &cobra.Command{
		Use:   "trim",
		Short: "Trim an existing seriatim transcript artifact by segment ID",
		RunE: func(cmd *cobra.Command, args []string) error {
			trimOpts := opts
			if !cmd.Flags().Changed("output-schema") {
				trimOpts.OutputSchema = ""
			}

			cfg, err := config.NewTrimConfig(trimOpts)
			if err != nil {
				return err
			}

			selector, err := triminternal.ParseSelector(cfg.Selector)
			if err != nil {
				return fmt.Errorf("invalid selector %q: %w", cfg.Selector, err)
			}

			data, err := os.ReadFile(cfg.InputFile)
			if err != nil {
				return fmt.Errorf("read --input-file %q: %w", cfg.InputFile, err)
			}

			artifact, err := triminternal.ParseArtifactJSON(data)
			if err != nil {
				return fmt.Errorf("--input-file %q: %w", cfg.InputFile, err)
			}
			inputSegmentCount := artifact.SegmentCount()
			inputSchema := artifact.Schema

			mode := triminternal.ModeKeep
			if cfg.Mode == "remove" {
				mode = triminternal.ModeRemove
			}

			trimmed, err := triminternal.ApplyArtifact(artifact, triminternal.Options{
				Mode:       mode,
				Selector:   selector,
				AllowEmpty: cfg.AllowEmpty,
			})
			if err != nil {
				return err
			}

			outputSchema := artifact.Schema
			if cfg.OutputSchema != "" {
				outputSchema = cfg.OutputSchema
			}

			outputArtifact, err := triminternal.ConvertArtifact(trimmed.Artifact, outputSchema)
			if err != nil {
				return err
			}

			if err := triminternal.ValidateArtifact(outputArtifact); err != nil {
				return fmt.Errorf("validate trimmed output: %w", err)
			}

			if err := writeOutputJSON(cfg.OutputFile, outputArtifact.Value()); err != nil {
				return err
			}

			if cfg.ReportFile != "" {
				audit := trimAuditReport{
					Operation:               "trim",
					InputFile:               cfg.InputFile,
					OutputFile:              cfg.OutputFile,
					InputSchema:             inputSchema,
					OutputSchema:            outputArtifact.Schema,
					Mode:                    cfg.Mode,
					Selector:                cfg.Selector,
					SelectedIDs:             selector.IDs(),
					AllowEmpty:              cfg.AllowEmpty,
					InputSegmentCount:       inputSegmentCount,
					RetainedSegmentCount:    len(trimmed.OldToNewID),
					RemovedSegmentCount:     len(trimmed.RemovedIDs),
					RemovedInputIDs:         append([]int(nil), trimmed.RemovedIDs...),
					OldToNewIDMapping:       orderedIDMapping(trimmed.OldToNewID),
					OverlapGroupsRecomputed: trimmed.OverlapGroupsRecomputed,
				}
				auditJSON, err := json.Marshal(audit)
				if err != nil {
					return fmt.Errorf("marshal trim audit report: %w", err)
				}

				rpt := report.Report{
					Metadata: report.Metadata{
						Application:   outputArtifact.Application(),
						Version:       outputArtifact.Version(),
						InputReader:   "trim-artifact",
						InputFiles:    []string{cfg.InputFile},
						OutputModules: []string{"json"},
					},
					Events: []report.Event{
						report.Info("trim", "trim", fmt.Sprintf("trimmed %d input segment(s) into %d output segment(s) with mode=%s", inputSegmentCount, outputArtifact.SegmentCount(), cfg.Mode)),
						report.Info("trim", "trim-audit", string(auditJSON)),
						report.Info("trim", "validate-output", fmt.Sprintf("validated %d output segment(s)", outputArtifact.SegmentCount())),
						report.Info("output", "json", "wrote transcript JSON"),
					},
				}
				if err := report.WriteJSON(cfg.ReportFile, rpt); err != nil {
					return err
				}
			}

			return nil
		},
	}

	flags := cmd.Flags()
	flags.StringVar(&opts.InputFile, "input-file", "", "input seriatim transcript artifact JSON file")
	flags.StringVar(&opts.OutputFile, "output-file", "", "output transcript JSON file")
	flags.StringVar(&opts.ReportFile, "report-file", "", "optional report JSON file")
	flags.StringVar(&opts.Keep, "keep", "", "segment ID selector to keep (for example: 1-10,15)")
	flags.StringVar(&opts.Remove, "remove", "", "segment ID selector to remove (for example: 1-10,15)")
	flags.StringVar(&opts.OutputSchema, "output-schema", "", "optional output JSON schema override: seriatim-minimal, seriatim-intermediate, or seriatim-full")
	flags.BoolVar(&opts.AllowEmpty, "allow-empty", false, "allow trimming to an empty transcript")

	return cmd
}

func writeOutputJSON(path string, value any) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	enc := json.NewEncoder(file)
	enc.SetIndent("", "  ")
	return enc.Encode(value)
}

func orderedIDMapping(mapping map[int]int) []trimIDMapping {
	keys := make([]int, 0, len(mapping))
	for oldID := range mapping {
		keys = append(keys, oldID)
	}
	sort.Ints(keys)

	pairs := make([]trimIDMapping, 0, len(keys))
	for _, oldID := range keys {
		pairs = append(pairs, trimIDMapping{
			OldID: oldID,
			NewID: mapping[oldID],
		})
	}
	return pairs
}