//go:build windows
// +build windows

package common

import (
	"context"
	"fmt"
	"path/filepath"
	"reflect"
	"strings"
	"syscall"
	"unsafe"

	"github.com/yusufpapurcu/wmi"
	"golang.org/x/sys/windows"
)

// for double values
type PDH_FMT_COUNTERVALUE_DOUBLE struct {
	CStatus     uint32
	DoubleValue float64
}

// for 64 bit integer values
type PDH_FMT_COUNTERVALUE_LARGE struct {
	CStatus    uint32
	LargeValue int64
}

// for long values
type PDH_FMT_COUNTERVALUE_LONG struct {
	CStatus   uint32
	LongValue int32
	padding   [4]byte
}

// windows system const
const (
	ERROR_SUCCESS        = 0
	ERROR_FILE_NOT_FOUND = 2
	DRIVE_REMOVABLE      = 2
	DRIVE_FIXED          = 3
	HKEY_LOCAL_MACHINE   = 0x80000002
	RRF_RT_REG_SZ        = 0x00000002
	RRF_RT_REG_DWORD     = 0x00000010
	PDH_FMT_LONG         = 0x00000100
	PDH_FMT_DOUBLE       = 0x00000200
	PDH_FMT_LARGE        = 0x00000400
	PDH_INVALID_DATA     = 0xc0000bc6
	PDH_INVALID_HANDLE   = 0xC0000bbc
	PDH_NO_DATA          = 0x800007d5

	STATUS_BUFFER_OVERFLOW      = 0x80000005
	STATUS_BUFFER_TOO_SMALL     = 0xC0000023
	STATUS_INFO_LENGTH_MISMATCH = 0xC0000004
)

const (
	ProcessBasicInformation = 0
	ProcessWow64Information = 26
	ProcessQueryInformation = windows.PROCESS_DUP_HANDLE | windows.PROCESS_QUERY_INFORMATION

	SystemExtendedHandleInformationClass = 64
)

var (
	Modkernel32 = windows.NewLazySystemDLL("kernel32.dll")
	ModNt       = windows.NewLazySystemDLL("ntdll.dll")
	ModPdh      = windows.NewLazySystemDLL("pdh.dll")
	ModPsapi    = windows.NewLazySystemDLL("psapi.dll")

	ProcGetSystemTimes                   = Modkernel32.NewProc("GetSystemTimes")
	ProcNtQuerySystemInformation         = ModNt.NewProc("NtQuerySystemInformation")
	ProcRtlGetNativeSystemInformation    = ModNt.NewProc("RtlGetNativeSystemInformation")
	ProcRtlNtStatusToDosError            = ModNt.NewProc("RtlNtStatusToDosError")
	ProcNtQueryInformationProcess        = ModNt.NewProc("NtQueryInformationProcess")
	ProcNtReadVirtualMemory              = ModNt.NewProc("NtReadVirtualMemory")
	ProcNtWow64QueryInformationProcess64 = ModNt.NewProc("NtWow64QueryInformationProcess64")
	ProcNtWow64ReadVirtualMemory64       = ModNt.NewProc("NtWow64ReadVirtualMemory64")

	PdhOpenQuery                = ModPdh.NewProc("PdhOpenQuery")
	PdhAddEnglishCounterW       = ModPdh.NewProc("PdhAddEnglishCounterW")
	PdhCollectQueryData         = ModPdh.NewProc("PdhCollectQueryData")
	PdhGetFormattedCounterValue = ModPdh.NewProc("PdhGetFormattedCounterValue")
	PdhCloseQuery               = ModPdh.NewProc("PdhCloseQuery")

	procQueryDosDeviceW = Modkernel32.NewProc("QueryDosDeviceW")
)

type FILETIME struct {
	DwLowDateTime  uint32
	DwHighDateTime uint32
}

// borrowed from net/interface_windows.go
func BytePtrToString(p *uint8) string {
	a := (*[10000]uint8)(unsafe.Pointer(p))
	i := 0
	for a[i] != 0 {
		i++
	}
	return string(a[:i])
}

// CounterInfo struct is used to track a windows performance counter
// copied from https://github.com/mackerelio/mackerel-agent/
type CounterInfo struct {
	PostName    string
	CounterName string
	Counter     windows.Handle
}

