//go:build windows

package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"
	"unsafe"
)

const (
	CreateToolhelp32SnapshotProcess = 0x00000002
	ProcessSetInformation           = 0x00000200
	ProcessQueryInformation         = 0x00000400
	PriorityClassIdle               = 0x00000040
	PriorityClassNormal             = 0x00000020
	PriorityClassHigh               = 0x00000080
	PriorityClassAboveNormal        = 0x00008000
	PriorityClassRealtime           = 0x00000100
	PriorityClassBelowNormal        = 0x00004000
	MessageBoxIconWarning           = 0x00000030
	ErrorAlreadyExists              = 183
	ErrorNoMoreFiles                = 18
	ErrorAccessDenied             = 5
	ErrorInvalidParameter         = 87
)

type processEntry32 struct {
	Size          uint32
	CntUsage      uint32
	PID           uint32
	DefaultHeapID uintptr
	ModuleID      uint32
	CntThreads    uint32
	ParentPID     uint32
	PrioClass     int32
	Flags         uint32
	ExeFile       [260]uint16
}

type processMemoryCounters struct {
	CBM                      uint32
	PageFaultCount           uint32
	PeakWorkingSetSize       uintptr
	WorkingSetSize           uintptr
	QuotaPeakPagedPoolUsage  uintptr
	QuotaPagedPoolUsage      uintptr
	QuotaPeakNonPagedPoolUsage  uintptr
	QuotaNonPagedPoolUsage    uintptr
	PeakPagefileUsage        uintptr
	PagefileUsage            uintptr
	PrivateUsage             uintptr
}

type processInfo struct {
	PID  uint32
	Name string
}

var (
	k32 = syscall.NewLazyDLL("kernel32.dll")
	ps  = syscall.NewLazyDLL("psapi.dll")

	procCreateSnap     = k32.NewProc("CreateToolhelp32Snapshot")
	procProcess32First = k32.NewProc("Process32FirstW")
	procProcess32Next  = k32.NewProc("Process32NextW")
	procGetMemInfo     = ps.NewProc("GetProcessMemoryInfo")
	procSetPriority    = k32.NewProc("SetPriorityClass")
	procGetPriority    = k32.NewProc("GetPriorityClass")
	procOpenProcess    = k32.NewProc("OpenProcess")
	procCloseHandle    = k32.NewProc("CloseHandle")
	procCreateMutex    = k32.NewProc("CreateMutexW")
	procMessageBox     = k32.NewProc("MessageBoxW")
)

func closeH(h syscall.Handle) {
	_, _, _ = procCloseHandle.Call(uintptr(h))
}

func openProc(pid uint32, acc uint32) (syscall.Handle, error) {
	r, _, e := procOpenProcess.Call(uintptr(acc), 0, uintptr(pid))
	if r == 0 {
		return 0, fmt.Errorf("OpenProcess(%d) failed: %w", pid, e)
	}
	return syscall.Handle(r), nil
}

func setPrio(pid uint32, cls uint32) error {
	h, err := openProc(pid, ProcessSetInformation)
	if err != nil {
		return err
	}
	defer closeH(h)
	r, _, e := procSetPriority.Call(uintptr(h), uintptr(cls))
	if r == 0 {
		return fmt.Errorf("SetPriorityClass(%d) failed: %w", pid, e)
	}
	return nil
}

func ignoreOpenError(err error) bool {
	return errors.Is(err, syscall.Errno(ErrorAccessDenied)) ||
		errors.Is(err, syscall.Errno(ErrorInvalidParameter))
}

func curPrioWithHandle(h syscall.Handle) (uint32, error) {
	r, _, e := procGetPriority.Call(uintptr(h))
	if r == 0 {
		return 0, fmt.Errorf("GetPriorityClass failed: %w", e)
	}
	return uint32(r), nil
}

func utf16Trim(p []uint16) string {
	for i, v := range p {
		if v == 0 {
			return syscall.UTF16ToString(p[:i])
		}
	}
	return syscall.UTF16ToString(p)
}

