//go:build windows

package main

import (
	"bytes"
	"fmt"
	"image/png"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
	"unsafe"
)

const (
	trayCallbackMessage = 0x8001
	trayQuitMessage     = 0x8002
	trayIconID          = 1
	menuExitID          = 1
)

type winPoint struct{ X, Y int32 }
type winRect struct{ Left, Top, Right, Bottom int32 }
type winMessage struct {
	Window         uintptr
	Message        uint32
	WParam, LParam uintptr
	Time           uint32
	Point          winPoint
	Private        uint32
}
type windowClass struct {
	Size, Style                        uint32
	Proc                               uintptr
	ClassExtra, WindowExtra            int32
	Instance, Icon, Cursor, Background uintptr
	MenuName, ClassName                *uint16
	SmallIcon                          uintptr
}
type notifyIconData struct {
	Size                uint32
	Window              uintptr
	ID, Flags, Callback uint32
	Icon                uintptr
	Tip                 [128]uint16
	State, StateMask    uint32
	Info                [256]uint16
	Version             uint32
	InfoTitle           [64]uint16
	InfoFlags           uint32
	GUID                [16]byte
	BalloonIcon         uintptr
}
type notifyIconIdentifier struct {
	Size   uint32
	Window uintptr
	ID     uint32
	GUID   [16]byte
}
type menuItemInfo struct {
	Size, Mask, Type, State, ID       uint32
	Submenu, Checked, Unchecked, Data uintptr
	Text                              *uint16
	TextLen                           uint32
	Bitmap                            uintptr
}
type bitmapHeader struct {
	Size                        uint32
	Width, Height               int32
	Planes, BitCount            uint16
	Compression, ImageSize      uint32
	XPels, YPels                int32
	ColorsUsed, ColorsImportant uint32
}
type iconInfo struct {
	Icon               int32
	HotspotX, HotspotY uint32
	Mask, Color        uintptr
}

func systemDLL(name string) *syscall.LazyDLL {
	var dir [32768]uint16
	r, _, err := k32.NewProc("GetSystemDirectoryW").Call(uintptr(unsafe.Pointer(&dir[0])), uintptr(len(dir)))
	if r == 0 || r >= uintptr(len(dir)) {
		panic(fmt.Sprintf("GetSystemDirectoryW: %v", err))
	}
	return syscall.NewLazyDLL(filepath.Join(syscall.UTF16ToString(dir[:r]), name))
}

var (
	u32               = systemDLL("user32.dll")
	shell32           = systemDLL("shell32.dll")
	gdi32             = systemDLL("gdi32.dll")
	pRegisterClass    = u32.NewProc("RegisterClassExW")
	pUnregisterClass  = u32.NewProc("UnregisterClassW")
	pCreateWindow     = u32.NewProc("CreateWindowExW")
	pDestroyWindow    = u32.NewProc("DestroyWindow")
	pDefWindowProc    = u32.NewProc("DefWindowProcW")
	pGetMessage       = u32.NewProc("GetMessageW")
	pTranslateMessage = u32.NewProc("TranslateMessage")
	pDispatchMessage  = u32.NewProc("DispatchMessageW")
	pPostQuit         = u32.NewProc("PostQuitMessage")
	pPostMessage      = u32.NewProc("PostMessageW")
	pCreatePopup      = u32.NewProc("CreatePopupMenu")
	pDestroyMenu      = u32.NewProc("DestroyMenu")
	pInsertMenuItem   = u32.NewProc("InsertMenuItemW")
	pTrackPopup       = u32.NewProc("TrackPopupMenu")
	pForeground       = u32.NewProc("SetForegroundWindow")
	pEndMenu          = u32.NewProc("EndMenu")
	pCursorPos        = u32.NewProc("GetCursorPos")
	pRegisterMessage  = u32.NewProc("RegisterWindowMessageW")
	pMessageFilter    = u32.NewProc("ChangeWindowMessageFilterEx")
	pCreateIcon       = u32.NewProc("CreateIconIndirect")
	pDestroyIcon      = u32.NewProc("DestroyIcon")
	pCreateDIB        = gdi32.NewProc("CreateDIBSection")
	pCreateBitmap     = gdi32.NewProc("CreateBitmap")
	pDeleteObject     = gdi32.NewProc("DeleteObject")
	pShellNotify      = shell32.NewProc("Shell_NotifyIconW")
	pNotifyRect       = shell32.NewProc("Shell_NotifyIconGetRect")
)

