// 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 (
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"

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

	flowdef "mxl-pattern-generator/internal/flow-def"
	"mxl-pattern-generator/internal/generator"
)

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

	videoWidth  uint
	videoHeight uint
	videoFPS    string
	videoUUID   string

	audioChannels     uint8
	audioSamplingFreq string
	audioUUID         string
}

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(fs *pflag.FlagSet) {
	fmt.Printf("%s %s\n", APP_NAME, APP_VER)
	fmt.Println("Usage: mxl-gen -d <domain> [-v <flowDef.json>] [-a <flowDef.json>] [options]")
	fmt.Println("   or: mxl-gen -d <domain> [--with <width px>] [--height <height px>] [--fps <framerate>] \\")
	fmt.Println("   [-c <channels amount>] [-f <sample rate>]")
	fmt.Println("   or: mxl-gen -d <domain> with default params")
	fmt.Println("Video and audio feeds params will be ignored, if flow definition file provided.")
	fmt.Println()
	fs.PrintDefaults()
}

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

func checkArgs(args appArgs) {
	printUsageAndExit := func() {
		printUsage()
		os.Exit(2)
	}
	// domain
	if args.domain == "" {
		fmt.Fprintf(os.Stderr, "Domain is required\n")
		printUsageAndExit()
	}
	fi, err := os.Stat(args.domain)
	if err != nil || !fi.IsDir() {
		fmt.Fprintf(os.Stderr, "Invalid MXL domain: %s\n", args.domain)
		printUsageAndExit()
	}
	if ok, err := mxl.IsTmpFs(args.domain); err != nil || !ok {
		fmt.Fprintf(os.Stderr, "Invalid MXL domain: %s\n", args.domain)
		fmt.Fprintln(os.Stderr, "Domain must be directory in tmps.")
		printUsageAndExit()
	}
	// FlowDef
	checkFlowDef := func(label, flowDef string) {
		fi, err := os.Stat(flowDef)
		if err != nil || fi.IsDir() {
			fmt.Fprintf(os.Stderr, "%s flow definition .json file is not accesible\n", label)
			printUsageAndExit()
		}
	}
	videoFlowDefProvided, audioFlowDefProvided := false, false
	if args.videoFlowDefFile != "" {
		checkFlowDef("Video", args.videoFlowDefFile)
		videoFlowDefProvided = true
	}
	if args.audioFlowDefFile != "" {
		checkFlowDef("Audio", args.audioFlowDefFile)
		audioFlowDefProvided = true
	}

	if !videoFlowDefProvided {
		if args.videoWidth == 0 || args.videoWidth%6 != 0 {
			// width%6 == 0 - because of v210 (6 pixels per 16-byte block)
			fmt.Fprintf(os.Stderr, "Video width must be > 0 and divisible by 6\n")
			printUsageAndExit()
		}
		if args.videoHeight == 0 {
			fmt.Fprintf(os.Stderr, "Video height must be > 0\n")
			printUsageAndExit()
		}
		if _, exists := frameRates[args.videoFPS]; !exists {
			fmt.Fprintf(os.Stderr, "FPS %s is not in available list.\n", args.videoFPS)
			fmt.Fprintln(os.Stderr, "If you need more complex solution, use flow definition .json instead.")
			fmt.Fprintln(os.Stderr, "Available list:")
			for key, _ := range frameRates {
				fmt.Fprintf(os.Stderr, "  %s\n", key)
			}
			printUsageAndExit()
		}
		if args.videoUUID != "" {
			if _, err := uuid.Parse(args.videoUUID); err != nil {
				fmt.Fprintf(os.Stderr, "Video UUID %s is not valid.\n", args.videoUUID)
				printUsageAndExit()
			}
		} else {
			args.videoUUID = uuid.NewString()
		}
		if args.pattern != "" {
			if _, exists := patterns[args.pattern]; !exists {
				fmt.Fprintf(os.Stderr, "Pattern %s is not in available list.\n", args.pattern)
				listPatterns(os.Stderr)
				os.Exit(2)
			}
		}
	}

	if !audioFlowDefProvided {
		if args.audioChannels == 0 {
			fmt.Fprintln(os.Stderr, "Audio channels amount must be > 0")
			printUsageAndExit()
		}
		if _, exists := samplingRates[args.audioSamplingFreq]; !exists {
			fmt.Fprintf(os.Stderr, "Sample rate %s is not in available list.\n", args.audioSamplingFreq)
			fmt.Fprintln(os.Stderr, "If you need more complex solution, use flow definition .json instead.")
			fmt.Fprintln(os.Stderr, "Available list:")
			for key, _ := range samplingRates {
				fmt.Fprintf(os.Stderr, "  %s\n", key)
			}
			printUsageAndExit()
		}
		if args.audioUUID != "" {
			if err := uuid.Validate(args.audioUUID); err != nil {
				fmt.Fprintf(os.Stderr, "Audio UUID %s is not valid.\n", args.audioUUID)
				printUsageAndExit()
			}
		} else {
			args.audioUUID = uuid.NewString()
		}
	}
}

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

