package collector

import (
	"bee/audit/internal/schema"
	"fmt"
	"log/slog"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

var nv5re = regexp.MustCompile(`(?i)^NV(\d+)$`)

// isNVLinkBridgeCandidate returns true for Mellanox PCIe devices that look like
// NVLink bridge mezzanine cards: narrow link (x2), no host net interfaces.
// These are the CPU-side PCIe control plane of the NVSwitch fabric on HGX/DGX systems.
func isNVLinkBridgeCandidate(bdf string, dev schema.HardwarePCIeDevice) bool {
	if !isMellanoxDevice(dev) {
		return false
	}
	if dev.LinkWidth == nil || *dev.LinkWidth > 2 {
		return false
	}
	if len(netIfacesByBDF(bdf)) > 0 {
		return false
	}
	return true
}

// confirmNVLinkBridgeDeviceName checks if the lspci DeviceName for bdf contains
// "NVLINK". This is a targeted single-device call, only executed for candidates
// already pre-filtered by isNVLinkBridgeCandidate.
func confirmNVLinkBridgeDeviceName(bdf string) bool {
	out, err := exec.Command("lspci", "-s", bdf, "-v").Output()
	if err != nil {
		return false
	}
	for _, line := range strings.Split(string(out), "\n") {
		if strings.Contains(strings.ToUpper(strings.TrimSpace(line)), "NVLINK") {
			return true
		}
	}
	return false
}

// markNVLinkBridge overwrites device_class and adds telemetry flags on a detected
// NVLink bridge card. Must be called before applyPCIeLinkSpeedWarning so that the
// correct severity (Critical) is applied.
func markNVLinkBridge(dev *schema.HardwarePCIeDevice) {
	class := "NVLinkBridge"
	dev.DeviceClass = &class
	if dev.Telemetry == nil {
		dev.Telemetry = map[string]any{}
	}
	dev.Telemetry["nvlink_bridge"] = true
}

// enrichNVLinkBridgesWithGPUTopo cross-references NVLink bridge PCIe status with
// the GPU-side NVLink topology reported by nvidia-smi. For each bridge device it
// adds nvlink_topo_all_active and nvlink_topo_min_links to the telemetry, and
// upgrades a degraded-link Warning to Critical when the fabric is also affected.
func enrichNVLinkBridgesWithGPUTopo(devs []schema.HardwarePCIeDevice) []schema.HardwarePCIeDevice {
	hasBridge := false
	for _, d := range devs {
		if d.DeviceClass != nil && *d.DeviceClass == "NVLinkBridge" {
			hasBridge = true
			break
		}
	}
	if !hasBridge {
		return devs
	}

	topo, err := queryNVIDIANVLinkTopo()
	if err != nil {
		slog.Info("nvlink-bridge: nvidia-smi topo unavailable, skipping cross-reference", "err", err)
		return devs
	}

	for i := range devs {
		if devs[i].DeviceClass == nil || *devs[i].DeviceClass != "NVLinkBridge" {
			continue
		}
		if devs[i].Telemetry == nil {
			devs[i].Telemetry = map[string]any{}
		}
		devs[i].Telemetry["nvlink_topo_all_active"] = topo.AllActive
		devs[i].Telemetry["nvlink_topo_min_links"] = topo.MinNVLinks
		devs[i].Telemetry["nvlink_topo_gpu_count"] = topo.GPUCount

		// If the bridge PCIe is already degraded AND the fabric is also degraded
		// (missing NVLink connections), escalate to Critical.
		if devs[i].Status != nil && *devs[i].Status == statusCritical && !topo.AllActive {
			devs[i].Telemetry["nvlink_fabric_affected"] = true
		}
	}

	slog.Info("nvlink-bridge: topo cross-reference applied",
		"gpu_count", topo.GPUCount,
		"all_active", topo.AllActive,
		"min_links", topo.MinNVLinks,
	)
	return devs
}

// nvlinkTopoResult summarises the GPU NVLink connectivity matrix.
type nvlinkTopoResult struct {
	GPUCount   int
	AllActive  bool // true if every GPU pair has at least one NVLink bond
	MinNVLinks int  // minimum NVLink bonds seen across any GPU pair (0 = some pair disconnected)
}

// queryNVIDIANVLinkTopo runs nvidia-smi topo -m and parses the NVLink matrix.
func queryNVIDIANVLinkTopo() (nvlinkTopoResult, error) {
	out, err := exec.Command("nvidia-smi", "topo", "-m").Output()
	if err != nil {
		return nvlinkTopoResult{}, err
	}
	return parseNVIDIATopologyMatrix(string(out)), nil
}

// locateGPUTopologyColumns finds the header line of a nvidia-smi topo -m
// matrix and the 0-based field indices (excluding the row label) that
// correspond to GPU columns. Returns headerIdx=-1 if fewer than 2 GPU columns
// are found.
func locateGPUTopologyColumns(lines []string) (headerIdx int, gpuColIndices []int, gpuCount int) {
	headerIdx = -1
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "GPU0") {
			parts := strings.Fields(trimmed)
			for j, col := range parts {
				if strings.HasPrefix(col, "GPU") {
					gpuColIndices = append(gpuColIndices, j)
				}
			}
			gpuCount = len(gpuColIndices)
			if gpuCount >= 2 {
				headerIdx = i
			}
			break
		}
	}
	if headerIdx < 0 {
		gpuColIndices = nil
		gpuCount = 0
	}
	return headerIdx, gpuColIndices, gpuCount
}

