package main

import (
	"log"
	"os"
	"runtime"
	"syscall"
	"time"
	"unsafe"

	"golang.org/x/sys/windows"
)

const (
	WM_DESTROY    = 0x0002
	WM_CLOSE      = 0x0010
	WM_QUIT       = 0x0012
	WM_MOUSEMOVE  = 0x0200
	WM_LBUTTONDOWN = 0x0201

	WS_POPUP        = 0x80000000
	WS_VISIBLE      = 0x10000000
	WS_CHILD        = 0x40000000
	WS_CLIPCHILDREN = 0x02000000
	WS_CLIPSIBLINGS = 0x04000000

	GWL_STYLE = ^uintptr(0) - 15

	PFD_DRAW_TO_WINDOW = 0x04
	PFD_SUPPORT_OPENGL = 0x20
	PFD_DOUBLEBUFFER   = 0x01

	GL_COLOR_BUFFER_BIT = 0x4060
	GL_TRIANGLE_STRIP   = 0x0005
	GL_ARRAY_BUFFER     = 0x8892
	GL_STATIC_DRAW      = 0x88E4
	GL_VERTEX_SHADER    = 0x8B31
	GL_FRAGMENT_SHADER  = 0x8B30
	GL_COMPILE_STATUS   = 0x8B81
	GL_INFO_LOG_LENGTH  = 0x8B84
	GL_RGBA             = 0x1908
	GL_UNSIGNED_BYTE    = 0x1401
	GL_FLOAT            = 0x1406
	GL_VERSION          = 0x1F02

	DIB_RGB_COLORS = 0
	SWP_NOSIZE       = 0x0001
	SWP_NOMOVE       = 0x0002
	SWP_NOZORDER     = 0x0004
	SWP_FRAMECHANGED = 0x0020
	HWND_BOTTOM      = 1
)

type bitmapInfoHeader struct {
	Size            uint32
	Width           int32
	Height          int32
	Planes          uint16
	BitCount        uint16
	Compression     uint32
	SizeImage       uint32
	XPelsPerMeter   int32
	YPelsPerMeter   int32
	ClrUsed         uint32
	ClrImportant    uint32
}

var (
	user32               = windows.NewLazySystemDLL("user32.dll")
	gdi32                = windows.NewLazySystemDLL("gdi32.dll")
	opengl32             = windows.NewLazySystemDLL("opengl32.dll")
	procFindWindowW      = user32.NewProc("FindWindowW")
	procFindWindowExW    = user32.NewProc("FindWindowExW")
	procSendMsgTimeout   = user32.NewProc("SendMessageTimeoutW")
	procSetParent        = user32.NewProc("SetParent")
	procMoveWindow       = user32.NewProc("MoveWindow")
	procShowWindow       = user32.NewProc("ShowWindow")
	procSetWindowLongPtrW = user32.NewProc("SetWindowLongPtrW")
	procSetWindowPos     = user32.NewProc("SetWindowPos")
	procCreateWindowExW  = user32.NewProc("CreateWindowExW")
	procDefWindowProcW   = user32.NewProc("DefWindowProcW")
	procRegisterClassExW = user32.NewProc("RegisterClassExW")
	procEnumDisplayMonitors = user32.NewProc("EnumDisplayMonitors")
	procEnumWindows         = user32.NewProc("EnumWindows")
	procGetClassNameW       = user32.NewProc("GetClassNameW")
	procGetWindowRect       = user32.NewProc("GetWindowRect")
	procSetProcessDPIAware  = user32.NewProc("SetProcessDPIAware")
	procPeekMessageW     = user32.NewProc("PeekMessageW")
	procTranslateMessage = user32.NewProc("TranslateMessage")
	procDispatchMessageW = user32.NewProc("DispatchMessageW")
	procPostQuitMessage  = user32.NewProc("PostQuitMessage")
	procGetDC            = user32.NewProc("GetDC")
	procReleaseDC        = user32.NewProc("ReleaseDC")
	procSetDIBitsToDevice = gdi32.NewProc("SetDIBitsToDevice")
	procChoosePixelFormat = gdi32.NewProc("ChoosePixelFormat")
	procSetPixelFormat   = gdi32.NewProc("SetPixelFormat")
	procSwapBuffers      = gdi32.NewProc("SwapBuffers")
	glFinish             = opengl32.NewProc("glFinish")
	procWglCreateContext = opengl32.NewProc("wglCreateContext")
	procWglMakeCurrent   = opengl32.NewProc("wglMakeCurrent")
	procWglDeleteContext = opengl32.NewProc("wglDeleteContext")
	procWglGetProcAddress = opengl32.NewProc("wglGetProcAddress")
	glViewport           = opengl32.NewProc("glViewport")
	glDrawArrays         = opengl32.NewProc("glDrawArrays")
	glGetString          = opengl32.NewProc("glGetString")
	glReadPixels         = opengl32.NewProc("glReadPixels")
	glClear              = opengl32.NewProc("glClear")
)

