package collector

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

type nvidiaGPUInfo struct {
	Index              int
	BDF                string
	Name               string
	Serial             string
	VBIOS              string
	TemperatureC       *float64
	PowerW             *float64
	ECCUncorrected     *int64
	ECCCorrected       *int64
	HWSlowdown         *bool
	PCIeLinkGenCurrent *int
	PCIeLinkGenMax     *int
	PCIeLinkWidthCur   *int
	PCIeLinkWidthMax   *int
}

// enrichPCIeWithNVIDIA enriches NVIDIA PCIe devices with data from nvidia-smi.
// If the driver/tool is unavailable, NVIDIA devices get Unknown status.
func enrichPCIeWithNVIDIA(devs []schema.HardwarePCIeDevice) []schema.HardwarePCIeDevice {
	if !hasNVIDIADevices(devs) {
		return devs
	}
	gpuByBDF, err := queryNVIDIAGPUs()
	if err != nil {
		slog.Info("nvidia: enrichment skipped", "err", err)
		return enrichPCIeWithNVIDIAData(devs, nil, false)
	}
	devs = enrichPCIeWithNVIDIAData(devs, gpuByBDF, true)
	return enrichPCIeWithNVIDIANVLinks(devs)
}

// enrichPCIeWithNVIDIANVLinks attaches per-link NVLink status (nvidia-smi
// nvlink -s) and error counters (nvidia-smi nvlink -e) to each GPU's
// HardwarePCIeDevice entry, keyed by the "nvidia_gpu_index" telemetry set by
// enrichPCIeWithNVIDIAData. Independent of NVSwitch/fabric-manager detection
// so it also covers direct GPU-to-GPU bridge boards with no switch present.
func enrichPCIeWithNVIDIANVLinks(devs []schema.HardwarePCIeDevice) []schema.HardwarePCIeDevice {
	statusByGPU, statusErr := nvlinkStatusFn()
	if statusErr != nil {
		slog.Info("nvidia: nvlink -s unavailable, skipping nvlink enrichment", "err", statusErr)
		return devs
	}
	errorsByGPU, errorsErr := nvlinkErrorsFn()
	if errorsErr != nil {
		slog.Info("nvidia: nvlink -e unavailable", "err", errorsErr)
	}

	for i := range devs {
		if devs[i].Telemetry == nil {
			continue
		}
		idx, ok := devs[i].Telemetry["nvidia_gpu_index"].(int)
		if !ok {
			continue
		}
		ports, ok := statusByGPU[idx]
		if !ok {
			continue
		}
		for j := range ports {
			if counters, ok := errorsByGPU[idx][ports[j].Index]; ok {
				ports[j].ReplayErrors = &counters.Replay
				ports[j].RecoveryErrors = &counters.Recovery
				ports[j].CRCErrors = &counters.CRC
			}
		}
		devs[i].NVLinks = ports
	}
	return devs
}

func hasNVIDIADevices(devs []schema.HardwarePCIeDevice) bool {
	for _, dev := range devs {
		if isNVIDIADevice(dev) {
			return true
		}
	}
	return false
}

func enrichPCIeWithNVIDIAData(devs []schema.HardwarePCIeDevice, gpuByBDF map[string]nvidiaGPUInfo, driverLoaded bool) []schema.HardwarePCIeDevice {
	enriched := 0
	for i := range devs {
		if !isNVIDIADevice(devs[i]) {
			continue
		}

		if !driverLoaded {
			setPCIeFallback(&devs[i])
			continue
		}

		bdf := ""
		if devs[i].BDF != nil {
			bdf = normalizePCIeBDF(*devs[i].BDF)
		}
		info, ok := gpuByBDF[bdf]
		if !ok {
			setPCIeFallback(&devs[i])
			continue
		}

		if v := strings.TrimSpace(info.Name); v != "" {
			devs[i].Model = &v
		}
		if v := strings.TrimSpace(info.Serial); v != "" {
			devs[i].SerialNumber = &v
		}
		if v := strings.TrimSpace(info.VBIOS); v != "" {
			devs[i].Firmware = &v
		}

		status := statusOK
		if info.ECCUncorrected != nil && *info.ECCUncorrected > 0 {
			status = statusWarning
			devs[i].ErrorDescription = stringPtr("GPU reports uncorrected ECC errors")
		}
		devs[i].Status = &status
		injectNVIDIATelemetry(&devs[i], info)
		enriched++
	}

	if driverLoaded {
		slog.Info("nvidia: enriched", "count", enriched)
	}
	return devs
}

