package modelplan

import (
	"strings"
	"testing"
)

const mig3g = "a100-80g-mig-3g.40gb"

func TestAMIGInstanceIsSizedByItsProfileNotItsCard(t *testing.T) {
	for class, want := range map[string]int{
		mig3g: 40, "h100-80g-mig-1g.10gb": 10, "h100-80g-mig-1g.10gb-me": 10, "a100-80g": 80, "a100-80g-mig-": 0,
	} {
		if got := MemFromClass(class); got != want {
			t.Errorf("MemFromClass(%q) = %d, want %d", class, got, want)
		}
	}
}

func TestAModelThatFitsOneMIGInstanceIsPlannedOnOne(t *testing.T) {
	// Four instances free on each of two nodes on a fabric: a 3B fits one,
	// and a 70B, which no single instance holds, is not spread over several
	// — a replica cannot use more than one.
	hw := []Hardware{{Class: mig3g, GPUMemGB: 40, GPUsPerNode: 4, Nodes: 2, Fabric: "ib-a"}}
	fits := Fits(coder3B, hw, 32768, []Precision{BF16})
	if len(fits) == 0 {
		t.Fatal("a 3B does not fit a 40 GB MIG instance")
	}
	for _, f := range fits {
		if f.TP != 1 || f.Nodes != 1 {
			t.Errorf("planned across MIG instances: %s", f.Describe())
		}
	}
	if fits := Fits(llama70B, hw, 8192, []Precision{BF16}); len(fits) != 0 {
		t.Errorf("a 70B was planned on MIG instances: %s", fits[0].Describe())
	}
}

func TestANodeWithACardInMIGModeOffersEachOfItsClasses(t *testing.T) {
	// Card 0 cut into two 3g.40gb instances, card 1 whole: two classes on
	// one node, each with its own count.
	nodes := []PoolNode{{ID: "a100", Class: mig3g, GPUs: 3, FreeGPUs: 3, HostMemGB: 512, HostFreeGB: 400,
		Kinds: []PoolKind{
			{Class: mig3g, GPUMemGB: 40, GPUs: 2, FreeGPUs: 2},
			{Class: "a100-80g", GPUMemGB: 80, GPUs: 1, FreeGPUs: 1},
		}}}
	arts := []Artifact{
		{Name: "q4.gguf", Format: "gguf", Bytes: 9 * gib},
		{Name: "f16.gguf", Format: "gguf", Bytes: 29 * gib},
		{Name: "big.gguf", Format: "gguf", Bytes: 60 * gib},
	}
	opts := Options(arts, shape14B, nodes, 16384)
	got := strings.Join(labels(opts), "; ")
	for _, want := range []string{"q4.gguf 1 × " + mig3g, "f16.gguf 1 × " + mig3g, "big.gguf 1 × a100-80g"} {
		if !strings.Contains(got, want) {
			t.Errorf("missing %q in %s", want, got)
		}
	}
	// Sixty gigabytes is more than an instance holds, and two instances are
	// not one bigger device.
	if strings.Contains(got, "big.gguf 2 × "+mig3g) {
		t.Errorf("planned across two MIG instances: %s", got)
	}
	for _, o := range opts {
		if o.Class == mig3g && (o.Hosts != 1 || !o.FreeNow) {
			t.Errorf("the MIG option counts %d hosts, free %v", o.Hosts, o.FreeNow)
		}
	}
}