type rect struct{ Left, Top, Right, Bottom int32 }
type wndClassEx struct {
	CbSize, Style                          uint32
	LpfnWndProc                            uintptr
	CbClsExtra, CbWndExtra                 int32
	HInstance, HIcon, HCursor, HbrBackground uintptr
	LpszMenuName, LpszClassName            *uint16
	HIconSm                                uintptr
}
type pfd struct {
	Size, Version                                              uint16
	Flags                                                      uint32
	PixelType, ColorBits, RedBits, RedShift, GreenBits, GreenShift,
	BlueBits, BlueShift, AlphaBits, AlphaShift, AccumBits, AccumRed,
	AccumGreen, AccumBlue, AccumAlpha, DepthBits, StencilBits,
	AuxBuffers, LayerType, Reserved                            byte
	LayerMask, VisibleMask, DamageMask                          uint32
}

var (
	wallpaperHwnd    uintptr
	glDC             uintptr
	glCtx            uintptr
	mouseX, mouseY   float32
	clickX, clickY   float32
	clickTime        float64
	screenW, screenH int32
	startTime        = float64(time.Now().UnixNano()) / 1e9
	frameCount       int
	pixelBuf         []byte
	bmi              bitmapInfoHeader
	glUseProgram, glUniform1fv, glUniform2fv uintptr
	useSwap          bool
)

