package discovery

import (
	"encoding/json"
	"fmt"
	"math"
	"regexp"
	"slices"
	"strconv"
	"strings"
)

// GPU is one accelerator on this node.
type GPU struct {
	Index int    `json:"index"`
	UUID  string `json:"uuid"`
	Name  string `json:"name"`
	// MemoryTotalMB is the total framebuffer size.
	MemoryTotalMB int `json:"memory_total_mb"`
	// MemoryUsedMB is what is already allocated. Non-zero on an idle node means
	// something other than our engines is using the card (plan §5.1).
	MemoryUsedMB int `json:"memory_used_mb"`
	// UtilizationPercent is how busy the card was over the last sample period.
	//
	// The number people actually watch, and the one nothing collected: memory
	// used says a model is loaded, and this says whether it is doing anything.
	// A fleet of cards at ninety percent occupancy and two percent utilisation
	// is a fleet holding models nobody is asking for, which is invisible
	// without it.
	UtilizationPercent int `json:"utilization_percent,omitempty"`
	// ComputeCapability gates quantization choices: an FP8 profile must not
	// land on a card below 8.9 (plan §6.2).
	ComputeCapability string `json:"compute_capability"`
	NUMANode          int    `json:"numa_node"`
	PCIBusID          string `json:"pci_bus_id,omitempty"`
	DriverVersion     string `json:"driver_version,omitempty"`
	CUDAVersion       string `json:"cuda_version,omitempty"`
	// NVLinkPeers are the indices of GPUs reachable over NVLink. Within a node
	// a TP group should stay inside one NVLink domain.
	NVLinkPeers []int `json:"nvlink_peers,omitempty"`
	// NICPaths is how far each network adapter is from this GPU, by HCA name:
	// PIX, PXB, PHB, NODE or SYS, as nvidia-smi topo reports it. What decides
	// which adapter a GPU's share of a multi-node job should go through.
	NICPaths map[string]string `json:"nic_paths,omitempty"`
	// Fabric is the GPU's multi-node NVLink membership, nil where the
	// driver reports none.
	Fabric *GPUFabric `json:"fabric,omitempty"`
	// Health carries anything that should keep the GPU out of placement.
	Health GPUHealth `json:"health"`
	// MIG is set when this is a MIG instance rather than a whole card: its
	// Index is then MIGIndex's, its UUID the instance's (MIG-…), and its
	// memory the instance's share. Utilisation stays unreported, because no
	// driver measures it per instance.
	MIG *MIGDevice `json:"mig,omitempty"`
}

// GPUHealth records conditions that make a GPU unsafe to schedule onto.
type GPUHealth struct {
	// ECCErrors is the count of uncorrected ECC errors.
	ECCErrors int `json:"ecc_errors"`
	// TemperatureC is the current die temperature.
	TemperatureC int `json:"temperature_c"`
	// Problems lists human-readable reasons to avoid this GPU.
	Problems []string `json:"problems,omitempty"`
}

// OK reports whether the GPU is fit for placement.
func (h GPUHealth) OK() bool { return len(h.Problems) == 0 }

// nonAlnum collapses everything that is not a letter or digit.
var nonAlnum = regexp.MustCompile(`[^a-z0-9]+`)

