//go:build windows

package main

import (
	"flag"
	"fmt"
	"math"
	"os"
	"os/signal"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"
	"unsafe"
)

const (
	CreateToolhelp32SnapshotProcess = 0x00000002
	ProcessSetInformation           = 0x00000200
	ProcessQueryLimitedInformation  = 0x00001000
	PriorityClassIdle               = 0x00000040
	PriorityClassNormal             = 0x00000020
	PriorityClassHigh               = 0x00000080
	PriorityClassAboveNormal        = 0x00008000
	PriorityClassRealtime           = 0x00000100
	PriorityClassBelowNormal        = 0x00004000
	ErrorNoMoreFiles                = 18
	ErrorAlreadyExists              = 183
)

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 processKey struct {
	PID     uint32
	Created uint64
}

type procInfo struct {
	PID    uint32
	Name   string
	RSS    uint64
	Key    processKey
	Prio   uint32
	Handle syscall.Handle
	CanSet bool
	Keep   bool
}

type trackedProc struct {
	name   string
	handle syscall.Handle
}

type savedPrio struct {
	name   string
	prio   uint32
	handle syscall.Handle
}

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")
	procGetProcessTimes = k32.NewProc("GetProcessTimes")
	procQueryImageName  = k32.NewProc("QueryFullProcessImageNameW")
	procCreateMutex     = k32.NewProc("CreateMutexW")
)

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(h syscall.Handle, pid uint32, cls uint32) error {
	r, _, e := procSetPriority.Call(uintptr(h), uintptr(cls))
	if r == 0 {
		return fmt.Errorf("SetPriorityClass(%d) failed: %w", pid, e)
	}
	return nil
}

func processCreated(h syscall.Handle) (uint64, error) {
	var created, exited, kernel, user syscall.Filetime
	r, _, e := procGetProcessTimes.Call(
		uintptr(h),
		uintptr(unsafe.Pointer(&created)),
		uintptr(unsafe.Pointer(&exited)),
		uintptr(unsafe.Pointer(&kernel)),
		uintptr(unsafe.Pointer(&user)),
	)
	if r == 0 {
		return 0, fmt.Errorf("GetProcessTimes failed: %w", e)
	}
	return uint64(created.HighDateTime)<<32 | uint64(created.LowDateTime), nil
}

func processName(h syscall.Handle) (string, error) {
	buf := make([]uint16, 32768)
	size := uint32(len(buf))
	r, _, e := procQueryImageName.Call(uintptr(h), 0, uintptr(unsafe.Pointer(&buf[0])), uintptr(unsafe.Pointer(&size)))
	if r == 0 {
		return "", fmt.Errorf("QueryFullProcessImageName failed: %w", e)
	}
	return filepath.Base(syscall.UTF16ToString(buf[:size])), nil
}

func processAlive(h syscall.Handle) (bool, error) {
	var created, exited, kernel, user syscall.Filetime
	r, _, e := procGetProcessTimes.Call(
		uintptr(h),
		uintptr(unsafe.Pointer(&created)),
		uintptr(unsafe.Pointer(&exited)),
		uintptr(unsafe.Pointer(&kernel)),
		uintptr(unsafe.Pointer(&user)),
	)
	if r == 0 {
		return false, fmt.Errorf("GetProcessTimes failed: %w", e)
	}
	return exited.HighDateTime == 0 && exited.LowDateTime == 0, nil
}

func parseMemSize(s string) (uint64, error) {
	s = strings.TrimSpace(s)
	if s == "" {
		return 0, fmt.Errorf("empty value")
	}
	multiplier := uint64(1)
	switch s[len(s)-1] {
	case 'k', 'K':
		multiplier = 1024
		s = s[:len(s)-1]
	case 'm', 'M':
		multiplier = 1024 * 1024
		s = s[:len(s)-1]
	case 'g', 'G':
		multiplier = 1024 * 1024 * 1024
		s = s[:len(s)-1]
	}
	n, err := strconv.ParseUint(s, 10, 64)
	if err != nil {
		return 0, err
	}
	if n > math.MaxUint64/multiplier {
		return 0, fmt.Errorf("value overflows uint64")
	}
	return n * multiplier, nil
}

