package renderer

import (
	"errors"
	"fmt"
	"mxl-player/internal/sdl"
	"unsafe"

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

// Raises by swapchain creation, whem windos is minimized
var ErrMinimized = errors.New("window minimized")

// ErrOutOfDate is returned by DrawFrame when the swapchain needs recreation.
var ErrOutOfDate = errors.New("swapchain out of date")

// shader push-constants block
type PushConstants struct {
	Width, Height, StrideBytes, WinW, WinH uint32
}

// all Vulkan state
type Renderer struct {
	pd     vk.PhysicalDevice
	dev    vk.Device
	queue  vk.Queue
	surf   vk.SurfaceKHR
	window uintptr

	format      vk.Format
	colorSpace  uint32
	presentMode vk.PresentMode
	renderPass  vk.RenderPass

	cmdPool  vk.CommandPool
	commands []vk.CommandBuffer

	staging   vk.AllocBuffer
	v210Buf   vk.AllocBuffer
	frameSize vk.DeviceSize

	vertMod        vk.ShaderModule
	fragMod        vk.ShaderModule
	decodeDSL      vk.DescriptorSetLayout
	decodeLayout   vk.PipelineLayout
	decodePipeline vk.Pipeline
	descPool       vk.DescriptorPool
	decodeSet      vk.DescriptorSet

	extent     vk.Extent2D
	swapchain  vk.SwapchainKHR
	swapImages []vk.Image
	swapViews  []vk.ImageView
	depthImg   vk.AllocImage
	depthView  vk.ImageView
	fbs        []vk.Framebuffer

	imageAvailable vk.Semaphore
	renderFinished vk.Semaphore
	inFlight       vk.Fence
}

// public params, needed to create a Renderer
type Config struct {
	PhysDevice                           vk.PhysicalDevice
	Device                               vk.Device
	Queue                                vk.Queue
	Surface                              vk.SurfaceKHR
	Window                               uintptr
	GraphicsFamily                       uint32
	VideoWidth, VideoHeight, VideoStride uint32
}

// creates Renderer with all Vulkan obkects initialized
func New(cfg Config) (*Renderer, error) {
	r := &Renderer{
		pd:     cfg.PhysDevice,
		dev:    cfg.Device,
		queue:  cfg.Queue,
		surf:   cfg.Surface,
		window: cfg.Window,
	}

	// Format: prefer 10-bit A2B10G10R10 (64), fallback to first
	// TODO: check other devices,
	// maybe it's not the best idea
	formats, _ := cfg.PhysDevice.SurfaceFormats(cfg.Surface)
	r.format = vk.FormatUndefined
	for _, f := range formats {
		if f.Format == vk.Format(64) && f.ColorSpace == vk.ColorSpaceSRGBNonlinear {
			r.format = f.Format
			r.colorSpace = f.ColorSpace
		}
	}
	if r.format == vk.FormatUndefined {
		r.format = formats[0].Format
		r.colorSpace = formats[0].ColorSpace
	}

	// Present mode. Prefer Mailbox, fallback to FIFO
	r.presentMode = vk.PresentModeFIFO
	modes, _ := cfg.PhysDevice.SurfacePresentModes(cfg.Surface)
	for _, m := range modes {
		if m == vk.PresentModeMailbox {
			r.presentMode = m
		}
	}

	// Render pass
	rp, err := cfg.Device.CreateColorDepthRenderPass(r.format, vk.FormatD32Sfloat)
	if err != nil {
		return nil, err
	}
	r.renderPass = rp

	// Command pool + buffers
	pool, err := cfg.Device.CreateCommandPool(cfg.GraphicsFamily)
	if err != nil {
		cfg.Device.DestroyRenderPass(rp)
		return nil, err
	}
	r.cmdPool = pool
	cmds, err := cfg.Device.AllocateCommandBuffers(pool, 1)
	if err != nil {
		cfg.Device.DestroyCommandPool(pool)
		cfg.Device.DestroyRenderPass(rp)
		return nil, err
	}
	r.commands = cmds

	// Buffers
	r.frameSize = vk.DeviceSize(cfg.VideoStride) * vk.DeviceSize(cfg.VideoHeight)
	r.staging, err = cfg.Device.CreateBuffer(cfg.PhysDevice, vk.BufferConfig{
		Size:       r.frameSize,
		Usage:      vk.BufferUsageTransferSrc,
		Properties: vk.MemoryHostVisible | vk.MemoryHostCoherent,
		Map:        true,
	})
	if err != nil {
		r.destroyPreSwapchain()
		return nil, err
	}

	// Shaders + pipeline
	r.vertMod, err = cfg.Device.CreateShaderModule(vertSPV)
	if err != nil {
		r.destroyPreSwapchain()
		return nil, err
	}
	r.fragMod, err = cfg.Device.CreateShaderModule(fragSPV)
	if err != nil {
		r.destroyPreSwapchain()
		return nil, err
	}
	r.decodeDSL, err = cfg.Device.CreateDescriptorSetLayout([]vk.DescriptorBinding{
		{Binding: 0, Type: vk.DescriptorStorageBuffer, Count: 1, Stages: vk.ShaderStageFragment},
	})
	if err != nil {
		r.destroyPreSwapchain()
		return nil, err
	}
	r.decodeLayout, err = cfg.Device.CreatePipelineLayout(
		[]vk.DescriptorSetLayout{r.decodeDSL}, vk.ShaderStageFragment, 20)
	if err != nil {
		r.destroyPreSwapchain()
		return nil, err
	}
	r.decodePipeline, err = cfg.Device.CreateGraphicsPipeline(vk.GraphicsPipelineConfig{
		Layout:       r.decodeLayout,
		RenderPass:   r.renderPass,
		VertexShader: r.vertMod,
		FragShader:   r.fragMod,
		Topology:     vk.TopologyTriangleList,
		PolygonMode:  vk.PolygonFill,
		CullMode:     vk.CullNone,
		FrontFace:    vk.FrontFaceCounterClockwise,
	})
	if err != nil {
		r.destroyPreSwapchain()
		return nil, err
	}

	// Descriptor pool + set
	r.descPool, err = cfg.Device.CreateDescriptorPool(1, map[vk.DescriptorType]uint32{
		vk.DescriptorStorageBuffer: 1,
	})
	if err != nil {
		r.destroyPreSwapchain()
		return nil, err
	}
	r.decodeSet, err = cfg.Device.AllocateDescriptorSet(r.descPool, r.decodeDSL)
	if err != nil {
		r.destroyPreSwapchain()
		return nil, err
	}
	cfg.Device.UpdateBufferDescriptor(r.decodeSet, 0, vk.DescriptorStorageBuffer, r.staging.Buffer, 0, vk.WholeSize)

	// Swapchain
	if err := r.createSwapchain(); err != nil {
		r.destroyPreSwapchain()
		return nil, err
	}

	// Sync objects
	r.imageAvailable, err = cfg.Device.CreateSemaphore()
	if err != nil {
		r.destroyAll()
		return nil, err
	}
	r.renderFinished, err = cfg.Device.CreateSemaphore()
	if err != nil {
		r.destroyAll()
		return nil, err
	}
	r.inFlight, err = cfg.Device.CreateFence(true)
	if err != nil {
		r.destroyAll()
		return nil, err
	}

	return r, nil
}

func (r *Renderer) createSwapchain() (err error) {
	defer func() {
		if err != nil {
			r.destroySwapchain()
		}
	}()
	caps, err := r.pd.SurfaceCapabilities(r.surf)
	if err != nil {
		return err
	}
	var w, h int32 = 1280, 720
	if !sdl.GetWindowSizeInPixels(r.window, &w, &h) {
		return fmt.Errorf("SDL_GetWindowSizeInPixels: %s", sdl.GetError())
	}
	clamp(&w, int32(caps.MinImageExtent.Width), int32(caps.MaxImageExtent.Width))
	clamp(&h, int32(caps.MinImageExtent.Height), int32(caps.MaxImageExtent.Height))
	if w == 0 || h == 0 {
		return ErrMinimized
	}
	r.extent = vk.Extent2D{Width: uint32(w), Height: uint32(h)}

	r.swapchain, err = r.dev.CreateSwapchain(vk.SwapchainConfig{
		Surface:       r.surf,
		MinImageCount: caps.MinImageCount,
		Format:        r.format,
		ColorSpace:    r.colorSpace,
		Extent:        r.extent,
		PresentMode:   r.presentMode,
		PreTransform:  caps.CurrentTransform,
	})
	if err != nil {
		return err
	}
	r.swapImages, err = r.dev.SwapchainImages(r.swapchain)
	if err != nil {
		return err
	}
	r.swapViews = make([]vk.ImageView, len(r.swapImages))
	for i, img := range r.swapImages {
		v, err := r.dev.CreateImageView(img, r.format, vk.AspectColor)
		if err != nil {
			return err
		}
		r.swapViews[i] = v
	}
	r.depthImg, err = r.dev.CreateImage2D(r.pd, vk.FormatD32Sfloat, r.extent, vk.ImageUsageDepthStencilAttachment)
	if err != nil {
		return err
	}
	r.depthView, err = r.dev.CreateImageView(r.depthImg.Image, vk.FormatD32Sfloat, vk.AspectDepth)
	if err != nil {
		return err
	}
	r.fbs = make([]vk.Framebuffer, len(r.swapViews))
	for i, cv := range r.swapViews {
		fb, err := r.dev.CreateFramebuffer(r.renderPass, []vk.ImageView{cv, r.depthView}, r.extent)
		if err != nil {
			return err
		}
		r.fbs[i] = fb
	}

	return nil
}

func (r *Renderer) destroySwapchain() {
	for _, fb := range r.fbs {
		r.dev.DestroyFramebuffer(fb)
	}
	r.dev.DestroyImageView(r.depthView)
	r.dev.DestroyImage(r.depthImg)
	for _, v := range r.swapViews {
		r.dev.DestroyImageView(v)
	}
	r.dev.DestroySwapchain(r.swapchain)
	r.fbs = nil
	r.swapViews = nil
	r.swapImages = nil
	r.depthView = 0
	r.depthImg = vk.AllocImage{}
	r.swapchain = 0
	r.extent = vk.Extent2D{}
}

func (r *Renderer) destroyPreSwapchain() {
	r.dev.DestroyDescriptorPool(r.descPool)
	r.dev.DestroyPipeline(r.decodePipeline)
	r.dev.DestroyPipelineLayout(r.decodeLayout)
	r.dev.DestroyDescriptorSetLayout(r.decodeDSL)
	r.dev.DestroyShaderModule(r.fragMod)
	r.dev.DestroyShaderModule(r.vertMod)
	r.dev.DestroyBuffer(r.v210Buf)
	r.dev.DestroyBuffer(r.staging)
	r.dev.DestroyCommandPool(r.cmdPool)
	r.dev.DestroyRenderPass(r.renderPass)
}

func (r *Renderer) destroyAll() {
	r.dev.DestroyFence(r.inFlight)
	r.dev.DestroySemaphore(r.renderFinished)
	r.dev.DestroySemaphore(r.imageAvailable)
	r.destroySwapchain()
	r.destroyPreSwapchain()
}

// Tears down everything. Call after Device.WaitIdle
func (r *Renderer) Destroy() {
	r.destroyAll()
}

// Persistently mapped CPU pointer for v210 uploads
func (r *Renderer) StagingMapped() unsafe.Pointer { return r.staging.Mapped }

// Current swapchain extent
func (r *Renderer) Extent() vk.Extent2D { return r.extent }

// Current v210 frame size in bytes
func (r *Renderer) FrameSize() vk.DeviceSize { return r.frameSize }

func clamp(val *int32, min, max int32) {
	if *val > max {
		*val = max
	} else if *val < min {
		*val = min
	}
}

// RecreateSwapchain destroys and recreates the swapchain + dependent views/fbs
func (r *Renderer) RecreateSwapchain() error {
	if err := r.dev.WaitIdle(); err != nil {
		return err
	}
	r.destroySwapchain()
	return r.createSwapchain()
}

// RecreateBuffers recreates the staging + v210 buffers for a new frame size
func (r *Renderer) RecreateBuffers(newSize vk.DeviceSize) error {
	if err := r.dev.WaitIdle(); err != nil {
		return err
	}
	r.dev.DestroyBuffer(r.v210Buf)
	r.dev.DestroyBuffer(r.staging)
	var err error
	r.staging, err = r.dev.CreateBuffer(r.pd, vk.BufferConfig{
		Size:       newSize,
		Usage:      vk.BufferUsageTransferSrc,
		Properties: vk.MemoryHostVisible | vk.MemoryHostCoherent,
		Map:        true,
	})
	if err != nil {
		return err
	}
	r.v210Buf, err = r.dev.CreateBuffer(r.pd, vk.BufferConfig{
		Size:       newSize,
		Usage:      vk.BufferUsageTransferDst | vk.BufferUsageStorageBuffer,
		Properties: vk.MemoryDeviceLocal,
		Map:        false,
	})
	if err != nil {
		return err
	}
	r.frameSize = newSize
	r.dev.UpdateBufferDescriptor(r.decodeSet, 0, vk.DescriptorStorageBuffer, r.staging.Buffer, 0, vk.WholeSize)
	return nil
}

// DrawFrame acquires an image, records commands, submits, and presents.
// Returns ErrOutOfDate if the swapchain needs recreation
func (r *Renderer) DrawFrame(videoW, videoH, stride uint32) error {
	imageIndex, res := r.dev.AcquireNextImage(r.swapchain, r.imageAvailable, ^uint64(0))
	if res == vk.ErrorOutOfDateKHR || res == vk.SuboptimalKHR {
		return ErrOutOfDate
	}
	if !res.Ok() {
		return fmt.Errorf("acquire: %v", res)
	}

	if err := r.dev.WaitFence(r.inFlight, ^uint64(0)); err != nil {
		return err
	}
	if err := r.dev.ResetFence(r.inFlight); err != nil {
		return err
	}

	cmd := r.commands[0]
	if err := cmd.Reset(); err != nil {
		return err
	}
	if err := cmd.Begin(vk.CommandBufferOneTimeSubmit); err != nil {
		return err
	}
	cmd.BeginRenderPass(
		r.renderPass,
		r.fbs[imageIndex],
		vk.Rect2D{Offset: vk.Offset2D{X: 0, Y: 0}, Extent: r.extent},
		[]vk.ClearValue{
			vk.ClearColor(0.0, 0.0, 0.0, 1.0),
			vk.ClearDepthStencil(1.0, 0),
		},
	)
	cmd.SetViewport(vk.Viewport{
		X: 0, Y: 0,
		Width:    float32(r.extent.Width),
		Height:   float32(r.extent.Height),
		MinDepth: 0, MaxDepth: 1,
	})
	cmd.SetScissor(vk.Rect2D{Offset: vk.Offset2D{X: 0, Y: 0}, Extent: r.extent})
	cmd.BindPipeline(r.decodePipeline)
	cmd.BindDescriptorSet(r.decodeLayout, 0, r.decodeSet)
	pc := PushConstants{
		Width:       videoW,
		Height:      videoH,
		StrideBytes: stride,
		WinW:        r.extent.Width,
		WinH:        r.extent.Height,
	}
	cmd.PushConstants(r.decodeLayout, vk.ShaderStageFragment, 0, unsafe.Pointer(&pc), 20)
	cmd.Draw(3, 1, 0, 0)
	cmd.EndRenderPass()
	if err := cmd.End(); err != nil {
		return err
	}

	if err := r.queue.Submit(vk.SubmitConfig{
		Wait:      r.imageAvailable,
		WaitStage: vk.StageColorAttachmentOutput,
		Command:   cmd,
		Signal:    r.renderFinished,
		Fence:     r.inFlight,
	}); err != nil {
		return err
	}

	res = r.queue.Present(r.swapchain, imageIndex, r.renderFinished)
	if res == vk.ErrorOutOfDateKHR || res == vk.SuboptimalKHR {
		return ErrOutOfDate
	}
	if !res.Ok() {
		return fmt.Errorf("present: %v", res)
	}
	return nil
}