func queryNVIDIAGPUs() (map[string]nvidiaGPUInfo, error) {
	out, err := exec.Command(
		"nvidia-smi",
		"--query-gpu=index,pci.bus_id,name,serial,vbios_version,temperature.gpu,power.draw,ecc.errors.uncorrected.aggregate.total,ecc.errors.corrected.aggregate.total,clocks_throttle_reasons.hw_slowdown,pcie.link.gen.current,pcie.link.gen.max,pcie.link.width.current,pcie.link.width.max",
		"--format=csv,noheader,nounits",
	).Output()
	if err != nil {
		return nil, err
	}
	return parseNVIDIASMIQuery(string(out))
}

func parseNVIDIASMIQuery(raw string) (map[string]nvidiaGPUInfo, error) {
	r := csv.NewReader(strings.NewReader(raw))
	r.TrimLeadingSpace = true
	r.FieldsPerRecord = -1

	records, err := r.ReadAll()
	if err != nil {
		return nil, err
	}

	result := make(map[string]nvidiaGPUInfo)
	for _, rec := range records {
		if len(rec) == 0 {
			continue
		}
		if len(rec) < 14 {
			return nil, fmt.Errorf("unexpected nvidia-smi columns: got %d, want 14", len(rec))
		}

		bdf := normalizePCIeBDF(rec[1])
		if bdf == "" {
			continue
		}

		info := nvidiaGPUInfo{
			Index:              parseRequiredInt(rec[0]),
			BDF:                bdf,
			Name:               strings.TrimSpace(rec[2]),
			Serial:             strings.TrimSpace(rec[3]),
			VBIOS:              strings.TrimSpace(rec[4]),
			TemperatureC:       parseMaybeFloat(rec[5]),
			PowerW:             parseMaybeFloat(rec[6]),
			ECCUncorrected:     parseMaybeInt64(rec[7]),
			ECCCorrected:       parseMaybeInt64(rec[8]),
			HWSlowdown:         parseMaybeBool(rec[9]),
			PCIeLinkGenCurrent: parseMaybeInt(rec[10]),
			PCIeLinkGenMax:     parseMaybeInt(rec[11]),
			PCIeLinkWidthCur:   parseMaybeInt(rec[12]),
			PCIeLinkWidthMax:   parseMaybeInt(rec[13]),
		}
		result[bdf] = info
	}

	return result, nil
}

func parseMaybeFloat(v string) *float64 {
	v = strings.TrimSpace(v)
	if v == "" || strings.EqualFold(v, "n/a") || strings.EqualFold(v, "not supported") || strings.EqualFold(v, "[not supported]") {
		return nil
	}
	n, err := strconv.ParseFloat(v, 64)
	if err != nil {
		return nil
	}
	return &n
}

func parseMaybeInt64(v string) *int64 {
	v = strings.TrimSpace(v)
	if v == "" || strings.EqualFold(v, "n/a") || strings.EqualFold(v, "not supported") || strings.EqualFold(v, "[not supported]") {
		return nil
	}
	n, err := strconv.ParseInt(v, 10, 64)
	if err != nil {
		return nil
	}
	return &n
}

