// 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
	overlayX    int
	overlayY    int

	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)
		fmt.Fprintf(os.Stderr, "Reason: %v\n", err)
		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{
	"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:      false,
	},
	"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:      false,
	},
	"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:      false,
	},
	"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:      false,
	},
	"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:      false,
	},
}

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.IntVar(&args.overlayX, "text-x", 0, "Text overlay X position in px")
	fs.IntVar(&args.overlayY, "text-y", 0, "Text overlay Y position in px")

	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")
	// 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]
	gen, err := generator.NewWGPUGenerator(vi.width, vi.height, videoPattern.kernelPath)
	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()
		if args.overlayX < 0 ||
			args.overlayX > int(vi.width) ||
			args.overlayY < 0 ||
			args.overlayY > int(vi.height) {
			log.Fatalf("Overlay X/Y pos can't be negative or greater, than video width/height")
		}
		overlay, err = generator.NewTextOverlay(
			args.textOverlay,
			int(vi.width),
			int(vi.height),
			args.overlayX,
			args.overlayY,
			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))
	}
}