package main

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

	vk "github.com/christerso/vulkan-go/vk"
	"github.com/ebitengine/purego"
)

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

const (
	initVideo       uint32 = 0x00000020
	windowVulkan    uint64 = 0x0000000010000000
	windowResizable uint64 = 0x0000000000000020

	sdlEventQuit             uint32 = 0x100
	sdlEventWindowResized    uint32 = 0x206
	sdlEventPixelSizeChanged uint32 = 0x207
	sdlEventKeyDown          uint32 = 0x300
	keyEscape                uint32 = 0x1B
)

var (
	sdlInit                        func(flags uint32) bool
	sdlQuit                        func()
	sdlGetError                    func() uintptr
	sdlCreateWindow                func(title *byte, w, h int32, flags uint64) uintptr
	sdlDestroyWindow               func(window uintptr)
	sdlVulkanGetInstanceExtensions func(count *uint32) uintptr
	sdlVulkanCreateSurface         func(window, instance, allocator uintptr, surface *uint64) bool
	sdlPollEvent                   func(event unsafe.Pointer) bool
	sdlGetWindowSizeInPixels       func(window uintptr, w, h *int32) bool
	sdlSetWindowFullscreen         func(windows uintptr, fullscreen bool) bool
)

func sdlError() string { return cstr(sdlGetError()) }

var sdlLoaded = false

func loadSDLMissing() error {
	if sdlLoaded {
		return nil
	}
	h, err := purego.Dlopen("libSDL3.so.0", purego.RTLD_NOW|purego.RTLD_GLOBAL)
	if err != nil {
		return fmt.Errorf("win: load SDL3: %w", err)
	}
	purego.RegisterLibFunc(&sdlInit, h, "SDL_Init")
	purego.RegisterLibFunc(&sdlQuit, h, "SDL_Quit")
	purego.RegisterLibFunc(&sdlGetError, h, "SDL_GetError")
	purego.RegisterLibFunc(&sdlCreateWindow, h, "SDL_CreateWindow")
	purego.RegisterLibFunc(&sdlDestroyWindow, h, "SDL_DestroyWindow")
	purego.RegisterLibFunc(&sdlVulkanGetInstanceExtensions, h, "SDL_Vulkan_GetInstanceExtensions")
	purego.RegisterLibFunc(&sdlVulkanCreateSurface, h, "SDL_Vulkan_CreateSurface")
	purego.RegisterLibFunc(&sdlPollEvent, h, "SDL_PollEvent")
	purego.RegisterLibFunc(&sdlGetWindowSizeInPixels, h, "SDL_GetWindowSizeInPixels")
	purego.RegisterLibFunc(&sdlSetWindowFullscreen, h, "SDL_SetWindowFullscreen")
	sdlLoaded = true
	return nil
}

// cstr reads a NUL-terminated C string at p.
func cstr(p uintptr) string {
	if p == 0 {
		return ""
	}
	var n int
	for *(*byte)(unsafe.Pointer(p + uintptr(n))) != 0 {
		n++
	}
	return string(unsafe.Slice((*byte)(unsafe.Pointer(p)), n))
}

// cbytes returns a NUL-terminated copy of s as *byte, kept alive by the caller.
func cbytes(s string) *byte {
	b := make([]byte, len(s)+1)
	copy(b, s)
	runtime.KeepAlive(b)
	return &b[0]
}

func main() {
	// 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 := loadSDLMissing(); err != nil {
		panic(err)
	}
	if !sdlInit(initVideo) {
		log.Fatalf("SDL_Init: %s", sdlError())
		return
	}

	windowHandler := sdlCreateWindow(cbytes(fmt.Sprintf("%s %s", APP_NAME, APP_VER)), WIN_WIDTH, WIN_HEIGHT, windowVulkan|windowResizable)
	if windowHandler == 0 {
		sdlQuit()
		log.Fatalf("SDL_CreateWindow: %s", sdlError())
		return
	}

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

	// Vulkan instance extensions SDL needs.
	var sdlExtCount uint32
	arr := sdlVulkanGetInstanceExtensions(&sdlExtCount)
	if arr == 0 {
		log.Fatal("sdlVulkanGetInstanceExtensions is 0")
		return
	}
	sdlExtensions := make([]string, sdlExtCount)
	for i := uint32(0); i < sdlExtCount; i++ {
		p := *(*uintptr)(unsafe.Pointer(arr + uintptr(i)*unsafe.Sizeof(uintptr(0))))
		sdlExtensions[i] = cstr(p)
	}
	// 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 !sdlVulkanCreateSurface(windowHandler, uintptr(vkInstance), 0, &vkSurface) {
		log.Fatalf("SDL_Vulkan_CreateSurface: %s", sdlError())
		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 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 !sdlGetWindowSizeInPixels(windowHandler, &w, &h) {
			return fmt.Errorf("SDL_GetWindowSizeInPixels: %s", sdlError())
		}
		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)
		}
		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{})
	staged := make(chan struct{})

	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
				}
				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 <- struct{}{}:
			case <-ctx.Done():
				return
			}
		}
	}()

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

	// Core Loop
	running := true
	resized := false
	fullscreen := false
	for running {
		var event [128]byte
		for sdlPollEvent(unsafe.Pointer(&event[0])) {
			eventType := *(*uint32)(unsafe.Pointer(&event[0]))
			switch eventType {
			case sdlEventQuit:
				running = false
			case sdlEventWindowResized, sdlEventPixelSizeChanged:
				resized = true
			case sdlEventKeyDown:
				key := *(*int32)(unsafe.Pointer(&event[28]))
				switch uint32(key) {
				case keyEscape:
					running = false
				// f
				case 0x66:
					fullscreen = !fullscreen
					sdlSetWindowFullscreen(windowHandler, fullscreen)
					resized = true
				}
			}
		}
		if !running {
			break
		}
		if resized {
			if err := recreateSwapChain(); err != nil {
				panic(err)
			}
			resized = false
		}
		// Acquire the next swapchain image
		imageIndex, r := vkDevice.AcquireNextImage(vkSwapchain, imageAvailable, ^uint64(0))
		if r == vk.ErrorOutOfDateKHR || r == vk.SuboptimalKHR {
			if err := recreateSwapChain(); err != nil {
				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)
		}

		// Grant the goroutine permission to write staging (GPU is idle now),
		// then wait for it to stage a new payload.
		select {
		case grant <- struct{}{}:
		case <-ctx.Done():
			running = false
			continue
		}
		select {
		case <-staged:
		case <-ctx.Done():
			running = false
			continue
		}

		// 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 {
				panic(err)
			}
			continue
		}
		if !r.Ok() {
			panic(fmt.Errorf("present: %v", r))
		}
	}
}