main.go

// Package main demonstrates RenderBundle for pre-recording render commands.
// RenderBundles are useful for static geometry that doesn't change between frames,
// reducing CPU overhead by pre-recording draw commands.
package main

import (
	"fmt"
	"log"
	"os"
	"syscall"
	"unsafe"

	"github.com/go-webgpu/webgpu/wgpu"
	"github.com/gogpu/gputypes"
	"golang.org/x/sys/windows"
)

const (
	windowWidth  = 800
	windowHeight = 600
	windowTitle  = "go-webgpu: RenderBundle Example"
)

// Win32 constants
const (
	csHRedraw                 = 0x0002
	csVRedraw                 = 0x0001
	wmDestroy                 = 0x0002
	wmSize                    = 0x0005
	idcArrow                  = 32512
	colorWindow               = 5
	swShowNormal              = 1
	pmRemove                  = 0x0001
	wsOverlappedWindow        = 0x00CF0000
	wsVisible                 = 0x10000000
	cwUseDefault       uint32 = 0x80000000
)

var (
	user32               = windows.NewLazyDLL("user32.dll")
	procRegisterClassExW = user32.NewProc("RegisterClassExW")
	procCreateWindowExW  = user32.NewProc("CreateWindowExW")
	procShowWindow       = user32.NewProc("ShowWindow")
	procUpdateWindow     = user32.NewProc("UpdateWindow")
	procPeekMessageW     = user32.NewProc("PeekMessageW")
	procTranslateMessage = user32.NewProc("TranslateMessage")
	procDispatchMessageW = user32.NewProc("DispatchMessageW")
	procDefWindowProcW   = user32.NewProc("DefWindowProcW")
	procPostQuitMessage  = user32.NewProc("PostQuitMessage")
	procLoadCursorW      = user32.NewProc("LoadCursorW")
	kernel32             = windows.NewLazyDLL("kernel32.dll")
	procGetModuleHandleW = kernel32.NewProc("GetModuleHandleW")
)

// WNDCLASSEXW represents the Win32 WNDCLASSEXW structure.
type WNDCLASSEXW struct {
	cbSize        uint32
	style         uint32
	lpfnWndProc   uintptr
	cbClsExtra    int32
	cbWndExtra    int32
	hInstance     windows.Handle
	hIcon         windows.Handle
	hCursor       windows.Handle
	hbrBackground windows.Handle
	lpszMenuName  *uint16
	lpszClassName *uint16
	hIconSm       windows.Handle
}

// MSG represents the Win32 MSG structure.
type MSG struct {
	hwnd    windows.HWND
	message uint32
	wParam  uintptr
	lParam  uintptr
	time    uint32
	pt      struct{ x, y int32 }
}

// Application state
type App struct {
	hwnd           windows.HWND
	hinstance      windows.Handle
	instance       *wgpu.Instance
	adapter        *wgpu.Adapter
	device         *wgpu.Device
	queue          *wgpu.Queue
	surface        *wgpu.Surface
	pipeline       *wgpu.RenderPipeline
	renderBundle   *wgpu.RenderBundle
	width          uint32
	height         uint32
	running        bool
	needsRecreate  bool
	surfaceTex     *wgpu.SurfaceTexture
	surfaceTexView *wgpu.TextureView
}

// Shader source (WGSL)
// NOTE: This shader uses a fallback approach for triangle coloring.
// Instead of @builtin(primitive_index) (which requires PRIMITIVE_INDEX GPU capability),
// we calculate the triangle index from vertex_index and pass color via varying.
// This works on ALL GPUs, including older hardware without primitive_index support.
const shaderSource = `
struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) color: vec3<f32>,
}

@vertex
fn vs_main(@builtin(vertex_index) idx: u32) -> VertexOutput {
    // Three triangles at different positions
    var positions = array<vec2<f32>, 9>(
        // Triangle 1 (left)
        vec2<f32>(-0.7, 0.3),
        vec2<f32>(-0.9, -0.3),
        vec2<f32>(-0.5, -0.3),
        // Triangle 2 (center)
        vec2<f32>(0.0, 0.5),
        vec2<f32>(-0.2, -0.1),
        vec2<f32>(0.2, -0.1),
        // Triangle 3 (right)
        vec2<f32>(0.7, 0.3),
        vec2<f32>(0.5, -0.3),
        vec2<f32>(0.9, -0.3)
    );

    // Colors for each triangle
    var colors = array<vec3<f32>, 3>(
        vec3<f32>(1.0, 0.2, 0.2),  // Red
        vec3<f32>(0.2, 1.0, 0.2),  // Green
        vec3<f32>(0.2, 0.2, 1.0)   // Blue
    );

    // Calculate triangle index from vertex index (3 vertices per triangle)
    // This is the FALLBACK for @builtin(primitive_index)
    let tri_idx = idx / 3u;

    var out: VertexOutput;
    out.position = vec4<f32>(positions[idx], 0.0, 1.0);
    out.color = colors[tri_idx];
    return out;
}

@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    // Use color passed from vertex shader (works on all GPUs!)
    return vec4<f32>(in.color, 1.0);
}
`

