package platform

import (
	"fmt"
	"os"
	"sort"
	"strconv"
	"strings"
)

// satReadFile is a seam for tests to fake sysfs reads (numa_node files).
var satReadFile = os.ReadFile

// gpuBandwidthSocketGroups splits gpuIndices into per-socket groups (ordered
// by ascending NUMA node ID) for RunNvidiaBandwidthPack. A cross-socket
// peer-to-peer path is a different (and, on platforms without NVLink, far
// less exercised) fault domain than a same-socket one, so testing each
// socket's GPUs in isolation before testing all of them together isolates
// whether a failure is specific to the cross-socket path.
//
// Falls back to a single group containing all of gpuIndices — i.e. no split
// — whenever the NUMA node can't be resolved for every GPU, or all resolve
// to the same node: there's nothing meaningful to split in that case.
func gpuBandwidthSocketGroups(gpuIndices []int, logFunc func(string)) [][]int {
	nodes, err := gpuNUMANodes(gpuIndices)
	if err != nil {
		if logFunc != nil {
			logFunc(fmt.Sprintf("nvbandwidth: could not resolve GPU NUMA nodes (%v); running all GPUs as one group", err))
		}
		return [][]int{gpuIndices}
	}

	byNode := map[int][]int{}
	var unresolved []int
	for _, idx := range gpuIndices {
		node, ok := nodes[idx]
		if !ok {
			if logFunc != nil {
				logFunc(fmt.Sprintf("nvbandwidth: no NUMA node resolved for GPU %d; will fold it into a resolved socket group instead of dropping the split", idx))
			}
			unresolved = append(unresolved, idx)
			continue
		}
		byNode[node] = append(byNode[node], idx)
	}
	// Fewer than two resolved sockets means there's nothing to split either
	// way: every GPU's node is unknown, or every resolved GPU shares one
	// socket. A single unresolved GPU among an otherwise clean multi-socket
	// system shouldn't cost us the split, so only bail out here.
	if len(byNode) < 2 {
		return [][]int{gpuIndices}
	}

	sortedNodes := make([]int, 0, len(byNode))
	for node := range byNode {
		sortedNodes = append(sortedNodes, node)
	}
	sort.Ints(sortedNodes)

	groups := make([][]int, 0, len(sortedNodes))
	for _, node := range sortedNodes {
		groups = append(groups, dedupeSortedIndices(byNode[node]))
	}
	if len(unresolved) > 0 {
		// Fold into the last group rather than running unresolved GPUs in a
		// group of their own — a lone GPU can't run a GPU-to-GPU bandwidth
		// test by itself, and the point of the split is to isolate the
		// sockets we *do* know about, not to also isolate the unknown one.
		last := len(groups) - 1
		groups[last] = dedupeSortedIndices(append(groups[last], unresolved...))
	}
	return groups
}

// gpuNUMANodes resolves the NUMA node each of gpuIndices' GPU is attached to,
// via nvidia-smi's PCI bus ID and the device's sysfs numa_node attribute.
// A GPU missing from the returned map means its node couldn't be resolved.
func gpuNUMANodes(gpuIndices []int) (map[int]int, error) {
	out, err := satExecCommand("nvidia-smi", "--query-gpu=index,pci.bus_id", "--format=csv,noheader,nounits").Output()
	if err != nil {
		return nil, fmt.Errorf("nvidia-smi: %w", err)
	}

	want := make(map[int]struct{}, len(gpuIndices))
	for _, idx := range gpuIndices {
		want[idx] = struct{}{}
	}

	nodes := make(map[int]int, len(gpuIndices))
	for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
		fields := strings.SplitN(line, ",", 2)
		if len(fields) != 2 {
			continue
		}
		idx, err := strconv.Atoi(strings.TrimSpace(fields[0]))
		if err != nil {
			continue
		}
		if _, ok := want[idx]; !ok {
			continue
		}
		bdf := normalizeNvidiaBDF(strings.TrimSpace(fields[1]))
		if node, ok := readPCINumaNode(bdf); ok {
			nodes[idx] = node
		}
	}
	return nodes, nil
}

// normalizeNvidiaBDF converts nvidia-smi's 8-hex-digit-domain PCI bus ID
// ("00000000:05:00.0") to the 4-hex-digit-domain form sysfs paths use
// ("0000:05:00.0").
func normalizeNvidiaBDF(busID string) string {
	domain, rest, ok := strings.Cut(busID, ":")
	if !ok {
		return busID
	}
	if len(domain) > 4 {
		domain = domain[len(domain)-4:]
	}
	return domain + ":" + rest
}

// readPCINumaNode reads a PCI device's NUMA affinity from sysfs. Returns
// false if the attribute is missing/unreadable or reports -1 (no affinity —
// common on single-socket or non-NUMA systems).
func readPCINumaNode(bdf string) (int, bool) {
	data, err := satReadFile("/sys/bus/pci/devices/" + bdf + "/numa_node")
	if err != nil {
		return 0, false
	}
	node, err := strconv.Atoi(strings.TrimSpace(string(data)))
	if err != nil || node < 0 {
		return 0, false
	}
	return node, true
}