// CreateQuery with a PdhOpenQuery call
// copied from https://github.com/mackerelio/mackerel-agent/
func CreateQuery() (windows.Handle, error) {
	var query windows.Handle
	r, _, err := PdhOpenQuery.Call(0, 0, uintptr(unsafe.Pointer(&query)))
	if r != 0 {
		return 0, err
	}
	return query, nil
}

// CreateCounter with a PdhAddEnglishCounterW call
func CreateCounter(query windows.Handle, pname, cname string) (*CounterInfo, error) {
	var counter windows.Handle
	r, _, err := PdhAddEnglishCounterW.Call(
		uintptr(query),
		uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(cname))),
		0,
		uintptr(unsafe.Pointer(&counter)))
	if r != 0 {
		return nil, err
	}
	return &CounterInfo{
		PostName:    pname,
		CounterName: cname,
		Counter:     counter,
	}, nil
}

// GetCounterValue get counter value from handle
// adapted from https://github.com/mackerelio/mackerel-agent/
func GetCounterValue(counter windows.Handle) (float64, error) {
	var value PDH_FMT_COUNTERVALUE_DOUBLE
	r, _, err := PdhGetFormattedCounterValue.Call(uintptr(counter), PDH_FMT_DOUBLE, uintptr(0), uintptr(unsafe.Pointer(&value)))
	if r != 0 && r != PDH_INVALID_DATA {
		return 0.0, err
	}
	return value.DoubleValue, nil
}

type Win32PerformanceCounter struct {
	PostName    string
	CounterName string
	Query       windows.Handle
	Counter     windows.Handle
}

func NewWin32PerformanceCounter(postName, counterName string) (*Win32PerformanceCounter, error) {
	query, err := CreateQuery()
	if err != nil {
		return nil, err
	}
	counter := Win32PerformanceCounter{
		Query:       query,
		PostName:    postName,
		CounterName: counterName,
	}
	r, _, err := PdhAddEnglishCounterW.Call(
		uintptr(counter.Query),
		uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(counter.CounterName))),
		0,
		uintptr(unsafe.Pointer(&counter.Counter)),
	)
	if r != 0 {
		return nil, err
	}
	return &counter, nil
}

func (w *Win32PerformanceCounter) GetValue() (float64, error) {
	r, _, err := PdhCollectQueryData.Call(uintptr(w.Query))
	if r != 0 && err != nil {
		if r == PDH_NO_DATA {
			return 0.0, fmt.Errorf("%w: this counter has not data", err)
		}
		return 0.0, err
	}

	return GetCounterValue(w.Counter)
}

func ProcessorQueueLengthCounter() (*Win32PerformanceCounter, error) {
	return NewWin32PerformanceCounter("processor_queue_length", `\System\Processor Queue Length`)
}

// WMIQueryWithContext - wraps wmi.Query with a timed-out context to avoid hanging
func WMIQueryWithContext(ctx context.Context, query string, dst interface{}, connectServerArgs ...interface{}) error {
	if _, ok := ctx.Deadline(); !ok {
		ctxTimeout, cancel := context.WithTimeout(ctx, Timeout)
		defer cancel()
		ctx = ctxTimeout
	}

	errChan := make(chan error, 1)
	go func() {
		errChan <- wmi.Query(query, dst, connectServerArgs...)
	}()

	select {
	case <-ctx.Done():
		return ctx.Err()
	case err := <-errChan:
		return err
	}
}

