package platform

import (
	"regexp"
	"sort"
	"strconv"
	"strings"
)

// nvidiaCPUAffinityRe matches an nvidia-smi "topo -m" CPU Affinity cell, e.g.
// "0-95,192-287" or a plain "0". Matched by shape rather than column
// position: NIC count (and therefore column offsets) varies per host, but
// this is the first token after the GPU/NIC relation cells (X/PIX/NODE/SYS)
// that looks like a core range list, on every layout seen so far.
var nvidiaCPUAffinityRe = regexp.MustCompile(`^[0-9]+(-[0-9]+)?(,[0-9]+(-[0-9]+)?)*$`)

// NvidiaSocketGroup is every GPU index sharing one CPU Affinity range in an
// "nvidia-smi topo -m" matrix — a proxy for "these GPUs are on the same CPU
// socket/NUMA node".
type NvidiaSocketGroup struct {
	CPUAffinity string
	GPUIndices  []int
}

// ParseNvidiaSocketGroups groups GPU indices from an "nvidia-smi topo -m"
// matrix by CPU Affinity, so a scenario can pick "a pair on one socket,
// then the other, then one cross-socket pair" without gpu_indices hardcoded
// per host — topology (which GPUs share a socket) differs machine to
// machine, so a scenario file can't bake this in the way it can bake in "run
// on all GPUs".
func ParseNvidiaSocketGroups(raw string) []NvidiaSocketGroup {
	lines := strings.Split(nvidiaNVLinkANSIRe.ReplaceAllString(raw, ""), "\n")
	headerIdx := -1
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "GPU0") {
			headerIdx = i
			break
		}
	}
	if headerIdx < 0 {
		return nil
	}

	order := map[string]int{}
	groups := map[string][]int{}
	for _, line := range lines[headerIdx+1:] {
		trimmed := strings.TrimSpace(line)
		if !strings.HasPrefix(trimmed, "GPU") {
			continue
		}
		cells := strings.Fields(trimmed)
		if len(cells) < 2 {
			continue
		}
		rowGPU, err := strconv.Atoi(strings.TrimPrefix(cells[0], "GPU"))
		if err != nil {
			continue
		}
		affinity := ""
		for _, cell := range cells[1:] {
			if nvidiaCPUAffinityRe.MatchString(cell) {
				affinity = cell
				break
			}
		}
		if affinity == "" {
			continue
		}
		if _, ok := order[affinity]; !ok {
			order[affinity] = len(order)
		}
		groups[affinity] = append(groups[affinity], rowGPU)
	}

	out := make([]NvidiaSocketGroup, 0, len(groups))
	for affinity, indices := range groups {
		sort.Ints(indices)
		out = append(out, NvidiaSocketGroup{CPUAffinity: affinity, GPUIndices: indices})
	}
	sort.Slice(out, func(i, j int) bool {
		return out[i].GPUIndices[0] < out[j].GPUIndices[0]
	})
	return out
}

// NvidiaBandwidthTestGroup is one stage of a progressive multi-GPU bandwidth
// test: a label and the GPU indices to run it against.
type NvidiaBandwidthTestGroup struct {
	Label      string
	GPUIndices []int
}

// NvidiaProgressiveBandwidthGroups turns socket groups into an escalating
// test plan: a pair within each socket that has one, then one cross-socket
// pair (first two sockets' lowest-indexed GPU each), then every GPU. Lets a
// scenario narrow down whether a failure needs the full GPU count or already
// reproduces on a single cross-socket pair, instead of only ever testing
// "all GPUs at once".
func NvidiaProgressiveBandwidthGroups(socketGroups []NvidiaSocketGroup) []NvidiaBandwidthTestGroup {
	var out []NvidiaBandwidthTestGroup
	var allGPUs []int
	var crossSocketPair []int

	for i, sg := range socketGroups {
		allGPUs = append(allGPUs, sg.GPUIndices...)
		if len(sg.GPUIndices) >= 2 {
			out = append(out, NvidiaBandwidthTestGroup{
				Label:      "same-socket-" + strconv.Itoa(i+1),
				GPUIndices: []int{sg.GPUIndices[0], sg.GPUIndices[1]},
			})
		}
		if len(crossSocketPair) < 2 {
			crossSocketPair = append(crossSocketPair, sg.GPUIndices[0])
		}
	}
	if len(crossSocketPair) == 2 {
		out = append(out, NvidiaBandwidthTestGroup{Label: "cross-socket", GPUIndices: crossSocketPair})
	}
	if len(allGPUs) > 0 {
		sort.Ints(allGPUs)
		out = append(out, NvidiaBandwidthTestGroup{Label: "all", GPUIndices: allGPUs})
	}
	return out
}