func allProcs() ([]processInfo, error) {
	snap, _, e := procCreateSnap.Call(CreateToolhelp32SnapshotProcess, 0)
	if snap == 0 {
		return nil, fmt.Errorf("CreateToolhelp32Snapshot failed: %w", e)
	}
	defer closeH(syscall.Handle(snap))

	pe := processEntry32{Size: uint32(unsafe.Sizeof(processEntry32{}))}
	r, _, e := procProcess32First.Call(snap, uintptr(unsafe.Pointer(&pe)))
	if r == 0 {
		return nil, fmt.Errorf("Process32First failed: %w", e)
	}

	var out []processInfo
	for {
		name := utf16Trim(pe.ExeFile[:])
		if name != "" {
			out = append(out, processInfo{PID: pe.PID, Name: name})
		}
		pe.Size = uint32(unsafe.Sizeof(processEntry32{}))
		r, _, e = procProcess32Next.Call(snap, uintptr(unsafe.Pointer(&pe)))
		if r == 0 {
			if e == syscall.Errno(ErrorNoMoreFiles) {
				break
			}
			return nil, fmt.Errorf("Process32Next failed: %w", e)
		}
	}
	return out, nil
}

func isAboveNormal(c uint32) bool {
	return c == PriorityClassAboveNormal ||
		c == PriorityClassHigh ||
		c == PriorityClassRealtime
}

func prioName(c uint32) string {
	switch c {
	case PriorityClassIdle:
		return "IDLE"
	case PriorityClassBelowNormal:
		return "BELOW_NORMAL"
	case PriorityClassNormal:
		return "NORMAL"
	case PriorityClassAboveNormal:
		return "ABOVE_NORMAL"
	case PriorityClassHigh:
		return "HIGH"
	case PriorityClassRealtime:
		return "REALTIME"
	default:
		return fmt.Sprintf("0x%X", c)
	}
}

func procName(pid uint32) (string, error) {
	snap, _, e := procCreateSnap.Call(CreateToolhelp32SnapshotProcess, 0)
	if snap == 0 {
		return "", fmt.Errorf("CreateToolhelp32Snapshot failed: %w", e)
	}
	defer closeH(syscall.Handle(snap))

	pe := processEntry32{Size: uint32(unsafe.Sizeof(processEntry32{}))}
	r, _, e := procProcess32First.Call(snap, uintptr(unsafe.Pointer(&pe)))
	if r == 0 {
		return "", fmt.Errorf("Process32First failed: %w", e)
	}

	for {
		if pe.PID == pid {
			return utf16Trim(pe.ExeFile[:]), nil
		}
		pe.Size = uint32(unsafe.Sizeof(processEntry32{}))
		r, _, e = procProcess32Next.Call(snap, uintptr(unsafe.Pointer(&pe)))
		if r == 0 {
			if e == syscall.Errno(ErrorNoMoreFiles) {
				break
			}
			return "", fmt.Errorf("Process32Next failed: %w", e)
		}
	}
	return "", nil
}

func isAlive(pid uint32, expectedName string) bool {
	h, err := openProc(pid, ProcessQueryInformation)
	if err != nil || h == 0 {
		return false
	}
	closeH(h)
	name, err := procName(pid)
	if err != nil || name == "" {
		return false
	}
	return name == expectedName
}

func ensureSingleInstance() syscall.Handle {
	h, _, err := procCreateMutex.Call(0, 0, uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("Global\\autopriority_v1"))))
	if h == 0 {
		fmt.Fprintf(os.Stderr, "autoPriority: failed to create mutex: %v\n", err)
		os.Exit(1)
	}
	if err == syscall.Errno(ErrorAlreadyExists) {
		procMessageBox.Call(0,
			uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("autoPriority is already running."))),
			uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr("autoPriority"))),
			MessageBoxIconWarning,
		)
		os.Exit(1)
	}
	return syscall.Handle(h)
}

