// SPDX-License-Identifier: Apache-2.0
//
// Copyright The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//go:build windows

package gdi32

import (
	"errors"
	"fmt"
	"unsafe"

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

// https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/d3dkmthk/ne-d3dkmthk-_kmtqueryadapterinfotype
// https://github.com/nalilord/AMDPlugin/blob/bb405b6d58ea543ff630f3488384473bee79f447/Common/d3dkmthk.pas#L54
const (
	// KMTQAITYPE_GETSEGMENTSIZE pPrivateDriverData points to a D3DKMT_SEGMENTSIZEINFO structure that contains information about the size of memory and aperture segments.
	KMTQAITYPE_GETSEGMENTSIZE = 3
	// KMTQAITYPE_ADAPTERADDRESS pPrivateDriverData points to a D3DKMT_ADAPTERADDRESS structure that contains information about the physical location on the PCI bus of the adapter.
	KMTQAITYPE_ADAPTERADDRESS = 6
	// KMTQAITYPE_ADAPTERREGISTRYINFO pPrivateDriverData points to a D3DKMT_ADAPTERREGISTRYINFO structure that contains registry information about the graphics adapter.
	KMTQAITYPE_ADAPTERREGISTRYINFO = 8
)

var ErrNoGPUDevices = errors.New("no GPU devices found")

func GetGPUDeviceByLUID(adapterLUID windows.LUID) (GPUDevice, error) {
	open := D3DKMT_OPENADAPTERFROMLUID{
		AdapterLUID: adapterLUID,
	}

	if err := D3DKMTOpenAdapterFromLuid(&open); err != nil {
		return GPUDevice{}, fmt.Errorf("D3DKMTOpenAdapterFromLuid failed: %w", err)
	}

	errs := make([]error, 0)

	gpuDevice, err := GetGPUDevice(open.HAdapter)
	if err != nil {
		errs = append(errs, fmt.Errorf("GetGPUDevice failed: %w", err))
	}

	if err := D3DKMTCloseAdapter(&D3DKMT_CLOSEADAPTER{
		HAdapter: open.HAdapter,
	}); err != nil {
		errs = append(errs, fmt.Errorf("D3DKMTCloseAdapter failed: %w", err))
	}

	if len(errs) > 0 {
		return gpuDevice, fmt.Errorf("errors occurred while getting GPU device: %w", errors.Join(errs...))
	}

	gpuDevice.LUID = adapterLUID

	return gpuDevice, nil
}

func GetGPUDevice(hAdapter D3DKMT_HANDLE) (GPUDevice, error) {
	var gpuDevice GPUDevice

	// Try segment size first
	var size D3DKMT_SEGMENTSIZEINFO

	query := D3DKMT_QUERYADAPTERINFO{
		hAdapter:              hAdapter,
		queryType:             KMTQAITYPE_GETSEGMENTSIZE,
		pPrivateDriverData:    unsafe.Pointer(&size),
		privateDriverDataSize: uint32(unsafe.Sizeof(size)),
	}

	if err := D3DKMTQueryAdapterInfo(&query); err != nil {
		return gpuDevice, fmt.Errorf("D3DKMTQueryAdapterInfo (segment size) failed: %w", err)
	}

	gpuDevice.DedicatedVideoMemorySize = size.DedicatedVideoMemorySize
	gpuDevice.DedicatedSystemMemorySize = size.DedicatedSystemMemorySize
	gpuDevice.SharedSystemMemorySize = size.SharedSystemMemorySize

	// Now try registry info
	var address D3DKMT_ADAPTERADDRESS

	query.queryType = KMTQAITYPE_ADAPTERADDRESS
	query.pPrivateDriverData = unsafe.Pointer(&address)
	query.privateDriverDataSize = uint32(unsafe.Sizeof(address))

	if err := D3DKMTQueryAdapterInfo(&query); err != nil {
		return gpuDevice, fmt.Errorf("D3DKMTQueryAdapterInfo (adapter address) failed: %w", err)
	}

	gpuDevice.BusNumber = address.BusNumber
	gpuDevice.DeviceNumber = address.DeviceNumber
	gpuDevice.FunctionNumber = address.FunctionNumber

	// Now try registry info
	var info D3DKMT_ADAPTERREGISTRYINFO

	query.queryType = KMTQAITYPE_ADAPTERREGISTRYINFO
	query.pPrivateDriverData = unsafe.Pointer(&info)
	query.privateDriverDataSize = uint32(unsafe.Sizeof(info))

	if err := D3DKMTQueryAdapterInfo(&query); err != nil && !errors.Is(err, windows.ERROR_FILE_NOT_FOUND) {
		return gpuDevice, fmt.Errorf("D3DKMTQueryAdapterInfo (info) failed: %w", err)
	}

	gpuDevice.AdapterString = windows.UTF16ToString(info.AdapterString[:])

	return gpuDevice, nil
}

func GetGPUDevices() ([]GPUDevice, error) {
	gpuDevices := make([]GPUDevice, 0, 2)

	// First call: Get the number of adapters
	enumAdapters := D3DKMT_ENUMADAPTERS2{
		NumAdapters: 0,
		PAdapters:   nil,
	}

	if err := D3DKMTEnumAdapters2(&enumAdapters); err != nil {
		return gpuDevices, fmt.Errorf("D3DKMTEnumAdapters2 (get count) failed: %w", err)
	}

	if enumAdapters.NumAdapters == 0 {
		return gpuDevices, ErrNoGPUDevices
	}

	// Second call: Get the actual adapter information
	pAdapters := make([]D3DKMT_ADAPTERINFO, enumAdapters.NumAdapters)
	enumAdapters.PAdapters = &pAdapters[0]

	if err := D3DKMTEnumAdapters2(&enumAdapters); err != nil {
		return gpuDevices, fmt.Errorf("D3DKMTEnumAdapters2 (get adapters) failed: %w", err)
	}

	var errs []error

	// Process each adapter
	for i := range enumAdapters.NumAdapters {
		adapter := pAdapters[i]
		// Validate handle before using it
		if adapter.HAdapter == 0 {
			errs = append(errs, fmt.Errorf("adapter %d has null handle", i))

			continue
		}

		func() {
			defer func() {
				if closeErr := D3DKMTCloseAdapter(&D3DKMT_CLOSEADAPTER{
					HAdapter: adapter.HAdapter,
				}); closeErr != nil {
					errs = append(errs, fmt.Errorf("failed to close adapter %v: %w", adapter.AdapterLUID, closeErr))
				}
			}()

			gpuDevice, err := GetGPUDevice(adapter.HAdapter)
			if err != nil {
				errs = append(errs, fmt.Errorf("failed to get GPU device for adapter %v: %w", adapter.AdapterLUID, err))

				return
			}

			gpuDevice.LUID = adapter.AdapterLUID

			gpuDevices = append(gpuDevices, gpuDevice)
		}()
	}

	if len(errs) > 0 {
		return gpuDevices, errors.Join(errs...)
	}

	if len(gpuDevices) == 0 {
		return gpuDevices, ErrNoGPUDevices
	}

	return gpuDevices, nil
}