//go:build windows

package game

import (
	"fmt"
	"unsafe"

	"golang.org/x/sys/windows"
)

// Per-process GPU engine utilization via PDH — the same counters Task
// Manager's "GPU engine" columns read. Unlike nvidia-smi's compute-apps
// this covers graphics work under WDDM, so games show up. The counter is
// added with PdhAddEnglishCounterW, which is independent of the Windows
// display language.

var (
	pdhDLL                          = windows.NewLazySystemDLL("pdh.dll")
	procPdhOpenQuery                = pdhDLL.NewProc("PdhOpenQueryW")
	procPdhAddEnglishCounter        = pdhDLL.NewProc("PdhAddEnglishCounterW")
	procPdhCollectQueryData         = pdhDLL.NewProc("PdhCollectQueryData")
	procPdhGetFormattedCounterArray = pdhDLL.NewProc("PdhGetFormattedCounterArrayW")
	procPdhCloseQuery               = pdhDLL.NewProc("PdhCloseQuery")
)

const (
	pdhFmtDouble = 0x00000200 // PDH_FMT_DOUBLE
	pdhMoreData  = 0x800007D2 // PDH_MORE_DATA
)

// pdhCountervalueItem mirrors PDH_FMT_COUNTERVALUE_ITEM (64-bit, double).
type pdhCountervalueItem struct {
	name    *uint16
	cStatus uint32
	_       uint32 // alignment padding
	value   float64
}

// gpuEngineSampler holds an open PDH query on the wildcard GPU Engine
// utilization counter. Keeping the query open across polls is what makes
// the rate-based values meaningful; the sampler lives as long as the
// process (PdhCloseQuery would only matter on unload).
type gpuEngineSampler struct {
	query   uintptr // PDH_HQUERY
	counter uintptr // PDH_HCOUNTER
	primed  bool
}

// openGPUEngineSampler opens a query on the per-process GPU engine
// utilization counter (all instances).
func openGPUEngineSampler() (*gpuEngineSampler, error) {
	var q uintptr
	if r, _, _ := procPdhOpenQuery.Call(0, 0, uintptr(unsafe.Pointer(&q))); r != 0 {
		return nil, fmt.Errorf("PdhOpenQuery: status %#x", r)
	}
	path, err := windows.UTF16PtrFromString(`\GPU Engine(*)\Utilization Percentage`)
	if err != nil {
		procPdhCloseQuery.Call(q)
		return nil, err
	}
	var c uintptr
	if r, _, _ := procPdhAddEnglishCounter.Call(q, uintptr(unsafe.Pointer(path)), 0, uintptr(unsafe.Pointer(&c))); r != 0 {
		procPdhCloseQuery.Call(q)
		return nil, fmt.Errorf("PdhAddEnglishCounter: status %#x", r)
	}
	return &gpuEngineSampler{query: q, counter: c}, nil
}

// sample collects the counter once and returns per-PID 3D-engine
// utilization in percent. The first call after open only primes the rate
// calculation and returns errNotPrimed. Processes can drive several 3D
// engines; their values are summed.
func (s *gpuEngineSampler) sample() (map[int]float64, error) {
	if r, _, _ := procPdhCollectQueryData.Call(s.query); r != 0 {
		return nil, fmt.Errorf("PdhCollectQueryData: status %#x", r)
	}
	if !s.primed {
		s.primed = true
		return nil, errNotPrimed
	}
	var size, count uint32
	r, _, _ := procPdhGetFormattedCounterArray.Call(s.counter, pdhFmtDouble,
		uintptr(unsafe.Pointer(&size)), uintptr(unsafe.Pointer(&count)), 0)
	if r == pdhMoreData && size == 0 {
		return nil, nil // no GPU engine instances at all
	}
	if r != pdhMoreData {
		return nil, fmt.Errorf("PdhGetFormattedCounterArray(size): status %#x", r)
	}
	buf := make([]byte, size)
	r, _, _ = procPdhGetFormattedCounterArray.Call(s.counter, pdhFmtDouble,
		uintptr(unsafe.Pointer(&size)), uintptr(unsafe.Pointer(&count)),
		uintptr(unsafe.Pointer(&buf[0])))
	if r != 0 {
		return nil, fmt.Errorf("PdhGetFormattedCounterArray: status %#x", r)
	}
	items := unsafe.Slice((*pdhCountervalueItem)(unsafe.Pointer(&buf[0])), int(count))
	out := make(map[int]float64)
	for i := range items {
		if items[i].cStatus != 0 || items[i].name == nil {
			continue
		}
		pid, engType, ok := parseGPUEngineInstance(windows.UTF16PtrToString(items[i].name))
		if !ok || engType != "3D" {
			continue // only the 3D engine marks game-like work
		}
		out[pid] += items[i].value
	}
	return out, nil
}