// Convert paths using native DOS format like:
//   "\Device\HarddiskVolume1\Windows\systemew\file.txt"
// into:
//   "C:\Windows\systemew\file.txt"
func ConvertDOSPath(p string) string {
	rawDrive := strings.Join(strings.Split(p, `\`)[:3], `\`)

	for d := 'A'; d <= 'Z'; d++ {
		szDeviceName := string(d) + ":"
		szTarget := make([]uint16, 512)
		ret, _, _ := procQueryDosDeviceW.Call(uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(szDeviceName))),
			uintptr(unsafe.Pointer(&szTarget[0])),
			uintptr(len(szTarget)))
		if ret != 0 && windows.UTF16ToString(szTarget[:]) == rawDrive {
			return filepath.Join(szDeviceName, p[len(rawDrive):])
		}
	}
	return p
}

type NtStatus uint32

func (s NtStatus) Error() error {
	if s == 0 {
		return nil
	}
	return fmt.Errorf("NtStatus 0x%08x", uint32(s))
}

func (s NtStatus) IsError() bool {
	return s>>30 == 3
}

type SystemExtendedHandleTableEntryInformation struct {
	Object                uintptr
	UniqueProcessId       uintptr
	HandleValue           uintptr
	GrantedAccess         uint32
	CreatorBackTraceIndex uint16
	ObjectTypeIndex       uint16
	HandleAttributes      uint32
	Reserved              uint32
}

type SystemExtendedHandleInformation struct {
	NumberOfHandles uintptr
	Reserved        uintptr
	Handles         [1]SystemExtendedHandleTableEntryInformation
}

// CallWithExpandingBuffer https://github.com/hillu/go-ntdll
func CallWithExpandingBuffer(fn func() NtStatus, buf *[]byte, resultLength *uint32) NtStatus {
	for {
		if st := fn(); st == STATUS_BUFFER_OVERFLOW || st == STATUS_BUFFER_TOO_SMALL || st == STATUS_INFO_LENGTH_MISMATCH {
			if int(*resultLength) <= cap(*buf) {
				(*reflect.SliceHeader)(unsafe.Pointer(buf)).Len = int(*resultLength)
			} else {
				*buf = make([]byte, int(*resultLength))
			}
			continue
		} else {
			if !st.IsError() {
				*buf = (*buf)[:int(*resultLength)]
			}
			return st
		}
	}
}

func NtQuerySystemInformation(
	SystemInformationClass uint32,
	SystemInformation *byte,
	SystemInformationLength uint32,
	ReturnLength *uint32,
) NtStatus {
	r0, _, _ := ProcNtQuerySystemInformation.Call(
		uintptr(SystemInformationClass),
		uintptr(unsafe.Pointer(SystemInformation)),
		uintptr(SystemInformationLength),
		uintptr(unsafe.Pointer(ReturnLength)))
	return NtStatus(r0)
}

// represents a win32 thread status
// see <winternl.h> for possible values
type THREAD_STATE uint32

const (
	StateInitialized THREAD_STATE = iota
	StateReady
	StateRunning
	StateStandby
	StateTerminated
	StateWait
	StateTransition
	StateUnknown
)

func (s THREAD_STATE) String() string {
	switch s {
	case StateInitialized:
		return "StateInitialized"
	case StateReady:
		return "StateReady"
	case StateRunning:
		return "StateRunning"
	case StateStandby:
		return "StateStandby"
	case StateTerminated:
		return "StateTerminated"
	case StateWait:
		return "StateWait"
	case StateTransition:
		return "StateTransition"
	case StateUnknown:
		return "StateUnknown"
	default:
		return "<invalid state>"
	}
}

// "BUSY" thread means a thread that is either running, transitioning to run or on a state that demans resources
// see https://www.brendangregg.com/blog/2017-08-08/linux-load-averages.html
func (s THREAD_STATE) Busy() bool {
	switch s {
	case StateReady, StateRunning, StateStandby, StateTransition:
		return true
	default:
		return false
	}
}

// SYSTEM_THREAD_INFORMATION contains thread information as it is returned by NtQuerySystemInformation() API call
// look for its structure & documentation at:
// https://learn.microsoft.com/en-us/windows/win32/api/winternl/nf-winternl-ntquerysysteminformation
type SYSTEM_THREAD_INFORMATION struct {
	Reserved1     [3]int64
	Reserved2     uint32
	StartAddress  uintptr
	UniqueProcess windows.Handle
	UniqueThread  windows.Handle
	Priority      int32
	BasePriority  int32
	Reserved3     uint32
	ThreadState   THREAD_STATE
	WaitReason    uint32
}

// SYSTEM_PROCESS_INFORMATION is a convenience struct to have first thread address at hand
// for this technique to access to heterogeneous data, see:
// https://justen.codes/breaking-all-the-rules-using-go-to-call-windows-api-2cbfd8c79724
type SYSTEM_PROCESS_INFORMATION struct {
	windows.SYSTEM_PROCESS_INFORMATION
	ThreadsTable [1]SYSTEM_THREAD_INFORMATION
}

// Stats are the stats this package offers
type Stats struct {
	ProcessCount    uint32
	ThreadCount     uint32
	ThreadsByStatus map[THREAD_STATE]uint32
	Load            uint32 // number of threads that contribute to system load, see https://www.brendangregg.com/blog/2017-08-08/linux-load-averages.html
}

func EmptyStats() *Stats {
	return &Stats{
		ThreadsByStatus: make(map[THREAD_STATE]uint32),
	}
}

// AddProc increments process count and returns itself.
func (s *Stats) AddProc() *Stats {
	s.ProcessCount += 1
	return s
}

// AddThread increments thread count, also updates ThreadsByStatus based on the status,
// finally if the state represents a busy thread, it increments the load.
// returns the current stats structure pointer.
func (s *Stats) AddThread(state THREAD_STATE) *Stats {
	s.ThreadCount += 1
	s.ThreadsByStatus[state] += 1

	if state.Busy() {
		s.Load += 1
	}

	return s
}

// SystemProcessInformationWalk is a helper structure to walk through the raw bytes
// that NtQuerySystemInformation produces and get correct structures
type SystemProcessInformationWalk struct {
	SizeInBytes uint32 // buffer size
	Offset      uint32 // current offset
	Buffer      []byte // buffer with the data
}

// Process returns the process under current offset
func (w *SystemProcessInformationWalk) Process() *SYSTEM_PROCESS_INFORMATION {
	return (*SYSTEM_PROCESS_INFORMATION)(unsafe.Pointer(&w.Buffer[w.Offset]))
}

// Next moves offset to the next process structure
// it returns true if there are still more PENDING processess to iterate
// it returns false if there are no more PENDING processess to iterate
// calling Next() when there are no more processes, has no effect
func (w *SystemProcessInformationWalk) Next() bool {
	proc := w.Process()

	if proc.NextEntryOffset == 0 || proc.NextEntryOffset+w.Offset > w.SizeInBytes {
		return false // reached the end
	}

	w.Offset += proc.NextEntryOffset

	return true
}

// Stats calculate stats for all processes and their threads
func (w *SystemProcessInformationWalk) Stats() *Stats {
	stats := EmptyStats()

	for {
		proc := w.Process()

		stats.AddProc()

		WalkThreads(proc, func(t SYSTEM_THREAD_INFORMATION) {
			stats.AddThread(t.ThreadState)
		})

		if ok := w.Next(); !ok {
			break
		}
	}

	return stats
}

// WalkThreads() iterates over all threads of current process and applies given function
func WalkThreads(proc *SYSTEM_PROCESS_INFORMATION, fn func(t SYSTEM_THREAD_INFORMATION)) {
	for i := 0; i < int(proc.NumberOfThreads); i++ {
		thread := *(*SYSTEM_THREAD_INFORMATION)(unsafe.Pointer(
			uintptr(unsafe.Pointer(&proc.ThreadsTable[0])) +
				uintptr(i)*unsafe.Sizeof(proc.ThreadsTable[0]),
		))

		fn(thread)
	}
}

// GetSystemProcessInformation retrieves information of all procecess and threads
// see: https://learn.microsoft.com/en-us/windows/win32/api/winternl/nf-winternl-ntquerysysteminformation
// look for SystemProcessInformation and related structures SYSTEM_PROCESS_INFORMATION and SYSTEM_THREAD_INFORMATION
// the returned structure has methods to walk through the structure
func GetSystemProcessInformation() (*SystemProcessInformationWalk, error) {
	var (
		oneKb     uint32 = 1024
		allocKb   uint32 = 1
		buffer    []byte
		usedBytes uint32
	)

	// iterating instead of calling common.CallWithExpandingBuffer hangs forever
	for {
		var allocBytes uint32 = allocKb * oneKb
		buffer = make([]byte, allocBytes)

		st := NtQuerySystemInformation(
			windows.SystemProcessInformation,
			&buffer[0],
			allocBytes,
			&usedBytes,
		)

		if st == NtStatus(windows.STATUS_INFO_LENGTH_MISMATCH) {
			allocKb *= 2
			continue
		}

		if st.IsError() {
			return nil, st.Error()
		}

		break
	}

	return &SystemProcessInformationWalk{
		SizeInBytes: usedBytes,
		Offset:      0,
		Buffer:      buffer,
	}, nil
}