go-mxl-pattern-generator/cmd/mxl-pattern/main.go

// MXL pattern generator feed rendered on the GPU via wgpu (WebGPU/Vulkan).
// Run from the repo root: make wgpu-gen
// (mixing libmxl cgo with wgpu/goffi needs the internal linker)
package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"sort"
	"strings"
	"syscall"

	"github.com/google/uuid"
	"github.com/qvest-digital/go-mxl/mxl"
	"github.com/spf13/pflag"

	"mxl-pattern-generator/internal/audio"
	"mxl-pattern-generator/internal/flowdef"
	"mxl-pattern-generator/internal/generator"
	"mxl-pattern-generator/internal/video"
)

const (
	APP_NAME = "MXL pattern generator"
	APP_VER  = "0.1.0"
)

type appArgs struct {
	showHelp bool

	domain           string
	videoFlowDefFile string
	audioFlowDefFile string
	pattern          string
	listPatterns     bool

	textOverlay string
	overlayX    int
	overlayY    int
	overlayPos  string

	videoWidth  uint
	videoHeight uint
	videoFPS    string
	videoUUID   string

	audioChannels     uint8
	audioSamplingFreq string
	audioLevel        string
	audioUUID         string
}

type parseResult struct {
	args      appArgs
	shouldRun bool
}

var frameRates = map[string]mxl.Rational{
	"23.97": {Num: 24000, Den: 1001},
	"24":    {Num: 24, Den: 1},
	"25":    {Num: 25, Den: 1},
	"29.97": {Num: 30000, Den: 1001},
	"30":    {Num: 30, Den: 1},
	"50":    {Num: 50, Den: 1},
	"59.94": {Num: 60000, Den: 1001},
	"60":    {Num: 60, Den: 1},
	"120":   {Num: 120, Den: 1},
	"240":   {Num: 240, Den: 1},
}

var samplingRates = map[string]mxl.Rational{
	"44.1": {Num: 44100, Den: 1},
	"48":   {Num: 48000, Den: 1},
	"96":   {Num: 96000, Den: 1},
	"192":  {Num: 192000, Den: 1},
}

func printHelp(w io.Writer, fs *pflag.FlagSet) {
	fmt.Fprintf(w, "%s %s\n", APP_NAME, APP_VER)
	fmt.Fprintln(w, "Usage: mxl-gen -d <domain> [-v <flowDef.json>] [-a <flowDef.json>] [options]")
	fmt.Fprintln(w, "   or: mxl-gen -d <domain> [--width <width px>] [--height <height px>] [--fps <framerate>] \\")
	fmt.Fprintln(w, "   [-c <channels amount>] [-f <sample rate>]")
	fmt.Fprintln(w, "   or: mxl-gen -d <domain> with default params")
	fmt.Fprintln(w, "Video and audio feed parameters are ignored when a flow definition file is provided.")
	fmt.Fprintln(w)
	fs.SetOutput(w)
	fs.PrintDefaults()
}

func printUsage(w io.Writer) {
	fmt.Fprintln(w, "Usage: mxl-gen -d <domain> [-v <flowDef.json>] [-a <flowDef.json>] [options]")
	fmt.Fprintln(w, "Try 'mxl-gen -h' for more information.")
}

func validateDomain(domain string) error {
	if domain == "" {
		return fmt.Errorf("domain is required")
	}
	fi, err := os.Stat(domain)
	if err != nil {
		return fmt.Errorf("invalid MXL domain %q: %w", domain, err)
	}
	if !fi.IsDir() {
		return fmt.Errorf("invalid MXL domain %q: not a directory", domain)
	}
	if ok, err := mxl.IsTmpFs(domain); err != nil || !ok {
		if err != nil {
			return fmt.Errorf("check MXL domain %q: %w", domain, err)
		}
		return fmt.Errorf("invalid MXL domain %q: directory must be on tmpfs", domain)
	}
	return nil
}

func validateFlowDefPath(label, path string) error {
	if path == "" {
		return nil
	}
	fi, err := os.Stat(path)
	if err != nil {
		return fmt.Errorf("%s flow definition %q is not accessible: %w", label, path, err)
	}
	if fi.IsDir() {
		return fmt.Errorf("%s flow definition %q is a directory", label, path)
	}
	return nil
}

