// Package fabrictest measures the fabric between two nodes, rail by rail.
//
// Every check on one node — ports up, a subnet manager, the right modules —
// can pass while the two nodes still cannot reach each other: a cable into the
// wrong switch, a partition key, a RoCE network with no route. A replica
// spanning them then starts, and hangs on its first all-reduce with nothing to
// say why. This runs what that all-reduce would ride on — an RDMA write from
// each adapter to its counterpart — and says how it went against the link's
// rate, before a replica does.
package fabrictest

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"strings"
	"time"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
)

// Asker sends one side of a test to a node and waits for its answer.
type Asker interface {
	FabricTest(ctx context.Context, node string, req *llmgwv1.FabricTestRequest) (*llmgwv1.FabricTestReply, error)
}

// Rail is one adapter on each node, paired by their position in the fabric's
// declaration: mlx5_0 with mlx5_0, the way NCCL pairs them.
type Rail struct {
	A, B         string // adapters
	PortA, PortB int
	// RateGbps is the slower of the two links.
	RateGbps int
}

// Result is one rail's test.
type Result struct {
	Rail
	Gbps float64 `json:"gbps"`
	// Verdict is "ok" at 80% of the link's rate or more, "slow" at 50%,
	// "fail" below that or when it did not run.
	Verdict string `json:"verdict"`
	Detail  string `json:"detail,omitempty"`
}

// Report is a whole test.
type Report struct {
	Fabric  string   `json:"fabric"`
	A, B    string   `json:"-"`
	Results []Result `json:"results"`
	// Notes say what the test could not cover.
	Notes []string `json:"notes,omitempty"`
}

// OK reports whether every rail passed.
func (r Report) OK() bool {
	for _, x := range r.Results {
		if x.Verdict != "ok" {
			return false
		}
	}
	return len(r.Results) > 0
}

// Plan pairs the adapters two nodes have on a fabric. With no fabric named,
// the one fabric they share is used.
func Plan(a, b inventory.Node, fabric string) (string, []Rail, []string, error) {
	if a.ID == b.ID {
		return "", nil, nil, errors.New("give two different nodes")
	}
	if fabric == "" {
		var common []string
		for _, f := range a.FabricNames() {
			if b.HasFabric(f) {
				common = append(common, f)
			}
		}
		switch len(common) {
		case 0:
			return "", nil, nil, fmt.Errorf("%s (%s) and %s (%s) declare no fabric in common",
				a.ID, orNone(a.FabricNames()), b.ID, orNone(b.FabricNames()))
		case 1:
			fabric = common[0]
		default:
			return "", nil, nil, fmt.Errorf("%s and %s share %s: name one with --fabric", a.ID, b.ID, strings.Join(common, ", "))
		}
	}
	for _, n := range []inventory.Node{a, b} {
		if !slices.Contains(n.Features, "fabric-test") {
			return "", nil, nil, fmt.Errorf("%s's agent (%s) is too old to run a bandwidth test: upgrade it", n.ID, orNone([]string{n.AgentVersion}))
		}
	}
	ha, hb := a.HCAsFor(fabric), b.HCAsFor(fabric)
	if len(ha) == 0 || len(hb) == 0 {
		return "", nil, nil, fmt.Errorf("%s is not declared on both %s and %s", fabric, a.ID, b.ID)
	}
	var notes []string
	if len(ha) != len(hb) {
		notes = append(notes, fmt.Sprintf("%s has %d adapters on %s and %s has %d: only the first %d are paired",
			a.ID, len(ha), fabric, b.ID, len(hb), min(len(ha), len(hb))))
	}
	var rails []Rail
	for i := range min(len(ha), len(hb)) {
		pa, ra := port(a, ha[i])
		pb, rb := port(b, hb[i])
		rails = append(rails, Rail{A: ha[i], B: hb[i], PortA: pa, PortB: pb, RateGbps: minPositive(ra, rb)})
	}
	return fabric, rails, notes, nil
}

// port is an adapter's first active port and its rate, or port 1 when the
// node reported none.
func port(n inventory.Node, hca string) (int, int) {
	for _, p := range n.IBPorts {
		if p.HCA == hca && strings.EqualFold(p.State, "ACTIVE") {
			return p.Port, p.RateGbps
		}
	}
	return 1, 0
}

