package engines

import (
	"slices"
	"testing"
)

func simSpec(rank, nnodes int) ShardSpec {
	return ShardSpec{
		ReplicaID: "rep-1", ShardID: "sh-1", Rank: rank, NNodes: nnodes,
		LeaderAddr: "10.0.0.5", LeaderPort: 20001,
		GPUIndices: []int{0, 1}, HTTPPort: 20000, MgmtInterface: "eth0",
		HCAs: []string{"mlx5_0", "mlx5_1"},
		Profile: Profile{
			Name: "tp4", Engine: "sim", TP: 4, GPUsPerNode: 2, Nodes: nnodes,
			ModelPath: "/models/x", ServedModelName: "big-model",
		},
	}
}

func TestOnlyTheLeaderServes(t *testing.T) {
	// The behaviour that made replicas fail while they were coming up perfectly
	// well: health checking a worker's port, which never opens.
	if !(&Sim{}).ServesHTTP(simSpec(0, 2)) {
		t.Error("the leader does not serve")
	}
	if (&Sim{}).ServesHTTP(simSpec(1, 2)) {
		t.Error("a worker claims to serve")
	}
}

func TestAMultiNodeShardCarriesItsRendezvous(t *testing.T) {
	cmds, err := (&Sim{cfg: SimConfig{Bin: "/usr/local/bin/llmgw"}}).Commands(simSpec(1, 2))
	if err != nil {
		t.Fatal(err)
	}
	if len(cmds) != 1 {
		t.Fatalf("%d commands", len(cmds))
	}
	c := cmds[0]

	if c.Args[0] != "sim-engine" {
		t.Errorf("the engine is a subcommand of llmgw: %v", c.Args)
	}
	for flag, want := range map[string]string{
		"--nnodes": "2", "--node-rank": "1", "--dist-init-addr": "10.0.0.5:20001",
	} {
		if got, _ := flagValue(c.Args, flag); got != want {
			t.Errorf("%s = %q, want %q", flag, got, want)
		}
	}
	// A worker must wait for the leader's port, or the launch is a race.
	if !c.WaitForLeaderPort {
		t.Error("a worker does not wait for the leader")
	}
}

func TestTheLeaderDoesNotWaitForItself(t *testing.T) {
	cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(0, 2))
	if cmds[0].WaitForLeaderPort {
		t.Error("the leader waits for a port it is about to open")
	}
}

func TestASingleNodeShardCarriesNoRendezvous(t *testing.T) {
	cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(0, 1))
	for _, flag := range []string{"--nnodes", "--node-rank", "--dist-init-addr"} {
		if slices.Contains(cmds[0].Args, flag) {
			t.Errorf("a single-node shard was given %s", flag)
		}
	}
}

func TestTheShardIsStillPinnedToItsFabric(t *testing.T) {
	// The simulated engine does not use NCCL, but the environment is generated
	// by the same code that generates it for the ones that do — so a cluster
	// exercised with this engine is exercising that too.
	cmds, _ := (&Sim{cfg: SimConfig{Bin: "x"}}).Commands(simSpec(1, 2))
	env := cmds[0].Env
	for _, want := range []string{
		"NCCL_IB_HCA=mlx5_0,mlx5_1", "NCCL_SOCKET_IFNAME=eth0",
		"CUDA_VISIBLE_DEVICES=0,1", "NCCL_IB_DISABLE=0",
	} {
		if !slices.Contains(env, want) {
			t.Errorf("the environment is missing %s", want)
		}
	}
}