//go:build nvml

// This file is the NVML-backed GPU discovery path, built only with the `nvml`
// tag because the bindings require cgo. Without the tag the agent falls back to
// nvidia-smi, which keeps the server side — and development machines without a
// C toolchain or a driver — building cleanly (plan §3.2).
//
// What NVML adds over nvidia-smi is NVLink topology. Nothing else exposes which
// GPUs share an NVLink domain, and that is what keeps a tensor-parallel group
// off the slow path inside a node (plan §6.2, constraint 3).

package discovery

import (
	"context"
	"fmt"

	"github.com/NVIDIA/go-nvml/pkg/nvml"
)

// nvmlAvailable reports that this build can use NVML.
const nvmlAvailable = true

// DiscoverGPUsNVML enumerates GPUs through NVML.
//
// It returns an error rather than falling back on its own: the caller decides
// whether to degrade to nvidia-smi, because on a node that is supposed to have
// working NVML a silent downgrade hides a broken driver installation.
func DiscoverGPUsNVML(_ context.Context) ([]GPU, error) {
	if ret := nvml.Init(); ret != nvml.SUCCESS {
		return nil, fmt.Errorf("nvml init: %s", nvml.ErrorString(ret))
	}
	defer func() { _ = nvml.Shutdown() }()

	driver, _ := nvml.SystemGetDriverVersion()
	cuda := ""
	if v, ret := nvml.SystemGetCudaDriverVersion(); ret == nvml.SUCCESS {
		// NVML reports 12040 for 12.4.
		cuda = fmt.Sprintf("%d.%d", v/1000, (v%1000)/10)
	}

	count, ret := nvml.DeviceGetCount()
	if ret != nvml.SUCCESS {
		return nil, fmt.Errorf("nvml device count: %s", nvml.ErrorString(ret))
	}

	// Bus id to index, so NVLink peers can be reported as indices. NVML gives
	// the remote end of a link as a PCI address, which is not what placement
	// works in.
	byBusID := make(map[string]int, count)
	gpus := make([]GPU, 0, count)

	for i := range count {
		dev, ret := nvml.DeviceGetHandleByIndex(i)
		if ret != nvml.SUCCESS {
			return nil, fmt.Errorf("nvml device %d: %s", i, nvml.ErrorString(ret))
		}

		g := GPU{Index: i, NUMANode: -1, DriverVersion: driver, CUDAVersion: cuda}
		g.UUID, _ = dev.GetUUID()
		g.Name, _ = dev.GetName()

		if mem, ret := dev.GetMemoryInfo(); ret == nvml.SUCCESS {
			g.MemoryTotalMB = int(mem.Total / (1 << 20))
			g.MemoryUsedMB = int(mem.Used / (1 << 20))
		}
		if major, minor, ret := dev.GetCudaComputeCapability(); ret == nvml.SUCCESS {
			g.ComputeCapability = fmt.Sprintf("%d.%d", major, minor)
		}
		if pci, ret := dev.GetPciInfo(); ret == nvml.SUCCESS {
			g.PCIBusID = busIDString(pci.BusId[:])
			byBusID[normalizeBusID(g.PCIBusID)] = i
		}
		if t, ret := dev.GetTemperature(nvml.TEMPERATURE_GPU); ret == nvml.SUCCESS {
			g.Health.TemperatureC = int(t)
		}
		// Volatile counters reset on driver reload, which is what we want: an
		// error from three months and one reboot ago is not a reason to refuse
		// a card today.
		if n, ret := dev.GetTotalEccErrors(nvml.MEMORY_ERROR_TYPE_UNCORRECTED, nvml.VOLATILE_ECC); ret == nvml.SUCCESS {
			g.Health.ECCErrors = int(n)
		}

		assessHealth(&g)
		// A card in MIG mode is its instances, not itself: CUDA will not
		// open the whole card while MIG is on.
		devs, _ := migFromNVML(g, nvmlCard{dev})
		gpus = append(gpus, devs...)
	}

	// A second pass, because a link's remote end may be a GPU enumerated
	// later. MIG instances have no NVLink of their own to use.
	for i := range gpus {
		if gpus[i].IsMIG() {
			continue
		}
		dev, ret := nvml.DeviceGetHandleByIndex(gpus[i].Index)
		if ret != nvml.SUCCESS {
			continue
		}
		gpus[i].NVLinkPeers = nvlinkPeers(dev, byBusID, gpus[i].Index)
	}
	return gpus, nil
}