func main() {
	runtime.LockOSThread()
	log.SetFlags(log.Ltime | log.Lmicroseconds)
	procSetProcessDPIAware.Call()

	if len(os.Args) > 1 && os.Args[1] == "probe" {
		probe()
		return
	}
	screenW, screenH = getDesktopSize()
	log.Printf("Desktop: %dx%d", screenW, screenH)

	// Allocate pixel buffer for GDI blit
	pixels := screenW * screenH * 4
	pixelBuf = make([]byte, pixels)
	log.Printf("Pixel buffer: %d bytes (%.1fMB)", pixels, float64(pixels)/1024/1024)

	bmi = bitmapInfoHeader{
		Size:        uint32(unsafe.Sizeof(bitmapInfoHeader{})),
		Width:       screenW,
		Height:      screenH, // positive = bottom-up (matches GL)
		Planes:      1,
		BitCount:    32,
		Compression: 0, // BI_RGB
	}

	mode := ""
	if len(os.Args) > 1 {
		mode = os.Args[1]
	}
	useSwap = len(os.Args) > 2 && os.Args[2] == "swap"

	createWindow()

	switch mode {
	case "windowed":
		// Diagnostic: skip embedding, show as normal fullscreen window
		procMoveWindow.Call(wallpaperHwnd, 0, 0, uintptr(screenW), uintptr(screenH), 1)
		log.Println("WINDOWED mode (no embed)")
	default:
		// 选择嵌入目标：progman 模式直接挂到 Progman（Win11 新内核），否则找 WorkerW
		var target uintptr
		switch mode {
		case "progman", "progman-top":
			target, _, _ = procFindWindowW.Call(uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("Progman"))), 0)
			var discard uintptr
			procSendMsgTimeout.Call(target, 0x052C, 0xD, 0, 0x0000, 1000, uintptr(unsafe.Pointer(&discard)))
		case "childww", "childww-top":
			// Progman 的直接子 WorkerW —— 可能才是 DWM 合成的壁纸层
			progman, _, _ := procFindWindowW.Call(uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("Progman"))), 0)
			var discard uintptr
			procSendMsgTimeout.Call(progman, 0x052C, 0xD, 0, 0x0000, 1000, uintptr(unsafe.Pointer(&discard)))
			target, _, _ = procFindWindowExW.Call(progman, 0, uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("WorkerW"))), 0)
		default:
			target = findWorkerW()
		}
		if target == 0 {
			log.Fatal("embed target not found")
		}
		procSetParent.Call(wallpaperHwnd, target)
		procSetWindowLongPtrW.Call(wallpaperHwnd, GWL_STYLE, uintptr(WS_CHILD|WS_VISIBLE|WS_CLIPCHILDREN|WS_CLIPSIBLINGS))
		procMoveWindow.Call(wallpaperHwnd, 0, 0, uintptr(screenW), uintptr(screenH), 1)
		// *-top: 不置底（默认 z-order，盖住图标也认）用于验证子窗口能否上屏
		// 其它: 置底到图标之下
		if mode != "progman-top" && mode != "childww-top" {
			procSetWindowPos.Call(wallpaperHwnd, HWND_BOTTOM, 0, 0, 0, 0, SWP_NOMOVE|SWP_NOSIZE|SWP_FRAMECHANGED)
		}
		log.Printf("Embedded into target=%x mode=%s", target, mode)
	}

	initGL()
	loadExtensions()
	progID := setupShaders()

	log.Println("Rendering...")
	targetFPS := 30.0
	frameInterval := 1.0 / targetFPS
	lastFrame := 0.0
	lastMove := 0.0

	type msg struct {
		HWnd    uintptr
		Message uint32
		WParam  uintptr
		LParam  uintptr
		Time    uint32
		Pt      struct{ X, Y int32 }
	}
	var m msg
	for {
		for {
			r, _, _ := procPeekMessageW.Call(uintptr(unsafe.Pointer(&m)), 0, 0, 0, 1)
			if r == 0 {
				break
			}
			if m.Message == WM_QUIT {
				cleanup()
				return
			}
			switch m.Message {
			case WM_MOUSEMOVE:
				mouseX = float32(int16(m.LParam&0xFFFF)) / float32(screenW)
				mouseY = 1.0 - float32(int16(m.LParam>>16)) / float32(screenH)
				lastMove = now()
			case WM_LBUTTONDOWN:
				clickX = float32(int16(m.LParam&0xFFFF)) / float32(screenW)
				clickY = 1.0 - float32(int16(m.LParam>>16)) / float32(screenH)
				clickTime = now()
			}
			procTranslateMessage.Call(uintptr(unsafe.Pointer(&m)))
			procDispatchMessageW.Call(uintptr(unsafe.Pointer(&m)))
		}

		t := now()
		if t-lastFrame < frameInterval {
			time.Sleep(time.Duration((frameInterval - (t - lastFrame)) * 0.8 * float64(time.Second)))
			continue
		}
		lastFrame = t
		if t-lastMove > 5.0 {
			targetFPS = 10.0
		} else {
			targetFPS = 30.0
		}
		frameInterval = 1.0 / targetFPS

		// GL render
		glViewport.Call(0, 0, uintptr(screenW), uintptr(screenH))
		glClear.Call(uintptr(GL_COLOR_BUFFER_BIT))
		syscall.SyscallN(glUseProgram, uintptr(progID))

		elapsed := float32(0)
		if clickTime > 0 {
			e := t - clickTime
			if e < 3.0 {
				elapsed = float32(e)
			} else {
				clickTime = 0
			}
		}

		fTime := float32(t)
		fScreen := [2]float32{float32(screenW), float32(screenH)}
		syscall.SyscallN(glUniform1fv, uintptr(0), uintptr(1), uintptr(unsafe.Pointer(&fTime)))
		syscall.SyscallN(glUniform2fv, uintptr(1), uintptr(1), uintptr(unsafe.Pointer(&fScreen[0])))
		syscall.SyscallN(glUniform2fv, uintptr(2), uintptr(1), uintptr(unsafe.Pointer(&mouseX)))
		syscall.SyscallN(glUniform1fv, uintptr(3), uintptr(1), uintptr(unsafe.Pointer(&elapsed)))
		syscall.SyscallN(glUniform2fv, uintptr(4), uintptr(1), uintptr(unsafe.Pointer(&clickX)))

		glDrawArrays.Call(GL_TRIANGLE_STRIP, 0, 4)

		if useSwap {
			// 双缓冲呈现：走 DWM flip，嵌入子窗口也能合成上屏
			procSwapBuffers.Call(glDC)
		} else {
			glFinish.Call()
			// Read framebuffer → GDI blit to window
			glReadPixels.Call(0, 0, uintptr(screenW), uintptr(screenH), GL_RGBA, GL_UNSIGNED_BYTE, uintptr(unsafe.Pointer(&pixelBuf[0])))
			winDC, _, _ := procGetDC.Call(wallpaperHwnd)
			procSetDIBitsToDevice.Call(
				winDC,
				0, 0,
				uintptr(screenW), uintptr(screenH),
				0, 0,
				0, uintptr(screenH),
				uintptr(unsafe.Pointer(&pixelBuf[0])),
				uintptr(unsafe.Pointer(&bmi)),
				DIB_RGB_COLORS,
			)
			procReleaseDC.Call(wallpaperHwnd, winDC)
		}

		frameCount++
		if frameCount <= 3 {
			log.Printf("FRAME#%d swap=%v", frameCount, useSwap)
		}
	}
}