func main() {
	fmt.Println("=== RenderBundle Example ===")
	fmt.Println()
	fmt.Println("RenderBundles pre-record draw commands for efficient replay.")
	fmt.Println("This is useful for static geometry that doesn't change.")
	fmt.Println()
	fmt.Println("Benefits:")
	fmt.Println("  - Reduced CPU overhead (commands recorded once)")
	fmt.Println("  - Better driver optimization opportunities")
	fmt.Println("  - Useful for static scene elements")
	fmt.Println()
	fmt.Println("Note: This example uses a FALLBACK approach for per-triangle coloring.")
	fmt.Println("Instead of @builtin(primitive_index) (requires PRIMITIVE_INDEX capability),")
	fmt.Println("we calculate triangle index from vertex_index in the vertex shader.")
	fmt.Println("This works on ALL GPUs, including older hardware.")
	fmt.Println()

	app := &App{
		width:   windowWidth,
		height:  windowHeight,
		running: true,
	}

	if err := app.init(); err != nil {
		log.Fatalf("Failed to initialize: %v", err)
	}
	defer app.cleanup()

	app.run()
}

// init initializes the application.
func (app *App) init() error {
	// Initialize WebGPU library
	if err := wgpu.Init(); err != nil {
		return fmt.Errorf("init wgpu: %w", err)
	}

	// Get HINSTANCE
	ret, _, _ := procGetModuleHandleW.Call(0)
	app.hinstance = windows.Handle(ret)

	// Create window
	if err := app.createWindow(); err != nil {
		return fmt.Errorf("create window: %w", err)
	}

	// Initialize WebGPU
	if err := app.initWebGPU(); err != nil {
		return fmt.Errorf("init webgpu: %w", err)
	}

	// Configure surface
	if err := app.configureSurface(); err != nil {
		return fmt.Errorf("configure surface: %w", err)
	}

	// Create render pipeline
	if err := app.createPipeline(); err != nil {
		return fmt.Errorf("create pipeline: %w", err)
	}

	// Create render bundle (pre-recorded commands)
	if err := app.createRenderBundle(); err != nil {
		return fmt.Errorf("create render bundle: %w", err)
	}

	return nil
}

// createWindow creates the main window.
func (app *App) createWindow() error {
	className, err := windows.UTF16PtrFromString("GoWebGPURenderBundle")
	if err != nil {
		return err
	}

	wndClass := WNDCLASSEXW{
		cbSize:        uint32(unsafe.Sizeof(WNDCLASSEXW{})),
		style:         csHRedraw | csVRedraw,
		lpfnWndProc:   syscall.NewCallback(app.wndProc),
		hInstance:     app.hinstance,
		lpszClassName: className,
	}

	cursor, _, _ := procLoadCursorW.Call(0, uintptr(idcArrow))
	wndClass.hCursor = windows.Handle(cursor)
	wndClass.hbrBackground = windows.Handle(colorWindow + 1)

	ret, _, _ := procRegisterClassExW.Call(uintptr(unsafe.Pointer(&wndClass)))
	if ret == 0 {
		return fmt.Errorf("RegisterClassExW failed")
	}

	titlePtr, err := windows.UTF16PtrFromString(windowTitle)
	if err != nil {
		return err
	}

	hwnd, _, _ := procCreateWindowExW.Call(
		0,
		uintptr(unsafe.Pointer(className)),
		uintptr(unsafe.Pointer(titlePtr)),
		uintptr(wsOverlappedWindow|wsVisible),
		uintptr(cwUseDefault),
		uintptr(cwUseDefault),
		uintptr(app.width),
		uintptr(app.height),
		0,
		0,
		uintptr(app.hinstance),
		0,
	)

	if hwnd == 0 {
		return fmt.Errorf("CreateWindowExW failed")
	}

	app.hwnd = windows.HWND(hwnd)

	_, _, _ = procShowWindow.Call(uintptr(app.hwnd), swShowNormal)
	_, _, _ = procUpdateWindow.Call(uintptr(app.hwnd))

	return nil
}