func validateVideoArgs(args *appArgs) error {
	if _, exists := patterns[args.pattern]; !exists {
		return fmt.Errorf("unknown video pattern %q (use --list-patterns to see available patterns)", args.pattern)
	}
	if args.videoFlowDefFile != "" {
		return nil
	}
	if args.videoWidth == 0 || args.videoWidth%6 != 0 {
		// v210 stores 6 pixels in each 16-byte block.
		return fmt.Errorf("video width must be greater than zero and divisible by 6")
	}
	if args.videoHeight == 0 {
		return fmt.Errorf("video height must be greater than zero")
	}
	if _, exists := frameRates[args.videoFPS]; !exists {
		return fmt.Errorf("unsupported video FPS %q (supported: %s); use a flow definition for other rates",
			args.videoFPS, sortedMapKeys(frameRates))
	}
	if args.videoUUID == "" {
		args.videoUUID = uuid.NewString()
		return nil
	}
	if _, err := uuid.Parse(args.videoUUID); err != nil {
		return fmt.Errorf("invalid video UUID %q: %w", args.videoUUID, err)
	}
	return nil
}

func validateAudioArgs(args *appArgs) error {
	if args.audioFlowDefFile == "" && args.audioChannels == 0 {
		return nil
	}
	if _, ok := audioLevels[args.audioLevel]; !ok {
		return fmt.Errorf(
			"unsupported audio level %q (supported: %s)",
			args.audioLevel,
			sortedMapKeys(audioLevels),
		)
	}
	if args.audioFlowDefFile != "" {
		return nil
	}
	if _, exists := samplingRates[args.audioSamplingFreq]; !exists {
		return fmt.Errorf("unsupported audio sample rate %q (supported: %s); use a flow definition for other rates",
			args.audioSamplingFreq, sortedMapKeys(samplingRates))
	}
	if args.audioUUID == "" {
		args.audioUUID = uuid.NewString()
		return nil
	}
	if err := uuid.Validate(args.audioUUID); err != nil {
		return fmt.Errorf("invalid audio UUID %q: %w", args.audioUUID, err)
	}
	return nil
}

func validateArgs(args *appArgs) error {
	checks := []func() error{
		func() error { return validateDomain(args.domain) },
		func() error { return validateFlowDefPath("video", args.videoFlowDefFile) },
		func() error { return validateFlowDefPath("audio", args.audioFlowDefFile) },
		func() error { return validateVideoArgs(args) },
		func() error { return validateAudioArgs(args) },
	}
	for _, check := range checks {
		if err := check(); err != nil {
			return err
		}
	}
	return nil
}