func loadExt(name string) uintptr {
	cname, _ := windows.BytePtrFromString(name)
	ptr, _, _ := procWglGetProcAddress.Call(uintptr(unsafe.Pointer(cname)))
	if ptr == 0 {
		log.Fatalf("ext not found: %s", name)
	}
	return ptr
}

func loadExtensions() {
	glUseProgram = loadExt("glUseProgram")
	glUniform1fv = loadExt("glUniform1fv")
	glUniform2fv = loadExt("glUniform2fv")
}

func setupShaders() uint32 {
	glCreateShader := loadExt("glCreateShader")
	glShaderSource := loadExt("glShaderSource")
	glCompileShader := loadExt("glCompileShader")
	glGetShaderiv := loadExt("glGetShaderiv")
	glGetShaderInfoLog := loadExt("glGetShaderInfoLog")

	compile := func(st uint32, src string) uint32 {
		s, _, _ := syscall.SyscallN(glCreateShader, uintptr(st))
		cstr, _ := windows.BytePtrFromString(src)
		p := uintptr(unsafe.Pointer(cstr))
		syscall.SyscallN(glShaderSource, s, 1, uintptr(unsafe.Pointer(&p)), 0)
		syscall.SyscallN(glCompileShader, s)
		var status int32
		syscall.SyscallN(glGetShaderiv, s, uintptr(GL_COMPILE_STATUS), uintptr(unsafe.Pointer(&status)))
		if status == 0 {
			var logLen int32
			syscall.SyscallN(glGetShaderiv, s, uintptr(GL_INFO_LOG_LENGTH), uintptr(unsafe.Pointer(&logLen)))
			if logLen > 0 {
				logBuf := make([]byte, logLen)
				syscall.SyscallN(glGetShaderInfoLog, s, uintptr(logLen), 0, uintptr(unsafe.Pointer(&logBuf[0])))
				log.Fatalf("shader compile failed:\n%s", string(logBuf[:logLen]))
			}
			log.Fatal("shader compile failed (no info log)")
		}
		return uint32(s)
	}
	vs := compile(GL_VERTEX_SHADER, vertexSrc)
	fs := compile(GL_FRAGMENT_SHADER, fragmentSrc)
	glCreateProgram := loadExt("glCreateProgram")
	prog, _, _ := syscall.SyscallN(glCreateProgram)
	syscall.SyscallN(loadExt("glAttachShader"), prog, uintptr(vs))
	syscall.SyscallN(loadExt("glAttachShader"), prog, uintptr(fs))
	syscall.SyscallN(loadExt("glLinkProgram"), prog)
	log.Println("Program:", prog)
	var vao, vbo uint32
	syscall.SyscallN(loadExt("glGenVertexArrays"), 1, uintptr(unsafe.Pointer(&vao)))
	syscall.SyscallN(loadExt("glGenBuffers"), 1, uintptr(unsafe.Pointer(&vbo)))
	quad := [8]float32{-1, -1, 1, -1, -1, 1, 1, 1}
	syscall.SyscallN(loadExt("glBindVertexArray"), uintptr(vao))
	syscall.SyscallN(loadExt("glBindBuffer"), uintptr(GL_ARRAY_BUFFER), uintptr(vbo))
	syscall.SyscallN(loadExt("glBufferData"), uintptr(GL_ARRAY_BUFFER), uintptr(len(quad)*4), uintptr(unsafe.Pointer(&quad[0])), uintptr(GL_STATIC_DRAW))
	syscall.SyscallN(loadExt("glEnableVertexAttribArray"), 0)
	syscall.SyscallN(loadExt("glVertexAttribPointer"), 0, 2, uintptr(GL_FLOAT), 0, 0, 0)
	return uint32(prog)
}