// wndProc is the window procedure callback.
func (app *App) wndProc(hwnd windows.HWND, msg uint32, wParam, lParam uintptr) uintptr {
	switch msg {
	case wmDestroy:
		app.running = false
		_, _, _ = procPostQuitMessage.Call(0)
		return 0
	case wmSize:
		newWidth := uint32(lParam & 0xFFFF)
		newHeight := uint32((lParam >> 16) & 0xFFFF)
		if newWidth != app.width || newHeight != app.height {
			app.width = newWidth
			app.height = newHeight
			app.needsRecreate = true
		}
		return 0
	}
	ret, _, _ := procDefWindowProcW.Call(
		uintptr(hwnd),
		uintptr(msg),
		wParam,
		lParam,
	)
	return ret
}

// initWebGPU initializes WebGPU resources.
func (app *App) initWebGPU() error {
	inst, err := wgpu.CreateInstance(nil)
	if err != nil {
		return fmt.Errorf("create instance: %w", err)
	}
	app.instance = inst

	adapter, err := inst.RequestAdapter(nil)
	if err != nil {
		return fmt.Errorf("request adapter: %w", err)
	}
	app.adapter = adapter

	device, err := adapter.RequestDevice(nil)
	if err != nil {
		return fmt.Errorf("request device: %w", err)
	}
	app.device = device

	app.queue = device.GetQueue()

	surface, err := inst.CreateSurfaceFromWindowsHWND(uintptr(app.hinstance), uintptr(app.hwnd))
	if err != nil {
		return fmt.Errorf("create surface: %w", err)
	}
	app.surface = surface

	return nil
}

// configureSurface configures the surface for rendering.
func (app *App) configureSurface() error {
	app.surface.Configure(&wgpu.SurfaceConfiguration{
		Device:      app.device,
		Format:      gputypes.TextureFormatBGRA8Unorm,
		Usage:       gputypes.TextureUsageRenderAttachment,
		Width:       app.width,
		Height:      app.height,
		AlphaMode:   gputypes.CompositeAlphaModeOpaque,
		PresentMode: gputypes.PresentModeFifo,
	})
	app.needsRecreate = false
	return nil
}

// createPipeline creates the render pipeline.
func (app *App) createPipeline() error {
	shader := app.device.CreateShaderModuleWGSL(shaderSource)
	if shader == nil {
		return fmt.Errorf("failed to create shader module")
	}
	defer shader.Release()

	pipeline := app.device.CreateRenderPipelineSimple(
		nil,
		shader, "vs_main",
		shader, "fs_main",
		gputypes.TextureFormatBGRA8Unorm,
	)
	if pipeline == nil {
		return fmt.Errorf("failed to create render pipeline")
	}

	app.pipeline = pipeline
	return nil
}

// createRenderBundle creates a pre-recorded render bundle.
func (app *App) createRenderBundle() error {
	fmt.Println("Creating RenderBundle with 3 triangles...")

	// Create render bundle encoder with matching format
	colorFormats := []gputypes.TextureFormat{gputypes.TextureFormatBGRA8Unorm}
	bundleEncoder := app.device.CreateRenderBundleEncoderSimple(
		colorFormats,
		gputypes.TextureFormatUndefined, // no depth
		1,                               // sample count
	)
	if bundleEncoder == nil {
		return fmt.Errorf("failed to create render bundle encoder")
	}

	// Record draw commands (these will be replayed every frame)
	bundleEncoder.SetPipeline(app.pipeline)
	bundleEncoder.Draw(9, 1, 0, 0) // 9 vertices = 3 triangles

	// Finish recording
	bundle := bundleEncoder.Finish(nil)
	bundleEncoder.Release()

	if bundle == nil {
		return fmt.Errorf("failed to finish render bundle")
	}

	app.renderBundle = bundle
	fmt.Println("RenderBundle created successfully!")
	return nil
}