// parseNVIDIATopologyMatrix extracts the minimum NVLink bond count from the
// nvidia-smi topo -m matrix.
//
// Format (abbreviated):
//
//	         GPU0  GPU1 ... NIC0 NIC1
//	GPU0      X   NV18 ... NODE NODE
//	GPU1     NV18   X  ... NODE NODE
//	NIC0     NODE  NODE...   X   PIX
//
// The header row starts with "GPU0"; its columns may include non-GPU entries
// (NIC, CPU) which are ignored. Only GPU×GPU cells containing NV# values are
// counted. X is self; non-NV tokens (NODE, SYS, PHB, PIX) are skipped.
func parseNVIDIATopologyMatrix(raw string) nvlinkTopoResult {
	lines := strings.Split(raw, "\n")
	headerIdx, gpuColIndices, gpuCount := locateGPUTopologyColumns(lines)
	if headerIdx < 0 {
		return nvlinkTopoResult{}
	}

	minLinks := -1 // -1 = no NV pair seen yet
	allActive := true

	for _, line := range lines[headerIdx+1:] {
		trimmed := strings.TrimSpace(line)
		if !strings.HasPrefix(trimmed, "GPU") {
			continue
		}
		cells := strings.Fields(trimmed)
		// cells[0] is the row label (e.g. "GPU0"); cells[1..] are column values.
		// gpuColIndices are 0-based within the header fields, so they map to
		// cells[idx+1] in the data rows (shift by 1 for the row label).
		for _, colIdx := range gpuColIndices {
			dataIdx := colIdx + 1
			if dataIdx >= len(cells) {
				continue
			}
			cell := cells[dataIdx]
			m := nv5re.FindStringSubmatch(cell)
			if len(m) != 2 {
				continue
			}
			n, err := strconv.Atoi(m[1])
			if err != nil {
				continue
			}
			if n == 0 {
				allActive = false
			}
			if minLinks < 0 || n < minLinks {
				minLinks = n
			}
		}
	}

	if minLinks < 0 {
		minLinks = 0
	}

	return nvlinkTopoResult{
		GPUCount:   gpuCount,
		AllActive:  allActive && minLinks > 0,
		MinNVLinks: minLinks,
	}
}