func createWindow() {
	var hinstance windows.Handle
	windows.GetModuleHandleEx(0, nil, &hinstance)
	className, _ := windows.UTF16PtrFromString("GLWallpaper")
	wndProc := windows.NewCallback(func(hwnd uintptr, msg uint32, wp, lp uintptr) uintptr {
		if msg == WM_DESTROY || msg == WM_CLOSE {
			procPostQuitMessage.Call(0)
			return 0
		}
		ret, _, _ := procDefWindowProcW.Call(hwnd, uintptr(msg), wp, lp)
		return ret
	})
	wc := wndClassEx{
		CbSize:        uint32(unsafe.Sizeof(wndClassEx{})),
		LpfnWndProc:   wndProc,
		HInstance:     uintptr(hinstance),
		LpszClassName: className,
	}
	procRegisterClassExW.Call(uintptr(unsafe.Pointer(&wc)))
	windowName, _ := windows.UTF16PtrFromString("GLWallpaper")
	wallpaperHwnd, _, _ = procCreateWindowExW.Call(
		0, uintptr(unsafe.Pointer(className)), uintptr(unsafe.Pointer(windowName)),
		WS_POPUP|WS_VISIBLE, 0, 0, uintptr(screenW), uintptr(screenH),
		0, 0, uintptr(hinstance), 0,
	)
	if wallpaperHwnd == 0 {
		log.Fatal("CreateWindow failed")
	}
	procShowWindow.Call(wallpaperHwnd, 5)
	log.Println("Window:", wallpaperHwnd)
}

func initGL() {
	glDC, _, _ = procGetDC.Call(wallpaperHwnd)
	if glDC == 0 {
		log.Fatal("GetDC failed")
	}
	flags := uint32(PFD_DRAW_TO_WINDOW | PFD_SUPPORT_OPENGL)
	if useSwap {
		flags |= PFD_DOUBLEBUFFER
	}
	desc := pfd{Size: uint16(unsafe.Sizeof(pfd{})), Version: 1, Flags: flags, PixelType: 0, ColorBits: 32, DepthBits: 24}
	pf, _, _ := procChoosePixelFormat.Call(glDC, uintptr(unsafe.Pointer(&desc)))
	if pf == 0 {
		log.Fatal("ChoosePixelFormat failed")
	}
	ok, _, _ := procSetPixelFormat.Call(glDC, pf, uintptr(unsafe.Pointer(&desc)))
	if ok == 0 {
		log.Fatal("SetPixelFormat failed")
	}
	glCtx, _, _ = procWglCreateContext.Call(glDC)
	if glCtx == 0 {
		log.Fatal("wglCreateContext failed")
	}
	mcR, _, _ := procWglMakeCurrent.Call(glDC, glCtx)
	ver, _, _ := glGetString.Call(GL_VERSION)
	log.Printf("MakeCurrent=%d GL=%s", mcR, windows.BytePtrToString((*byte)(unsafe.Pointer(ver))))
}