// Options for a run.
type Options struct {
	Seconds int
	// BasePort is perftest's TCP exchange port for the first rail; each rail
	// takes the next, so a rail whose server is still closing never meets the
	// next one's.
	BasePort int
	// Settle is how long the server is given to listen before the client
	// dials; the client retries besides.
	Settle time.Duration
}

// Run tests every rail, one after another: rails at once would share the
// PCIe and memory bandwidth they are measuring.
func Run(ctx context.Context, ask Asker, a, b inventory.Node, fabric string, rails []Rail, o Options) Report {
	if o.Seconds <= 0 {
		o.Seconds = 3
	}
	if o.BasePort == 0 {
		o.BasePort = 18515
	}
	rep := Report{Fabric: fabric, A: a.ID, B: b.ID}
	for i, r := range rails {
		rep.Results = append(rep.Results, runRail(ctx, ask, a, b, r, o, o.BasePort+i))
	}
	return rep
}

func runRail(ctx context.Context, ask Asker, a, b inventory.Node, r Rail, o Options, tcp int) Result {
	res := Result{Rail: r, Verdict: "fail"}
	if b.MgmtAddr == "" {
		res.Detail = b.ID + " reported no management address for the client to dial"
		return res
	}
	type answer struct {
		r   *llmgwv1.FabricTestReply
		err error
	}
	server := make(chan answer, 1)
	go func() {
		rep, err := ask.FabricTest(ctx, b.ID, &llmgwv1.FabricTestRequest{Role: "server", Hca: r.B,
			Port: int32(r.PortB), TcpPort: int32(tcp), Seconds: int32(o.Seconds)}) //nolint:gosec // small numbers
		server <- answer{rep, err}
	}()
	select {
	case <-ctx.Done():
		res.Detail = ctx.Err().Error()
		return res
	case <-time.After(o.Settle):
	}
	client, err := ask.FabricTest(ctx, a.ID, &llmgwv1.FabricTestRequest{Role: "client", Hca: r.A,
		Port: int32(r.PortA), Peer: b.MgmtAddr, TcpPort: int32(tcp), Seconds: int32(o.Seconds)}) //nolint:gosec // small numbers
	var srv answer
	select {
	case srv = <-server:
	case <-time.After(time.Duration(o.Seconds)*time.Second + time.Minute):
		srv.err = errors.New("the server side did not finish")
	}
	switch {
	case err != nil:
		res.Detail = a.ID + ": " + err.Error()
	case client.GetError() != "":
		res.Detail = a.ID + ": " + client.GetError()
		// The server's reason is often the better one: it saw no one arrive.
		if srv.r != nil && srv.r.GetError() != "" {
			res.Detail += "; " + b.ID + ": " + srv.r.GetError()
		}
	default:
		res.Gbps = client.GetGbps()
		res.Verdict, res.Detail = grade(res.Gbps, r.RateGbps)
	}
	return res
}

// grade compares a measurement with the link's rate. RDMA writes of 64 KiB
// reach 90–97% of the signalling rate on a healthy link; well under that is
// a degraded cable, a card in a narrow PCIe slot, or traffic crossing a
// socket.
func grade(gbps float64, rate int) (string, string) {
	if gbps <= 0 {
		return "fail", "no traffic got through"
	}
	if rate <= 0 {
		return "ok", "the link's rate is not known to compare with"
	}
	pct := gbps / float64(rate) * 100
	switch {
	case pct >= 80:
		return "ok", fmt.Sprintf("%.0f%% of %d Gb/s", pct, rate)
	case pct >= 50:
		return "slow", fmt.Sprintf("%.0f%% of %d Gb/s: check the cable, and the card's PCIe width and NUMA node", pct, rate)
	default:
		return "fail", fmt.Sprintf("%.0f%% of %d Gb/s", pct, rate)
	}
}

func minPositive(a, b int) int {
	switch {
	case a <= 0:
		return b
	case b <= 0:
		return a
	}
	return min(a, b)
}

func orNone(s []string) string {
	if len(s) == 0 {
		return "none"
	}
	return strings.Join(slices.Clone(s), ", ")
}