// Class derives the GPU class used for placement: same model, same memory.
//
// Two cards only belong in the same tensor-parallel group if they are the same
// class, so the derivation has to be stable across driver versions and naming
// quirks ("NVIDIA H100 80GB HBM3" and "H100-SXM5-80GB" are the same class).
func (g GPU) Class() string {
	name := strings.ToLower(g.Name)
	name = strings.TrimPrefix(name, "nvidia ")
	name = strings.TrimPrefix(name, "tesla ")

	// Keep the model token: the first word that contains a digit, which is how
	// NVIDIA names every accelerator (h100, a100, l40s, rtx4090).
	var model string
	for _, tok := range nonAlnum.Split(name, -1) {
		if tok == "" || !strings.ContainsFunc(tok, func(r rune) bool { return r >= '0' && r <= '9' }) {
			continue
		}
		// Skip memory and form-factor tokens, which are not the model.
		if strings.HasSuffix(tok, "gb") || tok == "80gb" || tok == "sxm4" || tok == "sxm5" || tok == "pcie" {
			continue
		}
		model = tok
		break
	}
	if model == "" {
		model = nonAlnum.ReplaceAllString(name, "-")
	}

	// Round to the nearest GiB: vendors report 81559 MiB for an "80GB" card.
	// A MIG instance is named from its card's memory and then its profile,
	// so two instances are one class only when both the card and the slice
	// are the same.
	mem := g.MemoryTotalMB
	if g.MIG != nil {
		mem = g.MIG.ParentMemoryTotalMB
	}
	gb := int(math.Round(float64(mem) / 1024))
	class := fmt.Sprintf("%s-%dg", model, gb)
	if g.MIG != nil {
		return migClass(class, g.MIG.Profile)
	}
	return class
}

// normalizeBusID makes bus ids comparable: NVML reports them uppercase from one
// call and lowercase from another, and sometimes with a shortened domain.
func normalizeBusID(s string) string {
	s = strings.ToLower(strings.TrimSpace(s))
	// "0000:1b:00.0" and "00000000:1B:00.0" are the same device.
	if i := strings.Index(s, ":"); i > 0 {
		s = strings.TrimLeft(s[:i], "0") + s[i:]
	}
	return strings.TrimPrefix(s, ":")
}

// MarshalJSON adds the derived class, so a consumer of the JSON does not have
// to reimplement the derivation to know which GPUs are interchangeable.
func (g GPU) MarshalJSON() ([]byte, error) {
	type alias GPU // prevents recursing back into this method
	return json.Marshal(struct {
		alias
		Class string `json:"class"`
	}{alias(g), g.Class()})
}

// Inventory is everything discovery found on this node.
type Inventory struct {
	GPUs []GPU `json:"gpus"`
	HCAs []HCA `json:"hcas"`
	// Source names the backend that produced the GPU list, so a consumer can
	// tell "no GPUs" from "discovery did not work".
	Source string `json:"source"`
	// MemoryTotalMB is the machine's RAM. Zero means unknown, which is what an
	// agent older than this field reports, and is treated as no constraint.
	MemoryTotalMB int `json:"memory_total_mb,omitempty"`
	// Weights are the model files this node can see. Empty means it was not
	// asked to look.
	Weights []LocalWeights `json:"weights,omitempty"`
	// Storage is the disk behind those weights, one entry per filesystem.
	// Empty means the same thing it means for Weights: not asked.
	Storage []Storage `json:"storage,omitempty"`
	// IMEX is the service a multi-node NVLink domain needs, read where the
	// node has GPUs.
	IMEX IMEX `json:"imex"`
}

// Classes returns the distinct GPU classes present, sorted.
func (inv Inventory) Classes() []string {
	seen := map[string]struct{}{}
	for _, g := range inv.GPUs {
		seen[g.Class()] = struct{}{}
	}
	out := make([]string, 0, len(seen))
	for c := range seen {
		out = append(out, c)
	}
	slices.Sort(out)
	return out
}

// parseCSVInt reads an integer from nvidia-smi's CSV output, tolerating the
// "[N/A]" and "Not Supported" values it emits for unsupported queries.
func parseCSVInt(s string) (int, bool) {
	s = strings.TrimSpace(s)
	s = strings.TrimSuffix(s, " MiB")
	s = strings.TrimSuffix(s, " W")
	if s == "" || strings.HasPrefix(s, "[") || strings.EqualFold(s, "Not Supported") {
		return 0, false
	}
	v, err := strconv.Atoi(s)
	if err != nil {
		return 0, false
	}
	return v, true
}