package platform

import (
	"reflect"
	"testing"
)

// realTopoTwoSocketSixGPU is a real "nvidia-smi topo -m" capture (CG480-S6053,
// 6 GPUs across 2 sockets: GPU0-3 on one, GPU4-5 on the other), ANSI escapes
// included as nvidia-smi actually emits them in the header row.
const realTopoTwoSocketSixGPU = "\t\x1b[4mGPU0\tGPU1\tGPU2\tGPU3\tGPU4\tGPU5\tNIC0\tNIC1\tNIC2\tNIC3\tCPU Affinity\tNUMA Affinity\tGPU NUMA ID\x1b[0m\n" +
	"GPU0\t X \tPIX\tNODE\tNODE\tSYS\tSYS\tNODE\tNODE\tPIX\tPIX\t0-95,192-287\t0\t\tN/A\n" +
	"GPU1\tPIX\t X \tNODE\tNODE\tSYS\tSYS\tNODE\tNODE\tPIX\tPIX\t0-95,192-287\t0\t\tN/A\n" +
	"GPU2\tNODE\tNODE\t X \tPIX\tSYS\tSYS\tPIX\tPIX\tNODE\tNODE\t0-95,192-287\t0\t\tN/A\n" +
	"GPU3\tNODE\tNODE\tPIX\t X \tSYS\tSYS\tPIX\tPIX\tNODE\tNODE\t0-95,192-287\t0\t\tN/A\n" +
	"GPU4\tSYS\tSYS\tSYS\tSYS\t X \tPIX\tSYS\tSYS\tSYS\tSYS\t96-191,288-383\t1\t\tN/A\n" +
	"GPU5\tSYS\tSYS\tSYS\tSYS\tPIX\t X \tSYS\tSYS\tSYS\tSYS\t96-191,288-383\t1\t\tN/A\n"

// realTopoSingleSocketFourGPU simulates a single-socket host (all GPUs share
// one CPU Affinity range) — no cross-socket stage should be produced.
const realTopoSingleSocketFourGPU = "GPU0\tGPU1\tGPU2\tGPU3\tCPU Affinity\tNUMA Affinity\n" +
	"GPU0\t X \tPIX\tNODE\tNODE\t0-31\t0\n" +
	"GPU1\tPIX\t X \tNODE\tNODE\t0-31\t0\n" +
	"GPU2\tNODE\tNODE\t X \tPIX\t0-31\t0\n" +
	"GPU3\tNODE\tNODE\tPIX\t X \t0-31\t0\n"

func TestParseNvidiaSocketGroupsTwoSockets(t *testing.T) {
	groups := ParseNvidiaSocketGroups(realTopoTwoSocketSixGPU)
	if len(groups) != 2 {
		t.Fatalf("groups=%+v, want 2", groups)
	}
	if !reflect.DeepEqual(groups[0].GPUIndices, []int{0, 1, 2, 3}) {
		t.Fatalf("group0=%v, want [0 1 2 3]", groups[0].GPUIndices)
	}
	if !reflect.DeepEqual(groups[1].GPUIndices, []int{4, 5}) {
		t.Fatalf("group1=%v, want [4 5]", groups[1].GPUIndices)
	}
	if groups[0].CPUAffinity == groups[1].CPUAffinity {
		t.Fatalf("expected distinct CPU affinities, got %q for both", groups[0].CPUAffinity)
	}
}

func TestNvidiaProgressiveBandwidthGroupsTwoSockets(t *testing.T) {
	socketGroups := ParseNvidiaSocketGroups(realTopoTwoSocketSixGPU)
	stages := NvidiaProgressiveBandwidthGroups(socketGroups)

	want := []NvidiaBandwidthTestGroup{
		{Label: "same-socket-1", GPUIndices: []int{0, 1}},
		{Label: "same-socket-2", GPUIndices: []int{4, 5}},
		{Label: "cross-socket", GPUIndices: []int{0, 4}},
		{Label: "all", GPUIndices: []int{0, 1, 2, 3, 4, 5}},
	}
	if len(stages) != len(want) {
		t.Fatalf("stages=%+v, want %+v", stages, want)
	}
	for i, w := range want {
		if stages[i].Label != w.Label || !reflect.DeepEqual(stages[i].GPUIndices, w.GPUIndices) {
			t.Fatalf("stage %d = %+v, want %+v", i, stages[i], w)
		}
	}
}

func TestNvidiaProgressiveBandwidthGroupsSingleSocket(t *testing.T) {
	socketGroups := ParseNvidiaSocketGroups(realTopoSingleSocketFourGPU)
	stages := NvidiaProgressiveBandwidthGroups(socketGroups)

	want := []NvidiaBandwidthTestGroup{
		{Label: "same-socket-1", GPUIndices: []int{0, 1}},
		{Label: "all", GPUIndices: []int{0, 1, 2, 3}},
	}
	if len(stages) != len(want) {
		t.Fatalf("stages=%+v, want %+v (no cross-socket stage on a single-socket host)", stages, want)
	}
	for i, w := range want {
		if stages[i].Label != w.Label || !reflect.DeepEqual(stages[i].GPUIndices, w.GPUIndices) {
			t.Fatalf("stage %d = %+v, want %+v", i, stages[i], w)
		}
	}
}

func TestParseNvidiaSocketGroupsEmptyOnGarbage(t *testing.T) {
	if got := ParseNvidiaSocketGroups("not a topology matrix"); len(got) != 0 {
		t.Fatalf("got %+v, want empty", got)
	}
}