func formatMemSize(b uint64) string {
	if b >= 1024*1024*1024 {
		return fmt.Sprintf("%.1fGB", float64(b)/(1024*1024*1024))
	}
	if b >= 1024*1024 {
		return fmt.Sprintf("%dMB", b/1024/1024)
	}
	return fmt.Sprintf("%dKB", b/1024)
}

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

func desiredPriority(rss, mem uint64, current uint32, gameMode, isGame bool) uint32 {
	if gameMode {
		if isGame {
			return PriorityClassHigh
		}
		return PriorityClassIdle
	}
	if rss >= mem {
		return PriorityClassHigh
	}
	if isAboveNormal(current) {
		return PriorityClassNormal
	}
	return 0
}

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 allProcs() ([]procInfo, error) {
	snap, _, e := procCreateSnap.Call(CreateToolhelp32SnapshotProcess, 0)
	if snap == ^uintptr(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 []procInfo
	for {
		name := syscall.UTF16ToString(pe.ExeFile[:])
		if name != "" {
			out = append(out, procInfo{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 singleInstance() (syscall.Handle, bool, error) {
	name, err := syscall.UTF16PtrFromString(`Local\autoPriority`)
	if err != nil {
		return 0, false, err
	}
	r, _, e := procCreateMutex.Call(0, 0, uintptr(unsafe.Pointer(name)))
	if r == 0 {
		return 0, false, fmt.Errorf("CreateMutex failed: %w", e)
	}
	return syscall.Handle(r), e == syscall.Errno(ErrorAlreadyExists), nil
}

func main() {
	memStr := flag.String("mem", "512M", "memory threshold (e.g. 512M, 1G, 2048M)")
	gameMemStr := flag.String("game-mem", "2G", "game memory threshold (e.g. 2G, 4G). Must be greater than -mem. 0 = disabled.")
	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
	}

	mem, err := parseMemSize(*memStr)
	if err != nil || mem == 0 {
		if err == nil {
			err = fmt.Errorf("must be greater than zero")
		}
		fmt.Fprintf(os.Stderr, "autoPriority: invalid -mem value: %v\n", err)
		os.Exit(1)
	}

	gameMem, err := parseMemSize(*gameMemStr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "autoPriority: invalid -game-mem value: %v\n", err)
		os.Exit(1)
	}
	if gameMem > 0 && gameMem <= mem {
		fmt.Fprintf(os.Stderr, "autoPriority: -game-mem must be greater than -mem\n")
		os.Exit(1)
	}

	mutex, exists, err := singleInstance()
	if err != nil {
		fmt.Fprintf(os.Stderr, "autoPriority: %v\n", err)
		os.Exit(1)
	}
	if exists {
		closeH(mutex)
		fmt.Fprintln(os.Stderr, "autoPriority: already running")
		return
	}
	defer closeH(mutex)

	logPath := filepath.Join(os.TempDir(), "autopriority.log")
	logFile, err := os.Create(logPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "autoPriority: cannot create log: %v\n", err)
		os.Exit(1)
	}
	defer logFile.Close()

	logf := func(format string, a ...any) {
		fmt.Fprintf(logFile, "[%s] ", time.Now().Format("02.01.2006 15:04:05"))
		fmt.Fprintf(logFile, format+"\n", a...)
	}

	logf("autoPriority started (mem=%s, game-mem=%s, interval=%s, dry-run=%v)",
		formatMemSize(mem), formatMemSize(gameMem), *interval, *dryRun)

	if !*dryRun {
		pid := uint32(os.Getpid())
		h, err := openProc(pid, ProcessSetInformation)
		if err == nil {
			err = setPrio(h, pid, PriorityClassIdle)
			closeH(h)
		}
		if err != nil {
			logf("warning: could not set own priority to IDLE: %v", err)
		} else {
			logf("own priority set to IDLE")
		}
	}

	blocked := make(map[processKey]string)
	unreadable := make(map[uint32]string)
	gameProcs := make(map[processKey]trackedProc)
	gameSaved := make(map[processKey]savedPrio)
	gameMode := false
	myPID := uint32(os.Getpid())

	readProcs := func() ([]procInfo, error) {
		procs, err := allProcs()
		if err != nil {
			return nil, err
		}

		live := make(map[uint32]string, len(procs))
		list := make([]procInfo, 0, len(procs))
		for i := range procs {
			p := &procs[i]
			live[p.PID] = p.Name
			if p.PID == myPID || p.PID == 0 {
				continue
			}

			h, setErr := openProc(p.PID, ProcessQueryLimitedInformation|ProcessSetInformation)
			p.CanSet = setErr == nil
			openErr := error(nil)
			if !p.CanSet {
				h, openErr = openProc(p.PID, ProcessQueryLimitedInformation)
			}
			if openErr != nil {
				if unreadable[p.PID] != p.Name {
					unreadable[p.PID] = p.Name
					logf("SKIP %s (PID %d): %v", p.Name, p.PID, openErr)
				}
				continue
			}

			name, readErr := processName(h)
			if readErr == nil {
				p.Name = name
			}
			created := uint64(0)
			if readErr == nil {
				created, readErr = processCreated(h)
			}
			var m processMemoryCounters
			if readErr == nil {
				m.CBM = uint32(unsafe.Sizeof(m))
				r, _, e := procGetMemInfo.Call(uintptr(h), uintptr(unsafe.Pointer(&m)), uintptr(unsafe.Sizeof(m)))
				if r == 0 {
					readErr = fmt.Errorf("GetProcessMemoryInfo failed: %w", e)
				}
			}
			var cur uintptr
			if readErr == nil {
				r, _, e := procGetPriority.Call(uintptr(h))
				cur = r
				if cur == 0 {
					readErr = fmt.Errorf("GetPriorityClass failed: %w", e)
				}
			}
			if readErr != nil {
				closeH(h)
				if unreadable[p.PID] != p.Name {
					unreadable[p.PID] = p.Name
					logf("SKIP %s (PID %d): %v", p.Name, p.PID, readErr)
				}
				continue
			}

			delete(unreadable, p.PID)
			p.Key = processKey{PID: p.PID, Created: created}
			p.RSS = uint64(m.WorkingSetSize)
			p.Prio = uint32(cur)
			p.Handle = h
			if !p.CanSet {
				if _, ok := blocked[p.Key]; !ok {
					blocked[p.Key] = p.Name
					logf("BLOCK %s (PID %d): %v (priority changes disabled)", p.Name, p.PID, setErr)
				}
			}
			list = append(list, *p)
		}

		for pid, name := range unreadable {
			if live[pid] != name {
				delete(unreadable, pid)
			}
		}
		return list, nil
	}

	restore := func(list []procInfo) {
		current := make(map[processKey]*procInfo, len(list))
		for i := range list {
			current[list[i].Key] = &list[i]
		}
		for key, saved := range gameSaved {
			alive, err := processAlive(saved.handle)
			if err != nil {
				logf("RESTORE %s (PID %d) status error: %v", saved.name, key.PID, err)
				continue
			}
			if !alive {
				closeH(saved.handle)
				delete(gameSaved, key)
				continue
			}
			cur, _, e := procGetPriority.Call(uintptr(saved.handle))
			if cur == 0 {
				logf("RESTORE %s (PID %d) priority error: %v", saved.name, key.PID, e)
				continue
			}
			if uint32(cur) != saved.prio {
				if err := setPrio(saved.handle, key.PID, saved.prio); err != nil {
					logf("RESTORE %s (PID %d) -> %s error: %v", saved.name, key.PID, prioName(saved.prio), err)
					continue
				}
				logf("RESTORE %s (PID %d) -> %s", saved.name, key.PID, prioName(saved.prio))
			}
			if p := current[key]; p != nil {
				p.Prio = saved.prio
			}
			closeH(saved.handle)
			delete(gameSaved, key)
		}
	}

	defer func() {
		for attempt := 0; attempt < 3 && len(gameSaved) > 0; attempt++ {
			restore(nil)
			if len(gameSaved) > 0 {
				time.Sleep(50 * time.Millisecond)
			}
		}
		if len(gameSaved) > 0 {
			logf("warning: %d process priorities could not be restored", len(gameSaved))
			for _, saved := range gameSaved {
				closeH(saved.handle)
			}
		}
		for _, game := range gameProcs {
			closeH(game.handle)
		}
		logf("autoPriority stopped")
		logFile.Sync()
	}()

	ticker := time.NewTicker(*interval)
	defer ticker.Stop()

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

	scan := func() {
		list, err := readProcs()
		if err != nil {
			logf("process scan error: %v", err)
			return
		}
		defer func() {
			for _, p := range list {
				if !p.Keep {
					closeH(p.Handle)
				}
			}
		}()

		current := make(map[processKey]string, len(list))
		for _, p := range list {
			current[p.Key] = p.Name
		}

		for key, name := range blocked {
			if current[key] != name {
				delete(blocked, key)
			}
		}
		for key, game := range gameProcs {
			alive, err := processAlive(game.handle)
			if err != nil {
				logf("GAME %s (PID %d) status error: %v", game.name, key.PID, err)
				continue
			}
			if !alive {
				closeH(game.handle)
				delete(gameProcs, key)
			}
		}
		for key, saved := range gameSaved {
			alive, err := processAlive(saved.handle)
			if err == nil && !alive {
				closeH(saved.handle)
				delete(gameSaved, key)
			}
		}

		for i := range list {
			p := &list[i]
			if gameMem > 0 && p.RSS >= gameMem {
				if _, ok := gameProcs[p.Key]; !ok {
					gameProcs[p.Key] = trackedProc{name: p.Name, handle: p.Handle}
					p.Keep = true
					logf("GAME DETECT %s (PID %d) RSS=%s", p.Name, p.PID, formatMemSize(p.RSS))
				}
			}
		}

		hasGame := gameMem > 0 && len(gameProcs) > 0
		enteringGame := hasGame && !gameMode
		if enteringGame {
			gameMode = true
		}
		if !hasGame && gameMode {
			restore(list)
			if len(gameSaved) == 0 {
				gameMode = false
				logf("GAME MODE OFF")
			} else {
				return
			}
		}
		if hasGame {
			sort.SliceStable(list, func(i, j int) bool {
				_, iGame := gameProcs[list[i].Key]
				_, jGame := gameProcs[list[j].Key]
				return iGame && !jGame
			})
		}

		for i := range list {
			p := &list[i]
			_, isBlocked := blocked[p.Key]
			_, isGame := gameProcs[p.Key]
			target := desiredPriority(p.RSS, mem, p.Prio, hasGame, isGame)
			action := "GAME"
			if !hasGame && target == PriorityClassHigh {
				action = "PROMOTE"
			} else if !hasGame {
				action = "DEMOTE"
			}

			if target == 0 || target == p.Prio || isBlocked {
				continue
			}
			if *dryRun {
				logf("[DRY-RUN] %s %s (PID %d) RSS=%s, %s -> %s", action, p.Name, p.PID, formatMemSize(p.RSS), prioName(p.Prio), prioName(target))
				continue
			}
			if err := setPrio(p.Handle, p.PID, target); err != nil {
				blocked[p.Key] = p.Name
				logf("BLOCK %s (PID %d): %v (priority changes disabled)", p.Name, p.PID, err)
				continue
			}
			if hasGame && !isGame && target == PriorityClassIdle {
				if _, saved := gameSaved[p.Key]; !saved {
					gameSaved[p.Key] = savedPrio{name: p.Name, prio: p.Prio, handle: p.Handle}
					p.Keep = true
				}
			}
			logf("%s %s (PID %d) RSS=%s, %s -> %s", action, p.Name, p.PID, formatMemSize(p.RSS), prioName(p.Prio), prioName(target))
			p.Prio = target
		}
		if enteringGame {
			logf("GAME MODE ON")
		}
	}

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