var patterns = map[string]pattern{
	"bars": {
		name:        "bars",
		description: "SMPTE 75% color bars",
		kernelPath:  "kernels/static/smpteBars.wgsl",
		motion:      false,
	},
	"bars-move": {
		name:        "bars-move",
		description: "SMPTE 75% color bars with moving square",
		kernelPath:  "kernels/dynamic/smpteBars.wgsl",
		motion:      true,
	},
	"gray-ramp": {
		name:        "gray-ramp",
		description: "Y gradient",
		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:      false,
	},
	"dvd": {
		name:        "dvd",
		description: "DVD logo pattern, but with MXL logo",
		kernelPath:  "kernels/dynamic/dvdLogo.wgsl",
		motion:      true,
	},
}

func listPatterns(f *os.File) {
	fmt.Fprintln(f, "List of available video patterns:")
	var maxNameWidth int = 0
	for name, _ := range patterns {
		l := len(name)
		if l > maxNameWidth {
			maxNameWidth = l
		}
	}
	for name, p := range patterns {
		fmt.Fprintf(f, "  %-*s - %s\n", maxNameWidth, name, p.description)
	}
}

func flagSetAddFlags(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", "bars", "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")
	fs.UintVar(&args.videoWidth, "width", 1920, "Video pattern width")
	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 [TODO]")
	// Audio pattern flags
	fs.Uint8VarP(&args.audioChannels, "channel", "c", 2, "Amount of audio channels. Each channel: num * 1kHz [TODO]")
	fs.StringVarP(&args.audioSamplingFreq, "freq", "f", "48", "Sampling frequency of test audio feed in kHz [TODO]")
	fs.StringVar(&args.audioUUID, "audio-id", "", "Audio UUID. Will be created, if not provided [TODO]")
}

