package agent

import (
	"context"
	"fmt"
	"os/exec"
	"slices"
	"strconv"
	"strings"
	"time"

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

// usableFabrics reports the fabrics this node can actually carry traffic on.
//
// Loud about what it drops. A fabric silently missing from a registration looks
// identical to one never configured, and the operator who typed the adapter
// name is the only person who can fix it — so the reason is logged here, on the
// machine where both halves are visible, rather than inferred later from a
// replica that will not converge.
func (a *Agent) usableFabrics(inv discovery.Inventory) []inventory.FabricAttachment {
	out := a.checkFabrics(inv.HCAs, true)
	a.mu.Lock()
	a.fabricsSent = fabricsFingerprint(out)
	a.mu.Unlock()
	return out
}

// checkFabrics is the declared fabrics this node can carry traffic on, from
// the adapters as they are now. Loud logs every problem; the heartbeat checks
// quietly and speaks only when the answer changes.
func (a *Agent) checkFabrics(hcas []discovery.HCA, loud bool) []inventory.FabricAttachment {
	declared := map[string][]string{}
	for _, f := range a.cfg.Fabrics {
		declared[f.Fabric] = f.HCAs
	}

	var out []inventory.FabricAttachment
	for _, c := range discovery.CheckFabrics(declared, hcas) {
		if loud {
			for _, p := range c.Problems {
				a.log.Error("a declared fabric adapter cannot carry traffic",
					"fabric", c.Fabric, "problem", p)
			}
		}
		if !c.OK() {
			if loud {
				a.log.Error("fabric not advertised; no declared adapter on it can carry traffic",
					"fabric", c.Fabric,
					"consequence", "no multi-node replica will be placed here")
			}
			continue
		}
		out = append(out, inventory.FabricAttachment{
			Fabric: c.Fabric, HCAs: c.Usable, RateGbps: c.RateGbps,
		})
	}
	return out
}

// fabricsFingerprint is enough to notice a fabric, an adapter or a rate
// changing.
func fabricsFingerprint(fs []inventory.FabricAttachment) string {
	var b strings.Builder
	for _, f := range fs {
		fmt.Fprintf(&b, "%s:%s:%d;", f.Fabric, strings.Join(f.HCAs, ","), f.RateGbps)
	}
	return b.String()
}

// sampleIB re-reads the adapters for the heartbeat: every port's state and
// counters, and the fabrics again when they no longer match what the control
// plane was told.
//
// Every heartbeat, because it is a few small files per port, and because a
// port that went down used to be invisible until the agent restarted — the
// node went on being offered multi-node work that would hang on its first
// collective.
func (a *Agent) sampleIB(hb *llmgwv1.Heartbeat) {
	a.mu.Lock()
	had := len(a.inv.HCAs) > 0
	a.mu.Unlock()
	if !had && len(a.cfg.Fabrics) == 0 {
		return
	}
	hcas, err := discovery.DiscoverIB(a.cfg.Discovery.IBRoot)
	if err != nil {
		return
	}
	a.attachSM(hcas)
	for _, h := range hcas {
		for _, p := range h.Ports {
			ip := &llmgwv1.IBPort{
				Hca: h.Name, Port: int32(p.Num), State: p.State, PhysState: p.PhysState,
				RateGbps: int32(p.RateGbps), LinkLayer: string(p.LinkLayer), Counters: p.Counters,
				NodeGuid: h.NodeGUID, PortGuid: p.PortGUID, Network: p.Network, GidPrefix: p.GIDPrefix,
			}
			a.mu.Lock()
			if p.SM != nil {
				ip.SmGuid, ip.SmLid, ip.SmPriority, ip.SmState = p.SM.GUID, int32(p.SM.LID), int32(p.SM.Priority), p.SM.State
			} else {
				ip.SmError = a.smErr[smKey(h.Name, p.Num)]
			}
			for _, r := range a.sms[smKey(h.Name, p.Num)] {
				ip.SubnetManagers = append(ip.SubnetManagers, &llmgwv1.SubnetManager{
					Guid: r.GUID, Lid: int32(r.LID), Priority: int32(r.Priority), State: r.State})
			}
			ip.SubnetManagersError = a.smsErr[smKey(h.Name, p.Num)]
			a.mu.Unlock()
			hb.IbPorts = append(hb.IbPorts, ip)
		}
	}
	a.mu.Lock()
	a.inv.HCAs = hcas // a reconnect registers what is true now
	a.mu.Unlock()
	if len(a.cfg.Fabrics) == 0 {
		return
	}
	now := a.checkFabrics(hcas, false)
	fp := fabricsFingerprint(now)
	a.mu.Lock()
	changed := fp != a.fabricsSent
	a.fabricsSent = fp
	a.mu.Unlock()
	if !changed {
		return
	}
	// Said once, with the reasons, the moment it happens.
	a.checkFabrics(hcas, true)
	a.log.Warn("the fabrics this node can carry traffic on changed", "now", fp)
	hb.FabricsChanged = true
	for _, f := range now {
		hb.Fabrics = append(hb.Fabrics, &llmgwv1.FabricAttachment{
			Fabric: f.Fabric, Hcas: f.HCAs, RateGbps: int32(f.RateGbps),
		})
	}
}

func smKey(hca string, port int) string { return fmt.Sprintf("%s/%d", hca, port) }

// scanFabric is discovery.ScanFabric, replaced in tests.
var scanFabric = discovery.ScanFabric

// querySMs is discovery.QuerySMs, replaced in tests.
var querySMs = func(hca string, port int) ([]discovery.SMRecord, error) {
	return discovery.QuerySMs(context.Background(), hca, port)
}

func splitSMKey(k string) (string, int) {
	i := strings.LastIndexByte(k, '/')
	port, _ := strconv.Atoi(k[i+1:])
	return k[:i], port
}

// querySM is discovery.QuerySM, replaced in tests.
var querySM = func(hca string, port int) (*discovery.SMInfo, error) {
	return discovery.QuerySM(context.Background(), hca, port)
}

// attachSM puts on each active InfiniBand port the subnet manager it sees,
// asking again once a minute and otherwise reusing the last answer.
func (a *Agent) attachSM(hcas []discovery.HCA) {
	a.mu.Lock()
	stale := time.Since(a.smAsked) > time.Minute || a.sm == nil
	a.mu.Unlock()
	if stale {
		sm, errs := map[string]*discovery.SMInfo{}, map[string]string{}
		_ = discovery.AttachSM(hcas, func(hca string, port int) (*discovery.SMInfo, error) {
			got, err := querySM(hca, port)
			if err != nil {
				errs[smKey(hca, port)] = err.Error()
				return nil, err
			}
			sm[smKey(hca, port)] = got
			return got, nil
		})
		// Every subnet manager on each fabric, through one port per master:
		// the same list through every port of one fabric is one query too
		// many for each.
		sms, smsErr := map[string][]discovery.SMRecord{}, map[string]string{}
		asked := map[string]bool{}
		for k, v := range sm {
			if asked[v.GUID] {
				continue
			}
			asked[v.GUID] = true
			hca, port := splitSMKey(k)
			if recs, err := querySMs(hca, port); err != nil {
				smsErr[k] = err.Error()
			} else {
				sms[k] = recs
			}
		}
		a.mu.Lock()
		for k, v := range sm {
			if old := a.sm[k]; old != nil && old.GUID != v.GUID {
				a.log.Warn("the master subnet manager changed", "port", k, "from", old.GUID, "to", v.GUID)
			}
		}
		a.sm, a.smErr, a.smAsked = sm, errs, time.Now()
		a.sms, a.smsErr = sms, smsErr
		a.mu.Unlock()
		return
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	for i := range hcas {
		for j := range hcas[i].Ports {
			if sm := a.sm[smKey(hcas[i].Name, hcas[i].Ports[j].Num)]; sm != nil && hcas[i].Ports[j].Active() {
				hcas[i].Ports[j].SM = sm
			}
		}
	}
}

// ensureDiags installs infiniband-diags on a node with adapters and without
// it. sminfo is what tells which physical fabric a port is on; without it,
// two separate fabrics given one name cannot be caught. Installed the way
// engine dependencies are, from the machine's own package sources.
func (a *Agent) ensureDiags(ctx context.Context) {
	a.mu.Lock()
	hasIB := len(a.inv.HCAs) > 0
	a.mu.Unlock()
	if !hasIB || a.prov == nil {
		return
	}
	if _, err := exec.LookPath("sminfo"); err == nil {
		return
	}
	a.log.Info("installing infiniband-diags, to tell which fabric each port is on")
	pkgs := map[string][]string{"deb": {"infiniband-diags"}, "rpm": {"infiniband-diags"}}
	if err := a.prov.InstallPackages(ctx, pkgs); err != nil {
		a.log.Warn("could not install infiniband-diags; fabrics cannot be told apart by their subnet manager",
			"err", err)
	}
}

// FeatureFabricTest is the agent answering FabricTestRequest.
const FeatureFabricTest = "fabric-test"

// runWriteBW is discovery.RunWriteBW, and haveWriteBW whether it is
// installed; both replaced in tests.
var (
	runWriteBW  = discovery.RunWriteBW
	haveWriteBW = func() bool { _, err := exec.LookPath("ib_write_bw"); return err == nil }
)

// fabricTest runs this node's side of a bandwidth test with another.
func (a *Agent) fabricTest(ctx context.Context, req *llmgwv1.FabricTestRequest) *llmgwv1.FabricTestReply {
	start := time.Now()
	reply := &llmgwv1.FabricTestReply{RequestId: req.GetRequestId()}
	defer func() { reply.TookMs = time.Since(start).Milliseconds() }()

	w := discovery.WriteBW{HCA: req.GetHca(), Port: int(req.GetPort()), GIDIndex: -1,
		TCPPort: int(req.GetTcpPort()), Seconds: int(req.GetSeconds())}
	if req.GetRole() == "client" {
		w.Peer = req.GetPeer()
	}
	a.mu.Lock()
	found := false
	for _, h := range a.inv.HCAs {
		for _, p := range h.Ports {
			if h.Name == w.HCA && p.Num == w.Port {
				found = true
				// RoCE needs its v2 GID named, as NCCL does; InfiniBand uses the LID.
				if p.LinkLayer == discovery.LinkEthernet && p.RoCEv2GID >= 0 {
					w.GIDIndex = p.RoCEv2GID
				}
			}
		}
	}
	a.mu.Unlock()
	if !found {
		reply.Error = fmt.Sprintf("this node has no port %s/%d", w.HCA, w.Port)
		return reply
	}
	if !haveWriteBW() && a.prov != nil {
		a.log.Info("installing perftest, for a bandwidth test with another node")
		pkgs := map[string][]string{"deb": {"perftest"}, "rpm": {"perftest"}}
		if err := a.prov.InstallPackages(ctx, pkgs); err != nil {
			reply.Error = "installing perftest: " + err.Error()
			return reply
		}
	}
	gbps, out, err := runWriteBW(ctx, w)
	reply.Gbps = gbps
	reply.Output = tailLines(out, 12)
	if err != nil {
		reply.Error = err.Error()
	}
	return reply
}

func tailLines(s string, n int) string {
	lines := strings.Split(strings.TrimRight(s, "\n"), "\n")
	if len(lines) > n {
		lines = lines[len(lines)-n:]
	}
	return strings.Join(lines, "\n")
}

// discoverGPUFabric and discoverIMEX are replaced in tests.
var (
	discoverGPUFabric = discovery.DiscoverGPUFabric
	discoverIMEX      = discovery.DiscoverIMEX
)

// sampleNVLink reads, once a minute, which NVLink domain each GPU is joined
// to and whether IMEX is up. nvidia-smi -q takes a second or two, too long
// for every heartbeat, and a fabric that was joined an hour ago rarely
// changes in five seconds — but when a fabric manager restarts it does, for
// every card at once.
func (a *Agent) sampleNVLink(hb *llmgwv1.Heartbeat) {
	a.mu.Lock()
	gpus := slices.Clone(a.inv.GPUs)
	due := time.Since(a.nvlAsked) > time.Minute
	if due {
		a.nvlAsked = time.Now()
	}
	a.mu.Unlock()
	if len(gpus) == 0 || !due {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	fab, err := discoverGPUFabric(ctx, a.cfg.Discovery.NvidiaSMI)
	if err != nil || len(fab) == 0 {
		return
	}
	discovery.ApplyFabric(gpus, fab)
	for _, g := range gpus {
		if g.Fabric == nil {
			continue
		}
		hb.GpuFabric = append(hb.GpuFabric, &llmgwv1.GPUFabricState{
			Index: int32(g.Index), State: g.Fabric.State, Status: g.Fabric.Status,
			ClusterUuid: g.Fabric.ClusterUUID, CliqueId: g.Fabric.CliqueID, Bandwidth: g.Fabric.Bandwidth,
		})
	}
	m := discoverIMEX(ctx)
	hb.Imex = &llmgwv1.IMEX{State: m.State, Channels: int32(m.Channels)}
}