func cleanup() {
	if glCtx != 0 {
		procWglMakeCurrent.Call(0, 0)
		procWglDeleteContext.Call(glCtx)
	}
	if glDC != 0 {
		procReleaseDC.Call(wallpaperHwnd, glDC)
	}
}

func classOf(hwnd uintptr) string {
	buf := make([]uint16, 256)
	procGetClassNameW.Call(hwnd, uintptr(unsafe.Pointer(&buf[0])), 256)
	return windows.UTF16ToString(buf)
}

func enumTops(tag string) {
	cb := windows.NewCallback(func(hwnd, _ uintptr) uintptr {
		cls := classOf(hwnd)
		if cls == "WorkerW" || cls == "Progman" {
			dv, _, _ := procFindWindowExW.Call(hwnd, 0, uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("SHELLDLL_DefView"))), 0)
			var r rect
			procGetWindowRect.Call(hwnd, uintptr(unsafe.Pointer(&r)))
			full := ""
			if r.Right-r.Left == screenW && r.Bottom-r.Top == screenH {
				full = " <FULLSCREEN>"
			}
			log.Printf("[%s] hwnd=%x class=%-8s DefView=%x rect=(%d,%d,%d,%d)%s", tag, hwnd, cls, dv, r.Left, r.Top, r.Right, r.Bottom, full)
		}
		return 1
	})
	procEnumWindows.Call(cb, 0)
}

// probe 用权威参数 WPARAM=0xD 发消息，多轮观察 WorkerW 是否分离出来
func probe() {
	screenW, screenH = getDesktopSize()
	progman, _, _ := procFindWindowW.Call(uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("Progman"))), 0)
	log.Printf("Progman=%x screen=%dx%d", progman, screenW, screenH)
	var discard uintptr
	procSendMsgTimeout.Call(progman, 0x052C, 0xD, 0, 0x0000, 1000, uintptr(unsafe.Pointer(&discard)))
	procSendMsgTimeout.Call(progman, 0x052C, 0xD, 1, 0x0000, 1000, uintptr(unsafe.Pointer(&discard)))
	log.Println("sent 0x052C WPARAM=0xD (LPARAM 0 and 1)")

	for i := 0; i < 4; i++ {
		time.Sleep(300 * time.Millisecond)
		log.Printf("--- round %d ---", i)
		enumTops("top")
	}
}

func findWorkerW() uintptr {
	progman, _, _ := procFindWindowW.Call(uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("Progman"))), 0)
	if progman == 0 {
		return 0
	}
	var discard uintptr
	procSendMsgTimeout.Call(progman, 0x052C, 0, 0, 0x0000, 1000, uintptr(unsafe.Pointer(&discard)))
	sdv, _, _ := procFindWindowExW.Call(progman, 0, uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("SHELLDLL_DefView"))), 0)
	if sdv != 0 {
		ww, _, _ := procFindWindowExW.Call(progman, sdv, uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("WorkerW"))), 0)
		if ww != 0 {
			return ww
		}
	}
	ww, _, _ := procFindWindowExW.Call(progman, 0, uintptr(unsafe.Pointer(windows.StringToUTF16Ptr("WorkerW"))), 0)
	return ww
}

func getDesktopSize() (int32, int32) {
	var mr, mb int32
	cb := windows.NewCallback(func(_, _, lprc, _ uintptr) uintptr {
		r := (*rect)(unsafe.Pointer(lprc))
		if r.Right > mr {
			mr = r.Right
		}
		if r.Bottom > mb {
			mb = r.Bottom
		}
		return 1
	})
	procEnumDisplayMonitors.Call(0, 0, cb, 0)
	return mr, mb
}

func now() float64 { return float64(time.Now().UnixNano())/1e9 - startTime }