// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package base

import (
	"testing"

	"github.com/hashicorp/nomad/client/lib/idset"
	"github.com/hashicorp/nomad/client/lib/numalib"
	"github.com/hashicorp/nomad/client/lib/numalib/hw"
	"github.com/hashicorp/nomad/plugins/base/proto"
	"github.com/shoenig/test/must"
)

func Test_nomadTopologyToProto(t *testing.T) {
	top := &numalib.Topology{
		Distances: numalib.SLIT{{10, 20}, {20, 10}},
		Cores: []numalib.Core{
			{
				SocketID:   0,
				NodeID:     1,
				ID:         2,
				Grade:      numalib.Performance,
				Disable:    false,
				BaseSpeed:  1000,
				MaxSpeed:   5000,
				GuessSpeed: 3000,
			},
		},
		OverrideTotalCompute:   90_000,
		OverrideWitholdCompute: 2000,
	}
	top.SetNodes(idset.From[hw.NodeID]([]hw.NodeID{0, 1}))

	pb := nomadTopologyToProto(top)
	must.Eq(t, &proto.ClientTopology{
		NodeIds: []uint32{0, 1},
		Distances: &proto.ClientTopologySLIT{
			Dimension: 2,
			Values:    []uint32{10, 20, 20, 10},
		},
		Cores: []*proto.ClientTopologyCore{
			{
				SocketId:   0,
				NodeId:     1,
				CoreId:     2,
				CoreGrade:  proto.CoreGrade_Performance,
				Disable:    false,
				BaseSpeed:  1000,
				MaxSpeed:   5000,
				GuessSpeed: 3000,
			},
		},
		OverrideTotalCompute:   90_000,
		OverrideWitholdCompute: 2000,
	}, pb)

	// make sure we don't panic in case of empty nodes, vide
	// https://github.com/hashicorp/nomad/issues/23385
	top2 := &numalib.Topology{}

	pb2 := nomadTopologyToProto(top2)
	must.Eq(t, &proto.ClientTopology{
		NodeIds:                []uint32{},
		Distances:              &proto.ClientTopologySLIT{Dimension: 0, Values: []uint32{}},
		Cores:                  nil,
		OverrideTotalCompute:   0,
		OverrideWitholdCompute: 0,
	}, pb2)
}

func Test_nomadTopologyFromProto(t *testing.T) {
	pb := &proto.ClientTopology{
		NodeIds: []uint32{0, 1},
		Distances: &proto.ClientTopologySLIT{
			Dimension: 2,
			Values:    []uint32{10, 20, 20, 10},
		},
		Cores: []*proto.ClientTopologyCore{
			{
				SocketId:   0,
				NodeId:     1,
				CoreId:     2,
				CoreGrade:  proto.CoreGrade_Performance,
				Disable:    false,
				BaseSpeed:  1000,
				MaxSpeed:   5000,
				GuessSpeed: 3000,
			},
		},
		OverrideTotalCompute:   90_000,
		OverrideWitholdCompute: 2000,
	}
	top := nomadTopologyFromProto(pb)
	expect := &numalib.Topology{
		Distances: numalib.SLIT{{10, 20}, {20, 10}},
		Cores: []numalib.Core{
			{
				SocketID:   0,
				NodeID:     1,
				ID:         2,
				Grade:      numalib.Performance,
				Disable:    false,
				BaseSpeed:  1000,
				MaxSpeed:   5000,
				GuessSpeed: 3000,
			},
		},
		OverrideTotalCompute:   90_000,
		OverrideWitholdCompute: 2000,
	}
	expect.SetNodes(idset.From[hw.NodeID]([]hw.NodeID{0, 1}))
	must.Eq(t, expect, top)
}

func Test_nomadTopologyDistancesToProto(t *testing.T) {
	slit := numalib.SLIT{
		{10, 20},
		{20, 10},
	}

	pb := nomadTopologyDistancesToProto(slit)
	must.Eq(t, 2, pb.Dimension)
	must.Eq(t, []uint32{10, 20, 20, 10}, pb.Values)
}

func Test_nomadTopologyDistanceFromProto(t *testing.T) {
	pb := &proto.ClientTopologySLIT{
		Dimension: 2,
		Values:    []uint32{10, 20, 20, 10},
	}

	slit := nomadTopologyDistancesFromProto(pb)
	must.Eq(t, numalib.SLIT{
		{10, 20},
		{20, 10},
	}, slit)
}

func Test_nomadTopologyCoresToProto(t *testing.T) {
	cores := []numalib.Core{
		{
			SocketID:   0,
			NodeID:     1,
			ID:         3,
			Grade:      numalib.Efficiency,
			Disable:    false,
			BaseSpeed:  1000,
			MaxSpeed:   3000,
			GuessSpeed: 2200,
		},
		{
			SocketID:   2,
			NodeID:     4,
			ID:         9,
			Grade:      numalib.Performance,
			Disable:    true,
			BaseSpeed:  1500,
			MaxSpeed:   5000,
			GuessSpeed: 3500,
		},
	}

	result := nomadTopologyCoresToProto(cores)
	must.Eq(t, []*proto.ClientTopologyCore{{
		SocketId:   0,
		NodeId:     1,
		CoreId:     3,
		CoreGrade:  proto.CoreGrade_Efficiency,
		Disable:    false,
		BaseSpeed:  1000,
		MaxSpeed:   3000,
		GuessSpeed: 2200,
	}, {
		SocketId:   2,
		NodeId:     4,
		CoreId:     9,
		CoreGrade:  proto.CoreGrade_Performance,
		Disable:    true,
		BaseSpeed:  1500,
		MaxSpeed:   5000,
		GuessSpeed: 3500,
	}}, result)
}

func Test_nomadTopologyCoresFromProto(t *testing.T) {
	pbcores := []*proto.ClientTopologyCore{{
		SocketId:   0,
		NodeId:     1,
		CoreId:     3,
		CoreGrade:  proto.CoreGrade_Efficiency,
		Disable:    false,
		BaseSpeed:  1000,
		MaxSpeed:   3000,
		GuessSpeed: 2200,
	}, {
		SocketId:   2,
		NodeId:     4,
		CoreId:     9,
		CoreGrade:  proto.CoreGrade_Performance,
		Disable:    true,
		BaseSpeed:  1500,
		MaxSpeed:   5000,
		GuessSpeed: 3500,
	}}

	cores := nomadTopologyCoresFromProto(pbcores)
	must.Eq(t, []numalib.Core{{
		SocketID:   0,
		NodeID:     1,
		ID:         3,
		Grade:      numalib.Efficiency,
		Disable:    false,
		BaseSpeed:  1000,
		MaxSpeed:   3000,
		GuessSpeed: 2200,
	}, {
		SocketID:   2,
		NodeID:     4,
		ID:         9,
		Grade:      numalib.Performance,
		Disable:    true,
		BaseSpeed:  1500,
		MaxSpeed:   5000,
		GuessSpeed: 3500,
	}}, cores)
}