// parseCrossNUMAPeers scans a nvidia-smi topo -m matrix for GPU pairs whose
// only path is "SYS" — traversing PCIe as well as the SMP interconnect
// between NUMA nodes (e.g. QPI/UPI). This is the slowest possible GPU-GPU
// path and, on servers where GPUs are only bridged pairwise via NVLink
// bridge (no switched NVLink fabric), it is exactly the hop that traffic
// between different bridge pairs has to cross. Returns a map from GPU index
// to the peer GPU indices reachable only via this cross-NUMA path.
func parseCrossNUMAPeers(raw string) map[int][]int {
	lines := strings.Split(raw, "\n")
	headerIdx, gpuColIndices, _ := locateGPUTopologyColumns(lines)
	if headerIdx < 0 {
		return nil
	}

	// colIdx (0-based within header fields) -> GPU index, in header order.
	colIdxToGPU := make(map[int]int, len(gpuColIndices))
	for gpuIdx, colIdx := range gpuColIndices {
		colIdxToGPU[colIdx] = gpuIdx
	}

	peers := make(map[int][]int)
	rowGPU := -1
	for _, line := range lines[headerIdx+1:] {
		trimmed := strings.TrimSpace(line)
		if !strings.HasPrefix(trimmed, "GPU") {
			continue
		}
		rowGPU++
		cells := strings.Fields(trimmed)
		for _, colIdx := range gpuColIndices {
			dataIdx := colIdx + 1
			if dataIdx >= len(cells) {
				continue
			}
			colGPU := colIdxToGPU[colIdx]
			if colGPU == rowGPU {
				continue
			}
			if strings.EqualFold(cells[dataIdx], "SYS") {
				peers[rowGPU] = append(peers[rowGPU], colGPU)
			}
		}
	}
	if len(peers) == 0 {
		return nil
	}
	return peers
}

// enrichGPUCrossNUMATopology flags GPUs that reach one or more peer GPUs only
// via a cross-NUMA-node PCIe hop ("SYS" in nvidia-smi topo -m). Unlike
// enrichNVLinkBridgesWithGPUTopo, this does not require an NVLink bridge PCIe
// device to be present: it applies to any multi-GPU box, since the weak point
// it detects is the path *between* NVLink-bridged pairs (or between GPUs with
// no NVLink at all), not the bridge itself.
func enrichGPUCrossNUMATopology(devs []schema.HardwarePCIeDevice) []schema.HardwarePCIeDevice {
	gpuByBDF, err := queryNVIDIAGPUs()
	if err != nil || len(gpuByBDF) < 2 {
		return devs
	}

	out, err := exec.Command("nvidia-smi", "topo", "-m").Output()
	if err != nil {
		slog.Info("gpu-topology: nvidia-smi topo unavailable, skipping cross-NUMA check", "err", err)
		return devs
	}
	peers := parseCrossNUMAPeers(string(out))
	if len(peers) == 0 {
		return devs
	}

	bdfToIndex := make(map[string]int, len(gpuByBDF))
	for bdf, info := range gpuByBDF {
		bdfToIndex[bdf] = info.Index
	}

	for i := range devs {
		if devs[i].BDF == nil {
			continue
		}
		idx, ok := bdfToIndex[normalizePCIeBDF(*devs[i].BDF)]
		if !ok {
			continue
		}
		peerList, ok := peers[idx]
		if !ok {
			continue
		}
		if devs[i].Telemetry == nil {
			devs[i].Telemetry = map[string]any{}
		}
		devs[i].Telemetry["nvlink_cross_numa_peers"] = peerList
		if devs[i].Status == nil || *devs[i].Status == statusOK {
			warn := statusWarning
			devs[i].Status = &warn
		}
		if devs[i].ErrorDescription == nil {
			devs[i].ErrorDescription = stringPtr(fmt.Sprintf(
				"GPU %d reaches GPU(s) %v only via a cross-NUMA-node PCIe path (SYS) — expect reduced bandwidth/increased latency for tensor-parallel workloads spanning these GPUs",
				idx, peerList))
		}
	}

	slog.Info("gpu-topology: cross-NUMA peers detected", "affected_gpus", len(peers))
	return devs
}