func main() {
	mem := flag.Uint64("mem", 512*1024*1024, "memory threshold in bytes")
	interval := flag.Duration("interval", time.Minute, "scan interval")
	dryRun := flag.Bool("dry-run", false, "log only, do not change priorities")
	flag.Parse()

	if *interval < 10*time.Second {
		*interval = 10 * time.Second
	}

	var logFile *os.File
	logPath := filepath.Join(os.TempDir(), "autopriority.log")
	os.Remove(logPath)
	f, err := os.OpenFile(logPath,
		os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0644)
	if err == nil {
		logFile = f
		defer logFile.Close()
	}

	ts := func() string {
		return time.Now().Format("02.01.2006 15:04:05")
	}

	logf := func(format string, a ...any) {
		if logFile == nil {
			return
		}
		fmt.Fprintf(logFile, "[%s] ", ts())
		fmt.Fprintf(logFile, format+"\n", a...)
		logFile.Sync()
	}

	logf("autoPriority started (mem=%d bytes, interval=%s, dry-run=%v)", *mem, *interval, *dryRun)

	mutex := ensureSingleInstance()
	defer closeH(mutex)

	if err := setPrio(uint32(os.Getpid()), PriorityClassIdle); err != nil {
		logf("warning: could not set own priority to IDLE: %v", err)
	} else {
		logf("own priority set to IDLE")
	}

	promoted := make(map[uint32]string)
	blocked := make(map[uint32]string)

	defer func() {
		for pid, name := range promoted {
			cur, err := procName(pid)
			if err != nil || cur != name {
				logf("SKIP RESTORE PID %d: process name mismatch (expected %s, got %s)", pid, name, cur)
				continue
			}
			if err := setPrio(pid, PriorityClassNormal); err != nil {
				logf("RESTORE %s (PID %d) -> NORMAL error: %v", name, pid, err)
			} else {
				logf("RESTORE %s (PID %d) -> NORMAL", name, pid)
			}
		}
		logf("autoPriority stopped")
	}()

	myPID := uint32(os.Getpid())
	ticker := time.NewTicker(*interval)
	defer ticker.Stop()

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)

	scan := func() {
		procs, err := allProcs()
		if err != nil {
			logf("allProcs error: %v", err)
			return
		}

		for pid, name := range promoted {
			if !isAlive(pid, name) {
				delete(promoted, pid)
			}
		}

		for pid, name := range blocked {
			if !isAlive(pid, name) {
				delete(blocked, pid)
			}
		}

		for _, p := range procs {
			if p.PID == myPID || p.PID == 0 {
				continue
			}
			if _, ok := promoted[p.PID]; ok {
				continue
			}
			if _, ok := blocked[p.PID]; ok {
				continue
			}

			h, err := openProc(p.PID, ProcessQueryInformation|ProcessSetInformation)
			if err != nil {
				if !ignoreOpenError(err) {
					logf("openProc(%s, PID %d) error: %v", p.Name, p.PID, err)
				}
				continue
			}

			var m processMemoryCounters
			m.CBM = uint32(unsafe.Sizeof(m))
			r, _, e := procGetMemInfo.Call(uintptr(h), uintptr(unsafe.Pointer(&m)), uintptr(unsafe.Sizeof(m)))
			if r == 0 {
				closeH(h)
				logf("GetProcessMemoryInfo(%d) error: %v", p.PID, e)
				continue
			}
			rssBytes := uint64(m.WorkingSetSize)
			if rssBytes < *mem {
				origPrio, prioErr := curPrioWithHandle(h)
				if prioErr == nil && isAboveNormal(origPrio) {
					if *dryRun {
						logf("[DRY-RUN] %s (PID %d) RSS=%dMB, priority=%s — would set NORMAL",
							p.Name, p.PID, rssBytes/1024/1024, prioName(origPrio))
					} else {
						sr, _, se := procSetPriority.Call(uintptr(h), uintptr(PriorityClassNormal))
						if sr != 0 {
							logf("DEMOTE %s (PID %d) RSS=%dMB, %s -> NORMAL",
								p.Name, p.PID, rssBytes/1024/1024, prioName(origPrio))
						} else {
							blocked[p.PID] = p.Name
logf("BLOCK %s (PID %d): %v (added to exclusion list)", p.Name, p.PID, se)
						}
					}
				}
				closeH(h)
				continue
			}

			rssMB := rssBytes / 1024 / 1024
			origPrio, prioErr := curPrioWithHandle(h)
			if prioErr != nil {
				closeH(h)
				continue
			}
			if origPrio == PriorityClassHigh {
				closeH(h)
				continue
			}
			origStr := prioName(origPrio)

			if *dryRun {
				closeH(h)
				logf("[DRY-RUN] %s (PID %d) RSS=%dMB, priority=%s — would set HIGH",
					p.Name, p.PID, rssMB, origStr)
				continue
			}

			sr, _, se := procSetPriority.Call(uintptr(h), uintptr(PriorityClassHigh))
			closeH(h)
			if sr == 0 {
				blocked[p.PID] = p.Name
				logf("BLOCK %s (PID %d): %v (added to exclusion list)", p.Name, p.PID, se)
				continue
			}

			promoted[p.PID] = p.Name
			logf("PROMOTE %s (PID %d) RSS=%dMB, %s -> HIGH",
				p.Name, p.PID, rssMB, origStr)
		}
	}

	scan()
	for {
		select {
		case <-stop:
			logf("received shutdown signal")
			return
		case <-ticker.C:
			scan()
		}
	}
}