func main() {
	var args appArgs
	flagSet := pflag.NewFlagSet("args", pflag.ContinueOnError)
	flagSet.SortFlags = false
	flagSet.Usage = func() { printUsage() }
	flagSetAddFlags(flagSet, &args)

	if err := flagSet.Parse(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, err)
		printUsage()
		os.Exit(2)
	}
	if args.showHelp {
		printHelp(flagSet)
		return
	}
	if args.listPatterns {
		listPatterns(os.Stderr)
		return
	}
	checkArgs(args)

	var mxlDomain string = "/dev/shm/mxl"
	type videoInfo struct {
		uuid   string
		width  uint
		height uint
		fps    mxl.Rational
	}
	var vi videoInfo
	var videoFlowDef string
	if args.videoFlowDefFile == "" {
		args.videoUUID = "8f1d2a4b-6c3e-4f5a-9b2c-1d7e8a3f0b5d" // TODO: remove before public release
		vi = videoInfo{
			uuid:   args.videoUUID,
			width:  args.videoWidth,
			height: args.videoHeight,
			fps:    frameRates[args.videoFPS],
		}
		flowDef, err := flowdef.NewFlowDefJSON(
			flowdef.TYPE_VIDEO,
			vi.uuid,
			vi.width,
			vi.height,
			uint(vi.fps.Num),
			uint(vi.fps.Den),
		)
		videoFlowDef = flowDef
		if err != nil {
			log.Fatalf("Could not create Flow Definition: %v", err)
		}
	} else {
		flowDef, err := flowdef.ReadFlowDefFile(args.videoFlowDefFile)
		if err != nil {
			log.Fatalf("Could not read video flow def .json: %s. Reason: %v", args.videoFlowDefFile, err)
		}
		videoFlowDef = flowDef
	}

	log.Printf("%s %s", APP_NAME, APP_VER)
	log.Printf("Domain: %s", mxlDomain)
	log.Printf("Video: %dx%d %d/%d", vi.width, vi.height, vi.fps.Num, vi.fps.Den)
	log.Printf("Video UUID: %s", vi.uuid)

	// TODO: if init failed -> CPU generator
	videoPattern := patterns[args.pattern]
	var genOpts []generator.WGPUOption
	if videoPattern.name == "dvd" {
		plane, err := generator.LoadLogoPNG("assets/mxl.png")
		if err != nil {
			log.Fatalf("logo load failed: %v", err)
		}
		genOpts = append(genOpts, generator.WithLogo(plane))
	}
	gen, err := generator.NewWGPUGenerator(vi.width, vi.height, videoPattern.kernelPath, genOpts...)
	if err != nil {
		log.Fatalf("wgpu init failed: %v", err)
	}
	defer gen.Close()

	// Static text overlay: rasterized + pre-packed once, stamped on each
	// frame after the GPU render (microseconds per frame).
	var overlay *generator.TextOverlay
	if args.textOverlay != "" {
		face, err := generator.LoadFace("assets/fonts/JetBrainsMonoNLNerdFontMono-Regular.ttf", 48)
		if err != nil {
			log.Fatalf("text overlay init failed: %v", err)
		}
		defer face.Close()
		overlay, err = generator.NewTextOverlay(args.textOverlay, int(vi.width), int(vi.height), face)
		if err != nil {
			log.Fatalf("text overlay init failed: %v", err)
		}
	}

	inst, err := mxl.NewInstance(mxlDomain, "")
	if err != nil {
		log.Fatalf("MXL Init Failed: %v", err)
	}
	defer inst.Close()

	writer, isCreated, err := inst.NewWriter(videoFlowDef)
	if err != nil {
		log.Fatalf("Failed to create MXL writer: %v", err)
	}
	if !isCreated {
		log.Printf("reusing existing flow: %s, domain: %s", vi.uuid, mxlDomain)
	}
	defer writer.Close()

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

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)

	// core loop
	var grainsWritten int64
	// animation clock: small counter, not the huge grain index.
	// Reason: current wgpu shaders limitations
	var tick uint32
	for {
		select {
		case <-stop:
			log.Printf("stopping after %d grains", grainsWritten)
			return
		default:
		}

		gwa, err := writer.OpenGrain(idx)
		if err != nil {
			log.Fatalf("OpenGrain(%d): %v", idx, err)
		}
		if err := gen.GenerateFrame(gwa.Payload, int(tick)); err != nil {
			log.Fatalf("GenerateFrame(%d): %v", idx, err)
		}
		if overlay != nil {
			if err := overlay.ApplyV210(gwa.Payload); err != nil {
				log.Fatalf("text overlay: %v", err)
			}
		}
		if err := gwa.Commit(gwa.TotalSlices, 0); err != nil {
			log.Fatalf("Commit(%d): %v", idx, err)
		}

		grainsWritten++
		idx++
		tick++
		if grainsWritten%100 == 0 {
			log.Printf("grains written=%d, index=%d", grainsWritten, idx)
		}
		// Pace ourselves to roughly the grain rate
		mxl.SleepNs(mxl.NsUntilIndex(idx, rate))
	}
}