package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"mxl-player/internal/sdl"
	"mxl-player/internal/source"
	"runtime"
	"time"
	"unsafe"

	vk "github.com/christerso/vulkan-go/vk"
	"github.com/qvest-digital/go-mxl/mxl"
)

const (
	APP_NAME         = "MXL Player"
	APP_VER          = "0.0.1"
	WIN_WIDTH  int32 = 1280
	WIN_HEIGHT int32 = 720
)

func main() {
	var errMinimized = errors.New("Window minimized")
	// TODO: remove flags defaults
	mxlDomain := flag.String("d", "/dev/shm/mxl", "MXL domain")
	mxlVideoFlowID := flag.String("v", "5fbec3b1-1b0f-417d-9059-8b94a47197ed", "MXL video flow UUID")
	flag.Parse()

	runtime.LockOSThread()
	if err := sdl.Load(); err != nil {
		panic(err)
	}
	if !sdl.Init(sdl.InitVideo) {
		log.Fatalf("SDL_Init: %s", sdl.GetError())
		return
	}

	windowHandler := sdl.CreateWindow(fmt.Sprintf("%s %s", APP_NAME, APP_VER), WIN_WIDTH, WIN_HEIGHT, sdl.WindowVulkan|sdl.WindowResizable)
	if windowHandler == 0 {
		sdl.Quit()
		log.Fatalf("SDL_CreateWindow: %s", sdl.GetError())
		return
	}

	// Vulkan init
	if err := vk.Load(); err != nil {
		panic(err)
	}

	// Vulkan instance extensions SDL needs.
	sdlExtensions := sdl.VulkanGetInstanceExtensions()
	if len(sdlExtensions) == 0 {
		log.Fatal("sdlVulkanGetInstanceExtensions is 0")
		return
	}
	// Vulkan Instance creation
	var vkLayers []string
	vkExtensions := append(sdlExtensions, vk.ExtDebugUtils)
	vkInstance, err := vk.CreateInstance(vk.InstanceConfig{
		ApplicationName: APP_NAME,
		EngineName:      "no engine",
		Extensions:      vkExtensions,
		Layers:          vkLayers,
	})
	if err != nil {
		log.Fatalf(err.Error())
		panic(err)
	}
	defer vkInstance.Destroy()
	// Vulkan surface
	var vkSurface uint64
	if !sdl.VulkanCreateSurface(windowHandler, uintptr(vkInstance), 0, &vkSurface) {
		log.Fatalf("SDL_Vulkan_CreateSurface: %s", sdl.GetError())
		return
	}
	vkSurf := vk.SurfaceKHR(vkSurface)
	defer vkInstance.DestroySurface(vkSurf)

	// Vulkan Physical Device
	devices, err := vkInstance.EnumeratePhysicalDevices()
	if err != nil || len(devices) == 0 {
		panic("No Vulkan devices")
	}
	for _, pd := range devices {
		info := pd.Info()
		fmt.Printf("%s (%s)\n", info.Name, info.Type)
	}
	vkPhysDevice := devices[0]
	// Vulkan Surface & Graphics Queue
	gfx, err := vkPhysDevice.GraphicsFamily()
	if err != nil {
		panic(err.Error())
	}
	if !vkPhysDevice.SurfaceSupport(gfx, vkSurf) {
		log.Fatalf("graphics queue cannot present")
		return
	}
	vkDevice, vkQueue, err := vkPhysDevice.CreateDevice(vk.DeviceConfig{
		GraphicsFamily: gfx,
		Extensions:     []string{"VK_KHR_swapchain"},
	})
	if err != nil {
		panic(err)
	}
	defer vkDevice.Destroy()
	// TODO: colorspace & render pass
	var vkFormat vk.Format = vk.FormatUndefined
	var vkColorSpace uint32
	formats, _ := vkPhysDevice.SurfaceFormats(vkSurf)
	for _, f := range formats {
		// Prefer 10-bit RGB (A2B10G10R10_UNORM = 64) to preserve V210's 10 bits
		if f.Format == vk.Format(64) && f.ColorSpace == vk.ColorSpaceSRGBNonlinear {
			vkFormat = f.Format
			vkColorSpace = f.ColorSpace
		}
	}
	if vkFormat == vk.FormatUndefined {
		vkFormat = formats[0].Format
		vkColorSpace = formats[0].ColorSpace
	}

	var vkPresentMode vk.PresentMode = vk.PresentModeFIFO
	modes, _ := vkPhysDevice.SurfacePresentModes(vkSurf)
	for _, m := range modes {
		if m == vk.PresentModeMailbox {
			vkPresentMode = m
		}
	}

	vkRenderPass, err := vkDevice.CreateColorDepthRenderPass(vkFormat, vk.FormatD32Sfloat)
	if err != nil {
		panic(err.Error())
	}
	defer vkDevice.DestroyRenderPass(vkRenderPass)

	// Vulkan Command Pool
	vkCommandPool, err := vkDevice.CreateCommandPool(gfx)
	if err != nil {
		panic(err.Error())
	}
	defer vkDevice.DestroyCommandPool(vkCommandPool)
	vkCommands, err := vkDevice.AllocateCommandBuffers(vkCommandPool, 1)
	if err != nil {
		panic(err.Error())
	}

	// MXL Source
	mxlSrc, err := source.Open(*mxlDomain, *mxlVideoFlowID)
	if err != nil {
		log.Fatalf("source: %v\n", err)
	}
	defer func() { _ = mxlSrc.Close() }()
	frameSize := vk.DeviceSize(mxlSrc.Stride()) * vk.DeviceSize(mxlSrc.Height())
	fmt.Printf("source: %dx%d stride=%d frameSize=%d\n",
		mxlSrc.Width(), mxlSrc.Height(), mxlSrc.Stride(), frameSize)
	// Staging buffer: host-visible, persistently mapped. The reader goroutine
	// writes V210 bytes here; the GPU copies from it.
	staging, err := vkDevice.CreateBuffer(vkPhysDevice, vk.BufferConfig{
		Size:       frameSize,
		Usage:      vk.BufferUsageTransferSrc,
		Properties: vk.MemoryHostVisible | vk.MemoryHostCoherent,
		Map:        true,
	})
	if err != nil {
		panic(err)
	}
	defer vkDevice.DestroyBuffer(staging)
	// Device-local V210 buffer: fast for the GPU to read (M3 compute), CPU can't
	// write it. Filled each frame by a CopyBuffer from staging.
	v210Buf, err := vkDevice.CreateBuffer(vkPhysDevice, vk.BufferConfig{
		Size:       frameSize,
		Usage:      vk.BufferUsageTransferDst | vk.BufferUsageStorageBuffer,
		Properties: vk.MemoryDeviceLocal,
		Map:        false,
	})
	if err != nil {
		panic(err)
	}
	defer vkDevice.DestroyBuffer(v210Buf)

	// Decode pipeline: fullscreen triangle, fragment reads V210 from the
	// staging buffer and writes 10-bit RGB to the swapchain color attachment.
	vertModule, err := vkDevice.CreateShaderModule(vertSPV)
	if err != nil {
		panic(err)
	}
	defer vkDevice.DestroyShaderModule(vertModule)

	fragModule, err := vkDevice.CreateShaderModule(fragSPV)
	if err != nil {
		panic(err)
	}
	defer vkDevice.DestroyShaderModule(fragModule)
	// Descriptor set layout: binding 0 = storage buffer (V210), fragment stage.
	decodeDSL, err := vkDevice.CreateDescriptorSetLayout([]vk.DescriptorBinding{
		{Binding: 0, Type: vk.DescriptorStorageBuffer, Count: 1, Stages: vk.ShaderStageFragment},
	})
	if err != nil {
		panic(err)
	}
	defer vkDevice.DestroyDescriptorSetLayout(decodeDSL)
	// Pipeline layout: the set layout + push constants {width,height,strideBytes}.
	decodeLayout, err := vkDevice.CreatePipelineLayout([]vk.DescriptorSetLayout{decodeDSL}, vk.ShaderStageFragment, 12)
	if err != nil {
		panic(err)
	}
	defer vkDevice.DestroyPipelineLayout(decodeLayout)
	decodePipeline, err := vkDevice.CreateGraphicsPipeline(vk.GraphicsPipelineConfig{
		Layout:       decodeLayout,
		RenderPass:   vkRenderPass,
		VertexShader: vertModule,
		FragShader:   fragModule,
		Topology:     vk.TopologyTriangleList,
		PolygonMode:  vk.PolygonFill,
		CullMode:     vk.CullNone,
		FrontFace:    vk.FrontFaceCounterClockwise,
	})
	if err != nil {
		panic(err)
	}
	defer vkDevice.DestroyPipeline(decodePipeline)
	// Descriptor pool + set, bound once to the staging buffer
	descPool, err := vkDevice.CreateDescriptorPool(1, map[vk.DescriptorType]uint32{
		vk.DescriptorStorageBuffer: 1,
	})
	if err != nil {
		panic(err)
	}
	defer vkDevice.DestroyDescriptorPool(descPool)

	decodeSet, err := vkDevice.AllocateDescriptorSet(descPool, decodeDSL)
	if err != nil {
		panic(err)
	}
	vkDevice.UpdateBufferDescriptor(decodeSet, 0, vk.DescriptorStorageBuffer, staging.Buffer, 0, vk.WholeSize)

	// Swapchain
	var (
		vkExtent    vk.Extent2D
		vkSwapchain vk.SwapchainKHR
		swapImages  []vk.Image
		swapViews   []vk.ImageView
		depthImg    vk.AllocImage
		depthView   vk.ImageView
		fbs         []vk.Framebuffer
	)
	destroySwapChain := func() {
		for _, fb := range fbs {
			vkDevice.DestroyFramebuffer(fb)
		}
		vkDevice.DestroyImageView(depthView)
		vkDevice.DestroyImage(depthImg)
		for _, v := range swapViews {
			vkDevice.DestroyImageView(v)
		}
		vkDevice.DestroySwapchain(vkSwapchain)
		fbs = nil
		swapViews = nil
		swapImages = nil
		depthView = 0
		depthImg = vk.AllocImage{}
		vkSwapchain = 0
		vkExtent = vk.Extent2D{}
	}
	createSwapChain := func() (err error) {
		defer func() {
			if err != nil {
				destroySwapChain()
			}
		}()
		caps, err := vkPhysDevice.SurfaceCapabilities(vkSurf)
		if err != nil {
			return err
		}
		//compute swapchain extent
		var w, h int32 = WIN_WIDTH, WIN_HEIGHT
		if !sdl.GetWindowSizeInPixels(windowHandler, &w, &h) {
			return fmt.Errorf("SDL_GetWindowSizeInPixels: %s", sdl.GetError())
		}
		if w < int32(caps.MinImageExtent.Width) {
			w = int32(caps.MinImageExtent.Width)
		}
		if h < int32(caps.MinImageExtent.Height) {
			h = int32(caps.MinImageExtent.Height)
		}
		if w > int32(caps.MaxImageExtent.Width) {
			w = int32(caps.MaxImageExtent.Width)
		}
		if h > int32(caps.MaxImageExtent.Height) {
			h = int32(caps.MaxImageExtent.Height)
		}
		if w == 0 || h == 0 {
			return errMinimized
		}
		vkExtent = vk.Extent2D{Width: uint32(w), Height: uint32(h)}

		vkSwapchain, err = vkDevice.CreateSwapchain(vk.SwapchainConfig{
			Surface:       vkSurf,
			MinImageCount: caps.MinImageCount,
			Format:        vkFormat,
			ColorSpace:    vkColorSpace,
			Extent:        vkExtent,
			PresentMode:   vkPresentMode,
			PreTransform:  caps.CurrentTransform,
		})
		if err != nil {
			return err
		}
		swapImages, err = vkDevice.SwapchainImages(vkSwapchain)
		if err != nil {
			return err
		}
		swapViews = make([]vk.ImageView, len(swapImages))
		for i, img := range swapImages {
			v, err := vkDevice.CreateImageView(img, vkFormat, vk.AspectColor)
			if err != nil {
				return err
			}
			swapViews[i] = v
		}
		depthImg, err = vkDevice.CreateImage2D(
			vkPhysDevice, vk.FormatD32Sfloat, vkExtent, vk.ImageUsageDepthStencilAttachment)
		if err != nil {
			return err
		}
		depthView, err = vkDevice.CreateImageView(depthImg.Image, vk.FormatD32Sfloat, vk.AspectDepth)
		if err != nil {
			return err
		}
		fbs = make([]vk.Framebuffer, len(swapViews))
		for i, cv := range swapViews {
			fb, err := vkDevice.CreateFramebuffer(vkRenderPass, []vk.ImageView{cv, depthView}, vkExtent)
			if err != nil {
				return err
			}
			fbs[i] = fb
		}
		return nil
	}
	recreateSwapChain := func() error {
		if err := vkDevice.WaitIdle(); err != nil {
			return err
		}
		destroySwapChain()
		return createSwapChain()
	}

	if err := createSwapChain(); err != nil {
		panic(err)
	}
	defer destroySwapChain()

	// syncs
	// Semaphores sync GPU work to GPU work (you never wait on these from the CPU). You need two:
	// - imageAvailable — AcquireNextImage signals it when a swapchain image is ready to render into.
	// Your draw submission waits on it before touching color.
	// - renderFinished — your draw submission signals it when rendering is done.
	// Present waits on it before showing the frame.
	// Fences sync GPU work to CPU. You need one:
	// - inFlight — your submission signals it when the GPU is done with the command buffer.
	// Your CPU waits on it before re-recording that same command buffer.
	// It's created signaled so the very first frame's wait returns immediately
	// (otherwise you'd deadlock waiting for work that hasn't been submitted yet).
	imageAvailable, err := vkDevice.CreateSemaphore()
	if err != nil {
		panic(err.Error())
	}
	defer vkDevice.DestroySemaphore(imageAvailable)

	renderFinished, err := vkDevice.CreateSemaphore()
	if err != nil {
		panic(err.Error())
	}
	defer vkDevice.DestroySemaphore(renderFinished)

	inFlight, err := vkDevice.CreateFence(true)
	if err != nil {
		panic(err.Error())
	}
	defer vkDevice.DestroyFence(inFlight)
	defer vkDevice.WaitIdle()

	// Reader goroutine: stages V210 bytes into the staging buffer.
	// Single-flight handshake: it only writes staging after the render thread
	// grants permission (post-WaitFence), so the GPU is never reading it
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	grant := make(chan struct{}, 1)
	staged := make(chan uint64)
	failed := make(chan struct{})

	reopen := func() error {
		if err := mxlSrc.Close(); err != nil {
			return fmt.Errorf("close old source: %w", err)
		}
		for {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
			}
			s, err := source.Open(*mxlDomain, *mxlVideoFlowID)
			if err == nil {
				newSize := vk.DeviceSize(s.Stride()) * vk.DeviceSize(s.Height())
				if newSize != frameSize {
					if err := vkDevice.WaitIdle(); err != nil {
						return err
					}
					vkDevice.DestroyBuffer(v210Buf)
					vkDevice.DestroyBuffer(staging)
					staging, err = vkDevice.CreateBuffer(vkPhysDevice, vk.BufferConfig{
						Size: newSize, Usage: vk.BufferUsageTransferSrc,
						Properties: vk.MemoryHostVisible | vk.MemoryHostCoherent, Map: true,
					})
					if err != nil {
						return err
					}
					v210Buf, err = vkDevice.CreateBuffer(vkPhysDevice, vk.BufferConfig{
						Size: newSize, Usage: vk.BufferUsageTransferDst | vk.BufferUsageStorageBuffer,
						Properties: vk.MemoryDeviceLocal, Map: false,
					})
					if err != nil {
						return err
					}
					frameSize = newSize
					vkDevice.UpdateBufferDescriptor(decodeSet, 0, vk.DescriptorStorageBuffer, staging.Buffer, 0, vk.WholeSize)
					fmt.Printf("source: resolution changed, frameSize=%d\n", newSize)
				}
				mxlSrc = s
				return nil
			}
			log.Printf("source: reopen retry: %v", err)
			time.Sleep(500 * time.Millisecond)
		}
	}

	go func() {
		for {
			select {
			case <-grant:
			case <-ctx.Done():
				return
			}
			f, err := mxlSrc.NextCtx(ctx, 200*time.Millisecond)
			if err != nil {
				if errors.Is(err, context.Canceled) {
					return
				}
				if errors.Is(err, mxl.ErrFlowInvalid) {
					log.Printf("source: flow invalid, reopening")
					if rerr := reopen(); rerr != nil {
						log.Printf("source: reopen failed: %v", rerr)
						cancel()
						return
					}
					select {
					case failed <- struct{}{}:
					case <-ctx.Done():
						return
					}
					continue
				}
				log.Printf("source: %v", err)
				cancel()
				return
			}
			// Copy the borrowed payload into staging BEFORE the next read
			// invalidates it. Within the grain's valid lifetime
			vk.CopyToMapped(staging.Mapped, f.Payload)
			// fmt.Println("staged", f.Index)
			select {
			case staged <- f.Index:
			case <-ctx.Done():
				return
			}
		}
	}()

	type decodePushConstants struct {
		Width, Height, StrideBytes, WinW, WinH uint32
	}

	// Core Loop
	running := true
	resized := false
	granted := false
	fullscreen := false
	var (
		lastIndex  uint64
		dropped    uint64
		frameCount uint64
		lastReport time.Time
	)
	for running {
		frameStart := time.Now()
		var event [128]byte
		for sdl.PollEvent(unsafe.Pointer(&event[0])) {
			eventType := *(*uint32)(unsafe.Pointer(&event[0]))
			switch eventType {
			case sdl.EventQuit:
				running = false
			case sdl.EventWindowResized, sdl.EventPixelSizeChanged:
				resized = true
			case sdl.EventKeyDown:
				key := *(*int32)(unsafe.Pointer(&event[28]))
				switch uint32(key) {
				case sdl.KeyEscape:
					running = false
				case sdl.KeyF:
					fullscreen = !fullscreen
					sdl.SetWindowFullscreen(windowHandler, fullscreen)
					resized = true
				}
			}
		}
		if !running {
			break
		}
		if resized {
			if err := recreateSwapChain(); err != nil {
				if errors.Is(err, errMinimized) {
					resized = true
					continue
				}
				panic(err)
			}
			resized = false
		}
		if !granted {
			select {
			case grant <- struct{}{}:
				granted = true
			case <-ctx.Done():
				running = false
				continue
			}
		}
		var shownIndex uint64
		select {
		case shownIndex = <-staged:
			granted = false
		case <-failed:
			granted = false
		case <-ctx.Done():
			running = false
			continue
		case <-time.After(100 * time.Millisecond):
			continue
		}

		// Acquire the next swapchain image. Right after we have a frame
		imageIndex, r := vkDevice.AcquireNextImage(vkSwapchain, imageAvailable, ^uint64(0))
		if r == vk.ErrorOutOfDateKHR || r == vk.SuboptimalKHR {
			if err := recreateSwapChain(); err != nil {
				if errors.Is(err, errMinimized) {
					resized = true
					continue
				}
				panic(err)
			}
			continue
		}
		if !r.Ok() {
			panic(fmt.Errorf("acquire: %v", r))
		}
		// Wait until GPU is done with our command buffer, then reset fence
		if err := vkDevice.WaitFence(inFlight, ^uint64(0)); err != nil {
			panic(err)
		}
		if err := vkDevice.ResetFence(inFlight); err != nil {
			panic(err)
		}

		if lastIndex != 0 && shownIndex > lastIndex {
			if g := shownIndex - lastIndex - 1; g > 0 {
				dropped += g
			}
		}
		lastIndex = shownIndex
		frameCount++
		if now := time.Now(); now.Sub(lastReport) >= time.Second {
			dt := now.Sub(lastReport).Seconds()
			fps := float64(frameCount) / dt
			fmt.Printf("fps=%.1f dropped=%d idx=%d frameTime=%.2fms\n",
				fps, dropped, shownIndex, float64(now.Sub(frameStart).Microseconds())/1000.0)
			frameCount = 0
			dropped = 0
			lastReport = now
		}

		// Record: copy staged V210 into device-local buffer, then clear
		cmd := vkCommands[0]
		if err := cmd.Reset(); err != nil {
			panic(err)
		}
		if err := cmd.Begin(vk.CommandBufferOneTimeSubmit); err != nil {
			panic(err)
		}
		cmd.CopyBuffer(staging.Buffer, v210Buf.Buffer, frameSize)
		cmd.BeginRenderPass(
			vkRenderPass,
			fbs[imageIndex],
			vk.Rect2D{Offset: vk.Offset2D{X: 0, Y: 0}, Extent: vkExtent},
			[]vk.ClearValue{
				vk.ClearColor(0.0, 0.4, 0.7, 1.0),
				vk.ClearDepthStencil(1.0, 0),
			},
		)
		cmd.SetViewport(vk.Viewport{
			X: 0, Y: 0,
			Width:    float32(vkExtent.Width),
			Height:   float32(vkExtent.Height),
			MinDepth: 0, MaxDepth: 1,
		})
		cmd.SetScissor(vk.Rect2D{Offset: vk.Offset2D{X: 0, Y: 0}, Extent: vkExtent})
		cmd.BindPipeline(decodePipeline)
		cmd.BindDescriptorSet(decodeLayout, 0, decodeSet)
		decodePC := decodePushConstants{
			Width:       mxlSrc.Width(),
			Height:      mxlSrc.Height(),
			StrideBytes: mxlSrc.Stride(),
			WinW:        vkExtent.Width,
			WinH:        vkExtent.Height,
		}
		cmd.PushConstants(decodeLayout, vk.ShaderStageFragment, 0, unsafe.Pointer(&decodePC), 20)
		cmd.Draw(3, 1, 0, 0)
		cmd.EndRenderPass()
		if err := cmd.End(); err != nil {
			panic(err)
		}
		// Submit: wait on imageAvailable, signal RenderFinished + inFlight
		if err := vkQueue.Submit(vk.SubmitConfig{
			Wait:      imageAvailable,
			WaitStage: vk.StageColorAttachmentOutput,
			Command:   cmd,
			Signal:    renderFinished,
			Fence:     inFlight,
		}); err != nil {
			panic(err)
		}
		// Present. Wait on renderFinished
		r = vkQueue.Present(vkSwapchain, imageIndex, renderFinished)
		if r == vk.ErrorOutOfDateKHR || r == vk.SuboptimalKHR {
			if err := recreateSwapChain(); err != nil {
				if errors.Is(err, errMinimized) {
					resized = true
					continue
				}
				panic(err)
			}
			continue
		}
		if !r.Ok() {
			panic(fmt.Errorf("present: %v", r))
		}
	}
}