//go:build windows

package main

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

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

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")
	procTerminate      = k32.NewProc("TerminateProcess")
)

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 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 parseMemSize(s string) (uint64, error) {
	s = strings.TrimSpace(s)
	if s == "" {
		return 0, fmt.Errorf("empty value")
	}
	var suffix byte
	switch s[len(s)-1] {
	case 'k', 'K':
		suffix = 'K'
		s = s[:len(s)-1]
	case 'm', 'M':
		suffix = 'M'
		s = s[:len(s)-1]
	case 'g', 'G':
		suffix = 'G'
		s = s[:len(s)-1]
	}
	n, err := strconv.ParseUint(s, 10, 64)
	if err != nil {
		return 0, err
	}
	switch suffix {
	case 'K':
		return n * 1024, nil
	case 'M':
		return n * 1024 * 1024, nil
	case 'G':
		return n * 1024 * 1024 * 1024, nil
	default:
		return n, 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 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 killOtherInstances() {
	myPID := os.Getpid()
	procs, err := allProcs()
	if err != nil {
		return
	}
	myName := filepath.Base(os.Args[0])
	for _, p := range procs {
		if p.PID != uint32(myPID) && filepath.Base(p.Name) == myName {
			h, err := openProc(p.PID, ProcessTerminate)
			if err == nil {
				procTerminate.Call(uintptr(h), 0)
				closeH(h)
			}
		}
	}
}

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

	killOtherInstances()
	time.Sleep(100 * time.Millisecond)

	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()
	}

	if gameMem > 0 {
		logf("autoPriority started (mem=%s, game-mem=%s, interval=%s, dry-run=%v)", formatMemSize(mem), formatMemSize(gameMem), *interval, *dryRun)
	} else {
		logf("autoPriority started (mem=%s, interval=%s, dry-run=%v)", formatMemSize(mem), *interval, *dryRun)
	}

	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)
	gameProcs := make(map[uint32]string)
	gameMode := false

	type savedPrio struct {
		name string
		prio uint32
	}
	gameSaved := make(map[uint32]savedPrio)

	defer func() {
		restoreProcs, _ := allProcs()
		restoreMap := make(map[uint32]string, len(restoreProcs))
		for _, p := range restoreProcs {
			restoreMap[p.PID] = p.Name
		}
		for pid, sv := range gameSaved {
			if cur, ok := restoreMap[pid]; !ok || cur != sv.name {
				continue
			}
			if err := setPrio(pid, sv.prio); err != nil {
				logf("RESTORE %s (PID %d) -> %s error: %v", sv.name, pid, prioName(sv.prio), err)
			} else {
				logf("RESTORE %s (PID %d) -> %s", sv.name, pid, prioName(sv.prio))
			}
		}
		for pid, name := range promoted {
			if cur, ok := restoreMap[pid]; !ok || 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)
			}
		}
		for pid, name := range gameProcs {
			if cur, ok := restoreMap[pid]; !ok || 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)

	trySetPrio := func(qh syscall.Handle, pid uint32, name string, cls uint32) bool {
		closeH(qh)
		sh, err := openProc(pid, ProcessSetInformation)
		if err != nil {
			blocked[pid] = name
			logf("BLOCK %s (PID %d): open SetInformation failed: %v (added to exclusion list)", name, pid, err)
			return false
		}
		defer closeH(sh)
		sr, _, se := procSetPriority.Call(uintptr(sh), uintptr(cls))
		if sr == 0 {
			blocked[pid] = name
			logf("BLOCK %s (PID %d): SetPriorityClass failed: %v (added to exclusion list)", name, pid, se)
			return false
		}
		time.Sleep(50 * time.Millisecond)
		return true
	}

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

		pidMap := make(map[uint32]string, len(procs))
		for _, p := range procs {
			pidMap[p.PID] = p.Name
		}

		for pid, name := range promoted {
			if cur, ok := pidMap[pid]; !ok || cur != name {
				delete(promoted, pid)
			}
		}

		for pid := range blocked {
			if _, ok := pidMap[pid]; !ok {
				delete(blocked, pid)
			}
		}

		for pid, name := range gameProcs {
			if cur, ok := pidMap[pid]; !ok || cur != name {
				delete(gameProcs, pid)
			}
		}

		type procRSS struct {
			info processInfo
			rss  uint64
		}
		var procList []procRSS

		for _, p := range procs {
			if p.PID == myPID || p.PID == 0 {
				continue
			}

			h, err := openProc(p.PID, ProcessQueryLimitedInformation)
			if err != nil {
				if _, ok := blocked[p.PID]; !ok {
					blocked[p.PID] = p.Name
					logf("BLOCK %s (PID %d): %v (added to exclusion list)", 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)
				if _, ok := blocked[p.PID]; !ok {
					logf("GetProcessMemoryInfo(%d) error: %v", p.PID, e)
				}
				continue
			}
			closeH(h)
			rssBytes := uint64(m.WorkingSetSize)

			procList = append(procList, procRSS{p, rssBytes})

			if gameMem > 0 && rssBytes >= gameMem {
				if _, ok := gameProcs[p.PID]; !ok {
					gameProcs[p.PID] = p.Name
				}
			}
		}

		hasGame := gameMem > 0 && len(gameProcs) > 0

		for _, pr := range procList {
			p := pr.info
			rssBytes := pr.rss
			isBlocked := false
			if _, ok := blocked[p.PID]; ok {
				isBlocked = true
			}

			if _, ok := gameProcs[p.PID]; ok {
				if hasGame {
					if isBlocked {
						logf("GAME %s (PID %d) RSS=%s — tracked as game (blocked)", p.Name, p.PID, formatMemSize(rssBytes))
						continue
					}
					h, err := openProc(p.PID, ProcessQueryLimitedInformation)
					if err != nil {
						continue
					}
					origPrio, prioErr := curPrioWithHandle(h)
					if prioErr != nil {
						closeH(h)
						continue
					}
					if origPrio == PriorityClassHigh {
						closeH(h)
						continue
					}
					if *dryRun {
						closeH(h)
						logf("[DRY-RUN] %s (PID %d) RSS=%s — would set HIGH (game)",
							p.Name, p.PID, formatMemSize(rssBytes))
					} else {
						if trySetPrio(h, p.PID, p.Name, PriorityClassHigh) {
							logf("GAME %s (PID %d) RSS=%s -> HIGH",
								p.Name, p.PID, formatMemSize(rssBytes))
						} else {
							logf("GAME %s (PID %d) RSS=%s — tracked as game (priority change failed)",
								p.Name, p.PID, formatMemSize(rssBytes))
						}
					}
				}
				continue
			}

			if hasGame {
				if isBlocked {
					continue
				}
				h, err := openProc(p.PID, ProcessQueryLimitedInformation)
				if err != nil {
					continue
				}
				origPrio, prioErr := curPrioWithHandle(h)
				if prioErr != nil {
					closeH(h)
					continue
				}
				if origPrio == PriorityClassIdle {
					closeH(h)
					continue
				}
				origStr := prioName(origPrio)
				if *dryRun {
					closeH(h)
					logf("[DRY-RUN] %s (PID %d) RSS=%s, priority=%s — would set IDLE (game)",
						p.Name, p.PID, formatMemSize(rssBytes), origStr)
				} else {
					if trySetPrio(h, p.PID, p.Name, PriorityClassIdle) {
						if _, saved := gameSaved[p.PID]; !saved {
							gameSaved[p.PID] = savedPrio{name: p.Name, prio: origPrio}
						}
						delete(promoted, p.PID)
						logf("GAME %s (PID %d) RSS=%s, %s -> IDLE",
							p.Name, p.PID, formatMemSize(rssBytes), origStr)
					}
				}
				continue
			}

			if isBlocked {
				continue
			}

			if _, ok := promoted[p.PID]; ok {
				continue
			}

			h, err := openProc(p.PID, ProcessQueryLimitedInformation)
			if err != nil {
				continue
			}

			if rssBytes < mem {
				origPrio, prioErr := curPrioWithHandle(h)
				if prioErr == nil && isAboveNormal(origPrio) {
					if *dryRun {
						closeH(h)
						logf("[DRY-RUN] %s (PID %d) RSS=%s, priority=%s — would set NORMAL",
							p.Name, p.PID, formatMemSize(rssBytes), prioName(origPrio))
					} else {
						if trySetPrio(h, p.PID, p.Name, PriorityClassNormal) {
							logf("DEMOTE %s (PID %d) RSS=%s, %s -> NORMAL",
								p.Name, p.PID, formatMemSize(rssBytes), prioName(origPrio))
						}
					}
				} else {
					closeH(h)
				}
				continue
			}

			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=%s, priority=%s — would set HIGH",
					p.Name, p.PID, formatMemSize(rssBytes), origStr)
				continue
			}

			if trySetPrio(h, p.PID, p.Name, PriorityClassHigh) {
				promoted[p.PID] = p.Name
				logf("PROMOTE %s (PID %d) RSS=%s, %s -> HIGH",
					p.Name, p.PID, formatMemSize(rssBytes), origStr)
			}
		}

		if gameMem > 0 {
			newHasGame := len(gameProcs) > 0
			if newHasGame && !gameMode {
				gameMode = true
				logf("GAME MODE ON")
			}
			if !newHasGame && gameMode {
				gameMode = false
				logf("GAME MODE OFF")
				for pid, sv := range gameSaved {
					cur, ok := pidMap[pid]
					if !ok || cur != sv.name {
						delete(gameSaved, pid)
						continue
					}
					if *dryRun {
						logf("[DRY-RUN] RESTORE %s (PID %d) -> %s", sv.name, pid, prioName(sv.prio))
					} else {
						if err := setPrio(pid, sv.prio); err != nil {
							logf("RESTORE %s (PID %d) -> %s error: %v", sv.name, pid, prioName(sv.prio), err)
						} else {
							logf("RESTORE %s (PID %d) -> %s", sv.name, pid, prioName(sv.prio))
						}
					}
					delete(gameSaved, pid)
				}
			}
		}
	}

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