func parseMaybeInt(v string) *int {
	v = strings.TrimSpace(v)
	if v == "" || strings.EqualFold(v, "n/a") || strings.EqualFold(v, "not supported") || strings.EqualFold(v, "[not supported]") {
		return nil
	}
	n, err := strconv.Atoi(v)
	if err != nil {
		return nil
	}
	return &n
}

func parseRequiredInt(v string) int {
	n, err := strconv.Atoi(strings.TrimSpace(v))
	if err != nil {
		return 0
	}
	return n
}

func pcieLinkGenLabel(gen int) string {
	return fmt.Sprintf("Gen%d", gen)
}

func parseMaybeBool(v string) *bool {
	v = strings.TrimSpace(strings.ToLower(v))
	switch v {
	case "active", "enabled", "true", "1":
		b := true
		return &b
	case "not active", "disabled", "false", "0":
		b := false
		return &b
	default:
		return nil
	}
}

func normalizePCIeBDF(bdf string) string {
	bdf = strings.TrimSpace(strings.ToLower(bdf))
	if bdf == "" {
		return ""
	}
	parts := strings.Split(bdf, ":")
	if len(parts) == 3 {
		domain := parts[0]
		if len(domain) > 4 {
			domain = domain[len(domain)-4:]
		}
		return domain + ":" + parts[1] + ":" + parts[2]
	}
	if len(parts) == 2 {
		return "0000:" + parts[0] + ":" + parts[1]
	}
	return bdf
}

func isNVIDIADevice(dev schema.HardwarePCIeDevice) bool {
	return dev.VendorID != nil && *dev.VendorID == NvidiaVendorID
}

func setPCIeFallback(dev *schema.HardwarePCIeDevice) {
	status := statusUnknown
	dev.Status = &status
}

func injectNVIDIATelemetry(dev *schema.HardwarePCIeDevice, info nvidiaGPUInfo) {
	if dev.Telemetry == nil {
		dev.Telemetry = map[string]any{}
	}
	dev.Telemetry["nvidia_gpu_index"] = info.Index
	if info.TemperatureC != nil {
		dev.TemperatureC = info.TemperatureC
	}
	if info.PowerW != nil {
		dev.PowerW = info.PowerW
	}
	if info.ECCUncorrected != nil {
		dev.ECCUncorrectedTotal = info.ECCUncorrected
	}
	if info.ECCCorrected != nil {
		dev.ECCCorrectedTotal = info.ECCCorrected
	}
	if info.HWSlowdown != nil {
		dev.HWSlowdown = info.HWSlowdown
	}
	// Override PCIe link speed/width with nvidia-smi driver values.
	// sysfs current_link_speed reflects the instantaneous physical link state and
	// can show Gen1 when the GPU is idle due to ASPM power management. The driver
	// knows the negotiated speed regardless of the current power state.
	if info.PCIeLinkGenCurrent != nil {
		s := pcieLinkGenLabel(*info.PCIeLinkGenCurrent)
		dev.LinkSpeed = &s
	}
	if info.PCIeLinkGenMax != nil {
		s := pcieLinkGenLabel(*info.PCIeLinkGenMax)
		dev.MaxLinkSpeed = &s
	}
	if info.PCIeLinkWidthCur != nil {
		dev.LinkWidth = info.PCIeLinkWidthCur
	}
	if info.PCIeLinkWidthMax != nil {
		dev.MaxLinkWidth = info.PCIeLinkWidthMax
	}
}

var (
	nvlinkGPUHeaderRe    = regexp.MustCompile(`^GPU (\d+):`)
	nvlinkSpeedLineRe    = regexp.MustCompile(`^Link (\d+):\s*([\d.]+)\s*GB/s`)
	nvlinkInactiveRe     = regexp.MustCompile(`^Link (\d+):\s*<inactive>`)
	nvlinkErrorCounterRe = regexp.MustCompile(`^Link (\d+):\s*(Replay|Recovery|CRC) Errors:\s*(\d+)`)
)