// releasePreviousFrame releases resources from the previous frame.
func (app *App) releasePreviousFrame() {
	if app.surfaceTexView != nil {
		app.surfaceTexView.Release()
		app.surfaceTexView = nil
	}
	if app.surfaceTex != nil && app.surfaceTex.Texture != nil {
		app.surfaceTex.Texture.Release()
		app.surfaceTex = nil
	}
}

// acquireSurfaceTexture gets the current surface texture.
func (app *App) acquireSurfaceTexture() error {
	surfaceTex, err := app.surface.GetCurrentTexture()
	if err != nil {
		if err == wgpu.ErrSurfaceLost || err == wgpu.ErrSurfaceNeedsReconfigure {
			app.needsRecreate = true
			return nil
		}
		return fmt.Errorf("get current texture: %w", err)
	}
	app.surfaceTex = surfaceTex

	view := surfaceTex.Texture.CreateView(nil)
	if view == nil {
		return fmt.Errorf("failed to create texture view")
	}
	app.surfaceTexView = view
	return nil
}

// render draws a frame.
func (app *App) render() error {
	if app.needsRecreate {
		if err := app.configureSurface(); err != nil {
			return fmt.Errorf("reconfigure surface: %w", err)
		}
	}

	app.releasePreviousFrame()

	if err := app.acquireSurfaceTexture(); err != nil {
		return err
	}
	if app.surfaceTexView == nil {
		return nil
	}

	encoder := app.device.CreateCommandEncoder(nil)
	if encoder == nil {
		return fmt.Errorf("failed to create command encoder")
	}
	defer encoder.Release()

	// Begin render pass
	pass := encoder.BeginRenderPass(&wgpu.RenderPassDescriptor{
		Label: "RenderBundle Pass",
		ColorAttachments: []wgpu.RenderPassColorAttachment{{
			View:    app.surfaceTexView,
			LoadOp:  gputypes.LoadOpClear,
			StoreOp: gputypes.StoreOpStore,
			ClearValue: wgpu.Color{
				R: 0.1,
				G: 0.1,
				B: 0.15,
				A: 1.0,
			},
		}},
	})
	if pass == nil {
		return fmt.Errorf("failed to begin render pass")
	}
	defer pass.Release()

	// Execute the pre-recorded render bundle!
	// This replays all the recorded draw commands efficiently.
	pass.ExecuteBundles([]*wgpu.RenderBundle{app.renderBundle})

	pass.End()

	cmdBuffer := encoder.Finish(nil)
	if cmdBuffer == nil {
		return fmt.Errorf("failed to finish command encoder")
	}
	defer cmdBuffer.Release()

	app.queue.Submit(cmdBuffer)
	app.surface.Present()

	return nil
}

// run is the main application loop.
func (app *App) run() {
	for app.running {
		var msg MSG
		for {
			ret, _, _ := procPeekMessageW.Call(
				uintptr(unsafe.Pointer(&msg)),
				0, 0, 0,
				pmRemove,
			)
			if ret == 0 {
				break
			}
			_, _, _ = procTranslateMessage.Call(uintptr(unsafe.Pointer(&msg)))
			_, _, _ = procDispatchMessageW.Call(uintptr(unsafe.Pointer(&msg)))
		}

		if err := app.render(); err != nil {
			fmt.Fprintf(os.Stderr, "Render error: %v\n", err)
			app.running = false
		}
	}
}

// cleanup releases all resources.
func (app *App) cleanup() {
	if app.surfaceTexView != nil {
		app.surfaceTexView.Release()
	}
	if app.surfaceTex != nil && app.surfaceTex.Texture != nil {
		app.surfaceTex.Texture.Release()
	}
	if app.renderBundle != nil {
		app.renderBundle.Release()
	}
	if app.pipeline != nil {
		app.pipeline.Release()
	}
	if app.surface != nil {
		app.surface.Release()
	}
	if app.queue != nil {
		app.queue.Release()
	}
	if app.device != nil {
		app.device.Release()
	}
	if app.adapter != nil {
		app.adapter.Release()
	}
	if app.instance != nil {
		app.instance.Release()
	}
}