func sortedMapKeys[V any](values map[string]V) string {
	keys := make([]string, 0, len(values))
	for key := range values {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	return strings.Join(keys, ", ")
}

type pattern struct {
	name        string
	description string
	kernelPath  string
	motion      bool
}

var patterns = map[string]pattern{
	"ebu75": {
		name:        "ebu75",
		description: "EBU 75% Color Bar Signal",
		kernelPath:  "kernels/static/ebu75.wgsl",
		motion:      false,
	},
	"ebu75-move": {
		name:        "ebu75-move",
		description: "EBU 75% Color Bar Signal with moving square",
		kernelPath:  "kernels/dynamic/ebu75.wgsl",
		motion:      true,
	},
	"ebu100": {
		name:        "ebu100",
		description: "EBU 100% Color Bar Signal",
		kernelPath:  "kernels/static/ebu100.wgsl",
		motion:      false,
	},
	"ebu100-move": {
		name:        "ebu100-move",
		description: "EBU 100% Color Bar Signal with moving square",
		kernelPath:  "kernels/dynamic/ebu100.wgsl",
		motion:      true,
	},
	"smpte": {
		name:        "smpte",
		description: "SMPTE RP-219 Color Bar Signal",
		kernelPath:  "kernels/static/smpteBars.wgsl",
		motion:      false,
	},
	"smpte-move": {
		name:        "smpte-move",
		description: "SMPTE RP-219 Color Bar Signal with moving square",
		kernelPath:  "kernels/dynamic/smpteBars.wgsl",
		motion:      true,
	},
	"gray-bars": {
		name:        "gray-bars",
		description: "13-step grayscale bars (Y 64..940)",
		kernelPath:  "kernels/static/yBars.wgsl",
		motion:      false,
	},
	"gray-bars-move": {
		name:        "gray-bars-move",
		description: "13-step grayscale bars (Y 64..940) with moving square",
		kernelPath:  "kernels/dynamic/yBars.wgsl",
		motion:      true,
	},
	"gray-ramp": {
		name:        "gray-ramp",
		description: "Y gradient (black -> 100% white)",
		kernelPath:  "kernels/static/yRamp.wgsl",
		motion:      false,
	},
	"gray-ramp-move": {
		name:        "gray-ramp-move",
		description: "Y gradient with moving square",
		kernelPath:  "kernels/dynamic/yRamp.wgsl",
		motion:      true,
	},
}

var audioLevels = map[string]float64{
	"ebu":   audio.LevelEBUDBFS,
	"smpte": audio.LevelSMPTEDBFS,
}

func listPatterns(w io.Writer) {
	fmt.Fprintln(w, "List of available video patterns:")
	names := make([]string, 0, len(patterns))
	maxNameWidth := 0
	for name := range patterns {
		names = append(names, name)
		l := len(name)
		if l > maxNameWidth {
			maxNameWidth = l
		}
	}
	sort.Strings(names)
	for _, name := range names {
		p := patterns[name]
		fmt.Fprintf(w, "  %-*s - %s\n", maxNameWidth, name, p.description)
	}
}

func addFlags(fs *pflag.FlagSet, args *appArgs) {
	// common flags
	fs.BoolVarP(&args.showHelp, "help", "h", false, "Show help message and exit")
	// MXL flags
	fs.StringVarP(&args.domain, "domain", "d", "", "MXL domain")
	fs.StringVarP(&args.videoFlowDefFile, "video", "v", "", "Video flow definition JSON file path")
	fs.StringVarP(&args.audioFlowDefFile, "audio", "a", "", "Audio flow definition JSON file path [TODO]")
	// Video pattern flags
	fs.StringVarP(&args.pattern, "pattern", "p", "ebu75", "Video pattern type")
	fs.BoolVar(&args.listPatterns, "list-patterns", false, "List video available video patterns and exit")

	fs.StringVarP(&args.textOverlay, "text", "t", "", "Text overlay above video pattern. Ignored if text-pos set")
	fs.IntVar(&args.overlayX, "text-x", 0, "Text overlay X position in px. Ignored if text-pos set")
	fs.IntVar(
		&args.overlayY,
		"text-y",
		0,
		"Text overlay Y position in px")
	fs.StringVar(
		&args.overlayPos,
		"text-pos",
		"",
		"Text overlay position with pre-defined values:\n"+
			"tl - top-left corner\n"+
			"tc - top-center\n"+
			"tr - top-right corner\n"+
			"cc - center of the frame\n"+
			"bl - bottom-left corner\n"+
			"bc - bottom-center\n"+
			"br - bottom-right corner",
	)

	fs.UintVar(&args.videoWidth, "width", 1920, "Video pattern width. Zero = no video [TODO: 0 width case]")
	fs.UintVar(&args.videoHeight, "height", 1080, "Video pattern height")
	fs.StringVar(&args.videoFPS, "fps", "25", "Video pattern FPS")
	fs.StringVar(&args.videoUUID, "video-id", "", "Video UUID. Will be created, if not provided")
	// Audio pattern flags
	fs.Uint8VarP(&args.audioChannels, "channel", "c", 0, "Amount of audio channels. Each channel: num * 1kHz")
	fs.StringVarP(&args.audioSamplingFreq, "freq", "f", "48", "Sampling frequency of test audio feed in kHz")
	fs.StringVar(&args.audioLevel, "audio-level", "ebu", "Audio alignment level: ebu (-18 dBFS) or smpte (-20 dBFS)")
	fs.StringVar(&args.audioUUID, "audio-id", "", "Audio UUID. Will be created, if not provided")
}

func parseArgs(argv []string, stdout, stderr io.Writer) (parseResult, error) {
	var args appArgs
	flagSet := pflag.NewFlagSet("args", pflag.ContinueOnError)
	flagSet.SortFlags = false
	flagSet.SetOutput(stderr)
	flagSet.Usage = func() { printUsage(stderr) }
	addFlags(flagSet, &args)

	if err := flagSet.Parse(argv); err != nil {
		return parseResult{}, err
	}
	if args.showHelp {
		printHelp(stdout, flagSet)
		return parseResult{args: args}, nil
	}
	if args.listPatterns {
		listPatterns(stdout)
		return parseResult{args: args}, nil
	}
	if flagSet.NArg() != 0 {
		return parseResult{}, fmt.Errorf("unexpected positional arguments: %v", flagSet.Args())
	}
	if err := validateArgs(&args); err != nil {
		return parseResult{}, err
	}
	return parseResult{args: args, shouldRun: true}, nil
}

func buildVideoConfig(args appArgs) (video.Config, error) {
	var definition flowdef.Video
	if args.videoFlowDefFile != "" {
		data, err := os.ReadFile(args.videoFlowDefFile)
		if err != nil {
			return video.Config{}, fmt.Errorf(
				"read video flow definition %q: %w",
				args.videoFlowDefFile,
				err,
			)
		}

		definition, err = flowdef.ParseV210Video(data)
		if err != nil {
			return video.Config{}, fmt.Errorf(
				"parse video flow definition %q: %w",
				args.videoFlowDefFile,
				err,
			)
		}
	} else {
		rate, ok := frameRates[args.videoFPS]
		if !ok {
			return video.Config{}, fmt.Errorf("unsupported video FPS %q", args.videoFPS)
		}

		var err error
		definition, err = flowdef.NewV210Video(
			args.videoUUID,
			args.videoWidth,
			args.videoHeight,
			flowdef.Rational{
				Numerator:   uint(rate.Num),
				Denominator: uint(rate.Den),
			},
		)
		if err != nil {
			return video.Config{}, fmt.Errorf(
				"build video flow definition: %w",
				err,
			)
		}
	}
	if _, ok := patterns[args.pattern]; !ok {
		return video.Config{}, fmt.Errorf("unknown video pattern %q", args.pattern)
	}

	return video.Config{
		Definition: definition,
		Pattern:    args.pattern,
		Overlay: video.OverlayConfig{
			Text:     args.textOverlay,
			X:        args.overlayX,
			Y:        args.overlayY,
			Position: args.overlayPos,
		},
	}, nil
}

// buildAudioConfig returns nil when audio is disabled. An explicit audio flow
// definition enables audio even when --channel is zero.
func buildAudioConfig(args appArgs) (*audio.Config, error) {
	if args.audioFlowDefFile == "" && args.audioChannels == 0 {
		return nil, nil
	}
	levelDBFS, ok := audioLevels[args.audioLevel]
	if !ok {
		return nil, fmt.Errorf("unsupported audio level %q", args.audioLevel)
	}

	var definition flowdef.Audio
	if args.audioFlowDefFile != "" {
		data, err := os.ReadFile(args.audioFlowDefFile)
		if err != nil {
			return nil, fmt.Errorf("read audio flow definition %q: %w", args.audioFlowDefFile, err)
		}
		definition, err = flowdef.ParseFloat32Audio(data)
		if err != nil {
			return nil, fmt.Errorf("parse audio flow definition %q: %w", args.audioFlowDefFile, err)
		}
	} else {
		rate, ok := samplingRates[args.audioSamplingFreq]
		if !ok {
			return nil, fmt.Errorf("unsupported audio sample rate %q", args.audioSamplingFreq)
		}
		var err error
		definition, err = flowdef.NewFloat32Audio(
			args.audioUUID,
			uint(args.audioChannels),
			flowdef.Rational{Numerator: uint(rate.Num), Denominator: uint(rate.Den)},
		)
		if err != nil {
			return nil, fmt.Errorf("build audio flow definition: %w", err)
		}
	}

	return &audio.Config{
		Definition: definition,
		LevelDBFS:  levelDBFS,
	}, nil
}

func main() {
	parsed, err := parseArgs(os.Args[1:], os.Stdout, os.Stderr)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		printUsage(os.Stderr)
		os.Exit(2)
	}
	if !parsed.shouldRun {
		return
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	if err := run(ctx, parsed.args); err != nil {
		log.Printf("%s: %v", APP_NAME, err)
		os.Exit(1)
	}
}

func run(ctx context.Context, args appArgs) (runErr error) {
	videoCfg, err := buildVideoConfig(args)
	if err != nil {
		return fmt.Errorf("video configuration: %w", err)
	}

	log.Printf("%s %s", APP_NAME, APP_VER)
	log.Printf("Domain: %s", args.domain)
	log.Printf("Video: %dx%d %d/%d",
		videoCfg.Width(), videoCfg.Height(), videoCfg.Rate().Num, videoCfg.Rate().Den)
	log.Printf("Video ID: %s", videoCfg.ID())

	inst, err := mxl.NewInstance(args.domain, "")
	if err != nil {
		return fmt.Errorf("initialize MXL domain %q: %w", args.domain, err)
	}
	defer func() {
		if err := inst.Close(); err != nil {
			runErr = errors.Join(runErr, fmt.Errorf("close MXL instance: %w", err))
		}
	}()

	return runVideo(ctx, inst, videoCfg)
}

func runVideo(ctx context.Context, inst *mxl.Instance, cfg video.Config) (runErr error) {
	pattern, ok := patterns[cfg.Pattern]
	if !ok {
		return fmt.Errorf("unknown video pattern %q", cfg.Pattern)
	}

	// TODO: fall back to a CPU generator if GPU initialization fails.
	gen, err := generator.NewWGPUGenerator(cfg.Width(), cfg.Height(), pattern.kernelPath)
	if err != nil {
		return fmt.Errorf("initialize wgpu video generator: %w", err)
	}
	defer func() {
		if err := gen.Close(); err != nil {
			runErr = errors.Join(runErr, fmt.Errorf("close video generator: %w", err))
		}
	}()

	overlay, err := buildTextOverlay(cfg)
	if err != nil {
		return err
	}

	flowJSON, err := json.Marshal(cfg.Definition)
	if err != nil {
		return fmt.Errorf("marshal video flow definition: %w", err)
	}
	writer, isCreated, err := inst.NewWriter(string(flowJSON))
	if err != nil {
		return fmt.Errorf("create video writer: %w", err)
	}
	defer func() {
		if err := writer.Close(); err != nil {
			runErr = errors.Join(runErr, fmt.Errorf("close video writer: %w", err))
		}
	}()
	if !isCreated {
		log.Printf("reusing existing video flow: %s", cfg.ID())
	}

	rate := writer.Config().Common.GrainRate
	idx := mxl.CurrentIndex(rate)
	log.Printf("writing video flow grainRate=%d/%d starting at idx=%d", rate.Num, rate.Den, idx)

	var grainsWritten int64
	var tick uint32 // Small animation counter; shaders cannot use the full grain index yet.
	for {
		select {
		case <-ctx.Done():
			log.Printf("stopping video after %d grains", grainsWritten)
			return nil
		default:
		}

		grain, err := writer.OpenGrain(idx)
		if err != nil {
			return fmt.Errorf("open video grain %d: %w", idx, err)
		}
		if err := gen.GenerateFrame(grain.Payload, int(tick)); err != nil {
			return cancelVideoGrain(grain, fmt.Errorf("generate frame for grain %d: %w", idx, err))
		}
		if overlay != nil {
			if err := overlay.ApplyV210(grain.Payload); err != nil {
				return cancelVideoGrain(grain, fmt.Errorf("apply text overlay to grain %d: %w", idx, err))
			}
		}
		if err := grain.Commit(grain.TotalSlices, 0); err != nil {
			return fmt.Errorf("commit video grain %d: %w", idx, err)
		}

		grainsWritten++
		idx++
		tick++
		if grainsWritten%100 == 0 {
			log.Printf("video grains written=%d, next index=%d", grainsWritten, idx)
		}
		mxl.SleepNs(mxl.NsUntilIndex(idx, rate))
	}
}

func buildTextOverlay(cfg video.Config) (overlay *generator.TextOverlay, resultErr error) {
	if cfg.Overlay.Text == "" {
		return nil, nil
	}
	if cfg.Overlay.X < 0 || cfg.Overlay.X > int(cfg.Width()) ||
		cfg.Overlay.Y < 0 || cfg.Overlay.Y > int(cfg.Height()) {
		return nil, fmt.Errorf("text overlay position (%d, %d) is outside the %dx%d video frame",
			cfg.Overlay.X, cfg.Overlay.Y, cfg.Width(), cfg.Height())
	}

	face, err := generator.LoadFace("assets/fonts/JetBrainsMonoNLNerdFontMono-Regular.ttf", 48)
	if err != nil {
		return nil, fmt.Errorf("load text overlay font: %w", err)
	}
	defer func() {
		if err := face.Close(); err != nil {
			resultErr = errors.Join(resultErr, fmt.Errorf("close text overlay font: %w", err))
		}
	}()

	overlay, err = generator.NewTextOverlay(
		cfg.Overlay.Text,
		int(cfg.Width()),
		int(cfg.Height()),
		cfg.Overlay.X,
		cfg.Overlay.Y,
		cfg.Overlay.Position,
		face,
	)
	if err != nil {
		return nil, fmt.Errorf("create text overlay: %w", err)
	}
	return overlay, nil
}

func cancelVideoGrain(grain *mxl.GrainWriteAccess, cause error) error {
	if err := grain.Cancel(); err != nil {
		return errors.Join(cause, fmt.Errorf("cancel video grain: %w", err))
	}
	return cause
}