type tray struct {
	hwnd     atomic.Uintptr
	menu     uintptr
	data     notifyIconData
	taskbar  uint32
	version4 bool
	menuOpen atomic.Bool
	done     chan struct{}
	err      error
	failure  error // message-loop thread only
	onExit   func()
	restored <-chan struct{}
	exitOnce sync.Once
	added    bool
}

func newTray(onExit func(), restored <-chan struct{}) (*tray, error) {
	t := &tray{done: make(chan struct{}), onExit: onExit, restored: restored}
	ready := make(chan error, 1)
	go func() {
		runtime.LockOSThread()
		defer runtime.UnlockOSThread()
		t.err = t.loop(ready)
		t.requestExit()
		close(t.done)
	}()
	if err := <-ready; err != nil {
		<-t.done
		return nil, err
	}
	return t, nil
}
func (t *tray) Close() error {
	select {
	case <-t.done:
		return t.err
	default:
	}
	if h := t.hwnd.Load(); h != 0 {
		r, _, e := pPostMessage.Call(h, trayQuitMessage, 0, 0)
		if r == 0 {
			return fmt.Errorf("PostMessage(close): %v", e)
		}
	}
	<-t.done
	return t.err
}

var trayWindows sync.Map
var trayWindowCallback = syscall.NewCallback(trayWindowProc)