// nvlinkErrorCounters holds the per-link error counters reported by
// "nvidia-smi nvlink -e" for one GPU.
type nvlinkErrorCounters struct {
	Replay, Recovery, CRC int64
}

// nvlinkStatusFn and nvlinkErrorsFn are swappable for testing.
var (
	nvlinkStatusFn = queryNVIDIANVLinkStatusByGPU
	nvlinkErrorsFn = queryNVIDIANVLinkErrorsByGPU
)

// queryNVIDIANVLinkStatusByGPU runs "nvidia-smi nvlink -s" and returns each
// GPU's NVLink ports keyed by GPU index (as printed in the "GPU N:" header,
// matching the index nvidia-smi --query-gpu also reports).
func queryNVIDIANVLinkStatusByGPU() (map[int][]schema.HardwareNVLinkPort, error) {
	out, err := exec.Command("nvidia-smi", "nvlink", "-s").Output()
	if err != nil {
		return nil, err
	}
	return parseNVIDIANVLinkStatusByGPU(string(out)), nil
}

func parseNVIDIANVLinkStatusByGPU(raw string) map[int][]schema.HardwareNVLinkPort {
	result := map[int][]schema.HardwareNVLinkPort{}
	currentGPU := -1
	for _, line := range strings.Split(raw, "\n") {
		trimmed := strings.TrimSpace(line)
		if m := nvlinkGPUHeaderRe.FindStringSubmatch(trimmed); m != nil {
			currentGPU, _ = strconv.Atoi(m[1])
			continue
		}
		if currentGPU < 0 {
			continue
		}
		if m := nvlinkInactiveRe.FindStringSubmatch(trimmed); m != nil {
			idx, _ := strconv.Atoi(m[1])
			result[currentGPU] = append(result[currentGPU], schema.HardwareNVLinkPort{Index: idx, Active: false})
			continue
		}
		if m := nvlinkSpeedLineRe.FindStringSubmatch(trimmed); m != nil {
			idx, _ := strconv.Atoi(m[1])
			port := schema.HardwareNVLinkPort{Index: idx, Active: true}
			if speed, err := strconv.ParseFloat(m[2], 64); err == nil {
				port.SpeedGBs = &speed
			}
			result[currentGPU] = append(result[currentGPU], port)
		}
	}
	return result
}

// queryNVIDIANVLinkErrorsByGPU runs "nvidia-smi nvlink -e" and returns
// per-link error counters keyed by GPU index then link index.
func queryNVIDIANVLinkErrorsByGPU() (map[int]map[int]nvlinkErrorCounters, error) {
	out, err := exec.Command("nvidia-smi", "nvlink", "-e").Output()
	if err != nil {
		return nil, err
	}
	return parseNVIDIANVLinkErrorsByGPU(string(out)), nil
}

func parseNVIDIANVLinkErrorsByGPU(raw string) map[int]map[int]nvlinkErrorCounters {
	result := map[int]map[int]nvlinkErrorCounters{}
	currentGPU := -1
	for _, line := range strings.Split(raw, "\n") {
		trimmed := strings.TrimSpace(line)
		if m := nvlinkGPUHeaderRe.FindStringSubmatch(trimmed); m != nil {
			currentGPU, _ = strconv.Atoi(m[1])
			continue
		}
		if currentGPU < 0 {
			continue
		}
		m := nvlinkErrorCounterRe.FindStringSubmatch(trimmed)
		if m == nil {
			continue
		}
		linkIdx, _ := strconv.Atoi(m[1])
		count, _ := strconv.ParseInt(m[3], 10, 64)
		if result[currentGPU] == nil {
			result[currentGPU] = map[int]nvlinkErrorCounters{}
		}
		c := result[currentGPU][linkIdx]
		switch m[2] {
		case "Replay":
			c.Replay = count
		case "Recovery":
			c.Recovery = count
		case "CRC":
			c.CRC = count
		}
		result[currentGPU][linkIdx] = c
	}
	return result
}