// nvlinkPeers returns the indices of GPUs reachable from dev over NVLink.
func nvlinkPeers(dev nvml.Device, byBusID map[string]int, self int) []int {
	seen := map[int]struct{}{}
	var out []int

	for link := range nvml.NVLINK_MAX_LINKS {
		pci, ret := dev.GetNvLinkRemotePciInfo(link)
		if ret != nvml.SUCCESS {
			// Unpopulated links return an error; that is the normal case for
			// most of the 36 possible links, so it is not worth reporting.
			continue
		}
		peer, ok := byBusID[normalizeBusID(busIDString(pci.BusId[:]))]
		if !ok || peer == self {
			// The remote end can be an NVSwitch rather than a GPU, in which
			// case it is not in the map.
			continue
		}
		if _, dup := seen[peer]; dup {
			continue
		}
		seen[peer] = struct{}{}
		out = append(out, peer)
	}
	return out
}

// busIDString turns NVML's fixed-size C char array into a Go string.
func busIDString(b []int8) string {
	buf := make([]byte, 0, len(b))
	for _, c := range b {
		if c == 0 {
			break
		}
		buf = append(buf, byte(c))
	}
	return string(buf)
}

// nvmlCard adapts an NVML device to what MIG discovery reads of it.
type nvmlCard struct{ dev nvml.Device }

func (c nvmlCard) MIGMode() (current, pending, ok bool) {
	cur, pend, ret := c.dev.GetMigMode()
	if ret != nvml.SUCCESS {
		// NOT_SUPPORTED on a card without MIG, which is most of them.
		return false, false, false
	}
	return cur == nvml.DEVICE_MIG_ENABLE, pend == nvml.DEVICE_MIG_ENABLE, true
}

func (c nvmlCard) MaxMIGDevices() int {
	n, ret := c.dev.GetMaxMigDeviceCount()
	if ret != nvml.SUCCESS {
		return 0
	}
	return n
}

func (c nvmlCard) MIGDevice(i int) (nvmlMIGInstance, bool) {
	h, ret := c.dev.GetMigDeviceHandleByIndex(i)
	if ret != nvml.SUCCESS {
		// NOT_FOUND for an index with no instance: the count is the most a
		// card can hold, not how many it has.
		return nil, false
	}
	return nvmlMIG{h}, true
}

// nvmlMIG adapts a MIG device handle.
type nvmlMIG struct{ dev nvml.Device }

func (m nvmlMIG) UUID() string {
	u, _ := m.dev.GetUUID()
	return u
}

func (m nvmlMIG) Name() string {
	n, _ := m.dev.GetName()
	return n
}

func (m nvmlMIG) Memory() (totalMB, usedMB int, ok bool) {
	mem, ret := m.dev.GetMemoryInfo()
	if ret != nvml.SUCCESS {
		return 0, 0, false
	}
	return int(mem.Total / (1 << 20)), int(mem.Used / (1 << 20)), true
}

func (m nvmlMIG) Instances() (gpuInstance, computeInstance int) {
	gi, ret := m.dev.GetGpuInstanceId()
	if ret != nvml.SUCCESS {
		gi = -1
	}
	ci, ret := m.dev.GetComputeInstanceId()
	if ret != nvml.SUCCESS {
		ci = -1
	}
	return gi, ci
}

func (m nvmlMIG) Attributes() (slices, memoryMB int, ok bool) {
	a, ret := m.dev.GetAttributes()
	if ret != nvml.SUCCESS {
		return 0, 0, false
	}
	return int(a.GpuInstanceSliceCount), int(a.MemorySizeMB), true
}