func (t *tray) requestExit() {
	t.exitOnce.Do(func() {
		if t.onExit != nil {
			t.onExit()
		}
	})
}
func (t *tray) fail(err error) {
	if t.failure == nil {
		t.failure = err
	}
	t.requestExit()
}
func pngBitmap(data []byte) (uintptr, error) {
	image, err := png.Decode(bytes.NewReader(data))
	if err != nil {
		return 0, err
	}
	if image.Bounds().Dx() != 16 || image.Bounds().Dy() != 16 {
		return 0, fmt.Errorf("expected a 16x16 Silk icon")
	}
	header := bitmapHeader{Size: uint32(unsafe.Sizeof(bitmapHeader{})), Width: 16, Height: -16, Planes: 1, BitCount: 32}
	var bits unsafe.Pointer
	bitmap, _, e := pCreateDIB.Call(0, uintptr(unsafe.Pointer(&header)), 0, uintptr(unsafe.Pointer(&bits)), 0, 0)
	if bitmap == 0 {
		return 0, fmt.Errorf("CreateDIBSection: %v", e)
	}
	if bits == nil {
		pDeleteObject.Call(bitmap)
		return 0, fmt.Errorf("CreateDIBSection returned no pixels")
	}
	pixels := unsafe.Slice((*byte)(bits), 16*16*4)
	for y := 0; y < 16; y++ {
		for x := 0; x < 16; x++ {
			r, g, b, a := image.At(x, y).RGBA()
			offset := (y*16 + x) * 4
			// RGBA() is premultiplied; native menu bitmaps require premultiplied BGRA.
			pixels[offset] = byte(b >> 8)
			pixels[offset+1] = byte(g >> 8)
			pixels[offset+2] = byte(r >> 8)
			pixels[offset+3] = byte(a >> 8)
		}
	}
	return bitmap, nil
}
func pngIcon(data []byte) (uintptr, error) {
	color, err := pngBitmap(data)
	if err != nil {
		return 0, err
	}
	defer pDeleteObject.Call(color)
	var maskBits [32]byte
	mask, _, e := pCreateBitmap.Call(16, 16, 1, 1, uintptr(unsafe.Pointer(&maskBits[0])))
	if mask == 0 {
		return 0, fmt.Errorf("CreateBitmap(mask): %v", e)
	}
	defer pDeleteObject.Call(mask)
	info := iconInfo{Icon: 1, Mask: mask, Color: color}
	icon, _, e := pCreateIcon.Call(uintptr(unsafe.Pointer(&info)))
	if icon == 0 {
		return 0, fmt.Errorf("CreateIconIndirect: %v", e)
	}
	return icon, nil
}
func (t *tray) addIcon() error {
	t.data.Flags = 0x87                                                // NIF_MESSAGE | NIF_ICON | NIF_TIP | NIF_SHOWTIP
	ok, _, _ := pShellNotify.Call(0, uintptr(unsafe.Pointer(&t.data))) // NIM_ADD
	if ok == 0 {
		ok, _, _ = pShellNotify.Call(1, uintptr(unsafe.Pointer(&t.data)))
	} // existing icon: NIM_MODIFY
	if ok == 0 {
		return fmt.Errorf("Shell_NotifyIcon: notification area unavailable")
	}
	t.added = true
	t.data.Version = 4
	ok, _, _ = pShellNotify.Call(4, uintptr(unsafe.Pointer(&t.data))) // NIM_SETVERSION, after EVERY add
	t.version4 = ok != 0
	return nil
}
func (t *tray) removeIcon() {
	if t.added {
		pShellNotify.Call(2, uintptr(unsafe.Pointer(&t.data)))
		t.added = false
	}
}
func (t *tray) loop(ready chan<- error) (result error) {
	sent := false
	defer func() {
		if !sent {
			ready <- result
		}
	}()
	icon, err := pngIcon(trayPNG)
	if err != nil {
		return err
	}
	defer pDestroyIcon.Call(icon)
	cross, err := pngBitmap(exitPNG)
	if err != nil {
		return err
	}
	defer pDeleteObject.Call(cross)
	menu, _, e := pCreatePopup.Call()
	if menu == 0 {
		return fmt.Errorf("CreatePopupMenu: %v", e)
	}
	t.menu = menu
	defer pDestroyMenu.Call(menu)
	label, _ := syscall.UTF16PtrFromString("Выход")
	item := menuItemInfo{Size: uint32(unsafe.Sizeof(menuItemInfo{})), Mask: 0x2 | 0x40 | 0x80, ID: menuExitID, Text: label, Bitmap: cross}
	ok, _, e := pInsertMenuItem.Call(menu, 0, 1, uintptr(unsafe.Pointer(&item)))
	if ok == 0 {
		return fmt.Errorf("InsertMenuItemW: %v", e)
	}
	instance, _, e := k32.NewProc("GetModuleHandleW").Call(0)
	if instance == 0 {
		return fmt.Errorf("GetModuleHandleW: %v", e)
	}
	className, _ := syscall.UTF16PtrFromString(fmt.Sprintf("autoPriority.Tray.%d", os.Getpid()))
	wc := windowClass{Size: uint32(unsafe.Sizeof(windowClass{})), Proc: trayWindowCallback, Instance: instance, ClassName: className, Icon: icon, SmallIcon: icon}
	atom, _, e := pRegisterClass.Call(uintptr(unsafe.Pointer(&wc)))
	if atom == 0 {
		return fmt.Errorf("RegisterClassExW: %v", e)
	}
	defer pUnregisterClass.Call(uintptr(unsafe.Pointer(className)), instance)
	// Hidden TOP-LEVEL window: a message-only window would miss TaskbarCreated.
	h, _, e := pCreateWindow.Call(0, uintptr(unsafe.Pointer(className)), 0, 0, 0, 0, 0, 0, 0, 0, instance, 0)
	if h == 0 {
		return fmt.Errorf("CreateWindowExW: %v", e)
	}
	t.hwnd.Store(h)
	trayWindows.Store(h, t)
	defer func() {
		if t.hwnd.Load() != 0 {
			pDestroyWindow.Call(h)
		}
		trayWindows.Delete(h)
	}()
	taskbarName, _ := syscall.UTF16PtrFromString("TaskbarCreated")
	taskbar, _, e := pRegisterMessage.Call(uintptr(unsafe.Pointer(taskbarName)))
	if taskbar == 0 {
		return fmt.Errorf("RegisterWindowMessageW: %v", e)
	}
	t.taskbar = uint32(taskbar)
	// Only the Explorer restart message crosses an elevated window's UIPI boundary.
	ok, _, e = pMessageFilter.Call(h, taskbar, 1, 0)
	if ok == 0 {
		return fmt.Errorf("ChangeWindowMessageFilterEx: %v", e)
	}
	t.data = notifyIconData{Size: uint32(unsafe.Sizeof(notifyIconData{})), Window: h, ID: trayIconID, Callback: trayCallbackMessage, Icon: icon}
	tip, _ := syscall.UTF16FromString("autoPriority")
	copy(t.data.Tip[:], tip)
	if err = t.addIcon(); err != nil {
		return err
	}
	defer t.removeIcon()
	ready <- nil
	sent = true
	for {
		var msg winMessage
		got, _, e := pGetMessage.Call(uintptr(unsafe.Pointer(&msg)), 0, 0, 0)
		if int32(got) == -1 {
			return fmt.Errorf("GetMessageW: %v", e)
		}
		if got == 0 {
			return t.failure
		}
		pTranslateMessage.Call(uintptr(unsafe.Pointer(&msg)))
		pDispatchMessage.Call(uintptr(unsafe.Pointer(&msg)))
	}
}
func trayEvent(version4 bool, wp, lp uintptr) (event, id uint32, point winPoint, anchored bool) {
	event = uint32(lp)
	id = uint32(wp)
	if version4 {
		event = uint32(lp) & 0xffff
		id = uint32(lp>>16) & 0xffff
		if event == 0x400 || event == 0x401 || (event >= 0x200 && event <= 0x20e) {
			point = winPoint{int32(int16(wp & 0xffff)), int32(int16((wp >> 16) & 0xffff))}
			anchored = true
		}
	}
	return
}
func (t *tray) popup(point winPoint, anchored bool) {
	if t.menuOpen.Swap(true) {
		return
	}
	defer t.menuOpen.Store(false)
	if !anchored {
		ok, _, _ := pCursorPos.Call(uintptr(unsafe.Pointer(&point)))
		if ok == 0 {
			point = winPoint{}
		}
	}
	h := t.hwnd.Load()
	if h == 0 {
		return
	}
	pForeground.Call(h)
	k32.NewProc("SetLastError").Call(0)
	command, _, e := pTrackPopup.Call(t.menu, 0x100|0x80|0x2, uintptr(point.X), uintptr(point.Y), 0, h, 0)
	pPostMessage.Call(h, 0, 0, 0) // documented focus workaround after TrackPopupMenu
	if command == menuExitID {
		t.requestExit()
	} else if command == 0 && e != syscall.Errno(0) {
		t.fail(fmt.Errorf("TrackPopupMenu: %v", e))
	}
	pShellNotify.Call(3, uintptr(unsafe.Pointer(&t.data))) // NIM_SETFOCUS
}
func trayWindowProc(h uintptr, message uint32, wp, lp uintptr) uintptr {
	value, ok := trayWindows.Load(h)
	if !ok {
		r, _, _ := pDefWindowProc.Call(h, uintptr(message), wp, lp)
		return r
	}
	t := value.(*tray)
	if message == t.taskbar {
		if err := t.addIcon(); err != nil {
			t.fail(err)
		}
		return 0
	}
	switch message {
	case trayCallbackMessage:
		event, id, point, anchored := trayEvent(t.version4, wp, lp)
		if id == trayIconID && trayActivation(t.version4, event) {
			t.popup(point, anchored)
		}
		return 0
	case 0x111:
		if uint32(wp)&0xffff == menuExitID {
			t.requestExit()
		}
		return 0 // WM_COMMAND
	case 0x10:
		t.requestExit()
		return 0 // WM_CLOSE requests cleanup, not forceful exit
	case 0x11:
		return 1 // WM_QUERYENDSESSION
	case 0x16: // WM_ENDSESSION: give the controller a bounded chance to restore
		if wp != 0 {
			t.requestExit()
			select {
			case <-t.restored:
			case <-time.After(3 * time.Second):
			}
		}
		return 0
	case 0x1f:
		pEndMenu.Call()
		return 0 // WM_CANCELMODE
	case trayQuitMessage:
		pEndMenu.Call()
		pDestroyWindow.Call(h)
		return 0
	case 0x2:
		t.removeIcon()
		t.hwnd.Store(0)
		t.requestExit()
		pPostQuit.Call(0)
		return 0 // WM_DESTROY
	case 0x82:
		trayWindows.Delete(h) // WM_NCDESTROY
	}
	r, _, _ := pDefWindowProc.Call(h, uintptr(message), wp, lp)
	return r
}
func showError(err error) {
	console, _, _ := k32.NewProc("GetConsoleWindow").Call()
	if console != 0 {
		return
	}
	text, _ := syscall.UTF16PtrFromString("autoPriority: " + err.Error())
	title, _ := syscall.UTF16PtrFromString("autoPriority")
	u32.NewProc("MessageBoxW").Call(0, uintptr(unsafe.Pointer(text)), uintptr(unsafe.Pointer(title)), 0x10)
}