package platform

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"

	"bee/audit/internal/collector"
)

// pcieLinkRetrainTimeout bounds how long we wait for a device to finish
// retraining (clear the Link Status "Link Training" bit) before giving up
// and reading whatever speed it settled on anyway. The PCIe spec allows up
// to 100ms for Gen1-3 and longer for higher generations with equalization;
// this is generous headroom above that.
const pcieLinkRetrainTimeout = 2 * time.Second

// pcieLinkFinding is one device's before/after link-speed retrain result.
type pcieLinkFinding struct {
	BDF          string
	Description  string
	VendorID     string
	ClassCode    string
	IsGPU        bool
	GPUVendor    string // "nvidia" or "amd", only set when IsGPU
	Skipped      string // non-empty reason this device wasn't retrained
	BeforeSpeed  string
	AfterSpeed   string
	MaxSpeed     string
	PortMaxSpeed string // bridge's own capability when MaxSpeed is limited by its downstream peer
	Width        int
	MaxWidth     int
	PortMaxWidth int // bridge's own capability when MaxWidth is limited by its downstream peer
	Degraded     bool
	NotPresent   bool // true when the slot trained to zero lanes: nothing is plugged in (or it fell off the bus), not a speed regression
}

// RunPCIeLinkCheckPack forces every enabled PCIe device to retrain its link
// (via the PCIe Link Control register's spec-defined Retrain Link bit — see
// PCIe base spec, Link Control Register, bit 5) and compares the
// post-retrain negotiated speed against the device's own reported maximum.
//
// This exists because a plain idle-time sysfs read of current_link_speed is
// not a reliable fault signal: NVIDIA GPUs (and other devices with runtime
// power management) deliberately downclock their PCIe link to save power
// while idle, which is indistinguishable from a real degraded slot/riser/
// cable without either sustained traffic or a forced retrain. Forcing a
// retrain sidesteps needing a device-specific load generator (bee-gpu-burn
// exists for GPUs; nothing plays that role for NICs, HBAs, or PCIe
// switches) — retraining is a PCIe-spec mechanism every endpoint supports,
// so this one check covers every PCIe device in the machine, not just
// GPUs. See bible-local/decisions/2026-08-24-pcie-gpu-gen1-idle-warning-unresolved.md
// for the history of narrower attempts that didn't generalize.
//
// Disabled devices (sysfs enable==0 — e.g. PCIe fabric-management endpoints
// the kernel never activates, per the 2026-06-12 decision) are left alone:
// they carry no data traffic, so there is nothing to verify and no reason
// to poke them.
func (s *System) RunPCIeLinkCheckPack(ctx context.Context, baseDir string, logFunc func(string)) (string, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if baseDir == "" {
		baseDir = "/var/log/bee-sat"
	}
	ts := time.Now().UTC().Format("20060102-150405")
	runDir := filepath.Join(baseDir, "pcie-link-"+ts)
	if err := os.MkdirAll(runDir, 0755); err != nil {
		return "", err
	}
	verboseLog := filepath.Join(runDir, "verbose.log")

	bdfs, err := listPCIDeviceBDFs()
	if err != nil {
		return "", fmt.Errorf("list PCI devices: %w", err)
	}

	var findings []pcieLinkFinding
	for _, bdf := range bdfs {
		if logFunc != nil {
			logFunc(fmt.Sprintf("=== %s ===", bdf))
		}
		f := retrainAndSamplePCIeDevice(ctx, verboseLog, bdf, logFunc)
		findings = append(findings, f)
	}

	summary := renderPCIeLinkCheckSummary(findings)
	if err := os.WriteFile(filepath.Join(runDir, "summary.txt"), []byte(summary), 0644); err != nil {
		return "", err
	}
	report := renderPCIeLinkCheckReport(findings)
	if err := os.WriteFile(filepath.Join(runDir, "pcie-link-report.txt"), []byte(report), 0644); err != nil {
		return "", err
	}
	return runDir, nil
}

// listPCIDeviceBDFs returns every BDF under /sys/bus/pci/devices, sorted for
// deterministic report ordering.
func listPCIDeviceBDFs() ([]string, error) {
	entries, err := os.ReadDir("/sys/bus/pci/devices")
	if err != nil {
		return nil, err
	}
	bdfs := make([]string, 0, len(entries))
	for _, e := range entries {
		bdfs = append(bdfs, e.Name())
	}
	sort.Strings(bdfs)
	return bdfs, nil
}

func retrainAndSamplePCIeDevice(ctx context.Context, verboseLog, bdf string, logFunc func(string)) pcieLinkFinding {
	f := pcieLinkFinding{BDF: bdf}

	vendor, _ := readPCIeSysfsHex(bdf, "vendor")
	class, _ := readPCIeSysfsHex(bdf, "class")
	f.VendorID = vendor
	f.ClassCode = class
	f.IsGPU, f.GPUVendor = classifyGPUFromVendorClass(vendor, class)
	f.Description = pcieDeviceDescription(ctx, verboseLog, bdf, logFunc)

	if enabled, ok := readPCIeSysfsInt(bdf, "enable"); ok && enabled == 0 {
		f.Skipped = "device disabled (no data traffic; link state has no operational impact)"
		return f
	}

	before, beforeOK := readPCIeSysfsString(bdf, "current_link_speed")
	maxSpeed, maxOK := readPCIeSysfsString(bdf, "max_link_speed")
	width, _ := readPCIeSysfsInt(bdf, "current_link_width")
	maxWidth, _ := readPCIeSysfsInt(bdf, "max_link_width")
	f.BeforeSpeed = before
	f.MaxSpeed = maxSpeed
	f.MaxWidth = maxWidth
	if !beforeOK || !maxOK {
		f.Skipped = "no PCIe link-speed attributes in sysfs (not a link-trained endpoint)"
		return f
	}

	if width == 0 {
		// A downstream switch/root port with nothing seated reads zero
		// trained lanes even before we touch it. Plenty of legitimate
		// configs leave slots like this unpopulated (not every server ships
		// every NIC/riser slot filled), so this is not by itself evidence
		// of anything wrong — retraining an empty slot can't produce a
		// meaningful speed reading, and there's no baseline here to say
		// "this used to have a card." Skip it exactly like a disabled
		// device: nothing to verify, no reason to fail the run over it.
		f.NotPresent = true
		f.Skipped = "no device present downstream (empty slot/riser — nothing to retrain)"
		return f
	}

	// A bridge/root port reports its own maximum capability in sysfs, not
	// the maximum mutually supported by the device at the other end of the
	// link.  Comparing a Gen4 x16 root port directly with a Gen3 x8 or Gen2
	// x4 endpoint therefore produces a false degradation even though the
	// link is running at the fastest rate the endpoint supports.  The child
	// device is tested separately, so use its advertised capability to
	// calculate the real target for this bridge-side view of the same link.
	if isPCIeBridgeClass(class) {
		if peerSpeed, peerWidth, ok := downstreamPCIeLinkCapability(bdf); ok {
			f.PortMaxSpeed = f.MaxSpeed
			f.PortMaxWidth = f.MaxWidth
			f.MaxSpeed = minPCIeLinkSpeed(f.MaxSpeed, peerSpeed)
			f.MaxWidth = minPositiveInt(f.MaxWidth, peerWidth)
		}
	}

	if err := retrainPCIeLink(ctx, verboseLog, bdf, logFunc); err != nil {
		f.Skipped = "retrain failed: " + err.Error()
		f.AfterSpeed = before
		f.Width = width
		// before/maxSpeed are already normalized "GenN" labels (see
		// readPCIeSysfsString) — compare directly, don't re-normalize.
		f.Degraded = before != maxSpeed
		return f
	}

	after, _ := readPCIeSysfsString(bdf, "current_link_speed")
	widthAfter, _ := readPCIeSysfsInt(bdf, "current_link_width")
	if widthAfter == 0 {
		// The device answered before the retrain but is gone immediately
		// after it (fell off the bus mid-check) — unlike the pre-retrain
		// case above, this had a live link a moment ago, so it's worth
		// surfacing rather than silently skipping.
		f.AfterSpeed = after
		f.Width = widthAfter
		f.NotPresent = true
		f.Degraded = true
		return f
	}
	f.AfterSpeed = after
	f.Width = widthAfter
	f.Degraded = after != maxSpeed
	return f
}

// retrainPCIeLink sets the Retrain Link bit (bit 5) of the PCI Express
// Capability's Link Control register via setpci, then polls the Link
// Status register's Link Training bit (bit 11) until it clears or
// pcieLinkRetrainTimeout elapses.
func retrainPCIeLink(ctx context.Context, verboseLog, bdf string, logFunc func(string)) error {
	linkCtrlOut, err := runSATCommandCtx(ctx, verboseLog, "setpci-read-"+bdf,
		[]string{"setpci", "-s", bdf, "CAP_EXP+0x10.w"}, nil, logFunc)
	if err != nil {
		return fmt.Errorf("read Link Control: %w", err)
	}
	cur, err := strconv.ParseUint(strings.TrimSpace(string(linkCtrlOut)), 16, 16)
	if err != nil {
		return fmt.Errorf("parse Link Control %q: %w", linkCtrlOut, err)
	}
	const retrainLinkBit = 0x0020
	newVal := uint16(cur) | retrainLinkBit

	if _, err := runSATCommandCtx(ctx, verboseLog, "setpci-retrain-"+bdf,
		[]string{"setpci", "-s", bdf, fmt.Sprintf("CAP_EXP+0x10.w=%04x", newVal)}, nil, logFunc); err != nil {
		return fmt.Errorf("write Retrain Link bit: %w", err)
	}

	const linkTrainingBit = 0x0800
	deadline := time.Now().Add(pcieLinkRetrainTimeout)
	for time.Now().Before(deadline) {
		statusOut, err := runSATCommandCtx(ctx, verboseLog, "setpci-status-"+bdf,
			[]string{"setpci", "-s", bdf, "CAP_EXP+0x12.w"}, nil, nil)
		if err == nil {
			if status, perr := strconv.ParseUint(strings.TrimSpace(string(statusOut)), 16, 16); perr == nil {
				if status&linkTrainingBit == 0 {
					return nil
				}
			}
		}
		time.Sleep(50 * time.Millisecond)
	}
	// Timed out waiting for training to clear; the caller still samples
	// whatever speed sysfs reports, which is the honest answer either way.
	return nil
}

func pcieDeviceDescription(ctx context.Context, verboseLog, bdf string, logFunc func(string)) string {
	out, err := runSATCommandCtx(ctx, verboseLog, "lspci-"+bdf, []string{"lspci", "-s", bdf}, nil, logFunc)
	if err != nil {
		return ""
	}
	line := strings.TrimSpace(string(out))
	if idx := strings.Index(line, "\n"); idx >= 0 {
		line = line[:idx]
	}
	if idx := strings.Index(line, " "); idx >= 0 {
		return strings.TrimSpace(line[idx+1:])
	}
	return line
}

func readPCIeSysfsString(bdf, attr string) (string, bool) {
	raw, err := os.ReadFile(filepath.Join("/sys/bus/pci/devices", bdf, attr))
	if err != nil {
		return "", false
	}
	v := strings.TrimSpace(string(raw))
	if v == "" {
		return "", false
	}
	return collector.NormalizePCILinkSpeed(v), true
}

func readPCIeSysfsInt(bdf, attr string) (int, bool) {
	raw, err := os.ReadFile(filepath.Join("/sys/bus/pci/devices", bdf, attr))
	if err != nil {
		return 0, false
	}
	v, err := strconv.Atoi(strings.TrimSpace(string(raw)))
	if err != nil || v < 0 {
		return 0, false
	}
	return v, true
}

func readPCIeSysfsHex(bdf, attr string) (string, bool) {
	raw, err := os.ReadFile(filepath.Join("/sys/bus/pci/devices", bdf, attr))
	if err != nil {
		return "", false
	}
	return strings.TrimSpace(string(raw)), true
}

// isPCIeBridgeClass matches PCI class 0x0604xx (PCI-to-PCI bridge).  These
// functions describe the upstream side of a downstream link, so their own
// max_link_* values must be capped by the peer's capability.
func isPCIeBridgeClass(classHex string) bool {
	c := strings.TrimPrefix(strings.ToLower(strings.TrimSpace(classHex)), "0x")
	return len(c) >= 4 && c[:4] == "0604"
}

// downstreamPCIeLinkCapability returns the strongest capability advertised
// by a bridge's immediate child functions.  In sysfs those functions are
// direct entries below the bridge device directory.  Multifunction devices
// expose several children for one physical link; taking the strongest values
// avoids understating the link because one auxiliary function omitted data.
func downstreamPCIeLinkCapability(bdf string) (speed string, width int, ok bool) {
	bridgeDir := filepath.Join("/sys/bus/pci/devices", bdf)
	return downstreamPCIeLinkCapabilityAt(bridgeDir)
}

func downstreamPCIeLinkCapabilityAt(bridgeDir string) (speed string, width int, ok bool) {
	entries, err := os.ReadDir(bridgeDir)
	if err != nil {
		return "", 0, false
	}
	for _, entry := range entries {
		if !isFullPCIBDF(entry.Name()) {
			continue
		}
		childDir := filepath.Join(bridgeDir, entry.Name())
		rawSpeed, speedErr := os.ReadFile(filepath.Join(childDir, "max_link_speed"))
		if speedErr != nil {
			continue
		}
		childSpeed := collector.NormalizePCILinkSpeed(strings.TrimSpace(string(rawSpeed)))
		if pcieGeneration(childSpeed) > pcieGeneration(speed) {
			speed = childSpeed
		}
		if rawWidth, widthErr := os.ReadFile(filepath.Join(childDir, "max_link_width")); widthErr == nil {
			if childWidth, parseErr := strconv.Atoi(strings.TrimSpace(string(rawWidth))); parseErr == nil && childWidth > width {
				width = childWidth
			}
		}
		ok = true
	}
	return speed, width, ok && speed != ""
}

func isFullPCIBDF(s string) bool {
	if len(s) != len("0000:00:00.0") || s[4] != ':' || s[7] != ':' || s[10] != '.' {
		return false
	}
	for i, r := range s {
		if i == 4 || i == 7 || i == 10 {
			continue
		}
		if !((r >= '0' && r <= '9') || (r >= 'a' && r <= 'f') || (r >= 'A' && r <= 'F')) {
			return false
		}
	}
	return true
}

func minPCIeLinkSpeed(a, b string) string {
	ga, gb := pcieGeneration(a), pcieGeneration(b)
	switch {
	case ga == 0:
		return b
	case gb == 0 || ga <= gb:
		return a
	default:
		return b
	}
}

func pcieGeneration(speed string) int {
	v := strings.TrimPrefix(strings.TrimSpace(speed), "Gen")
	gen, _ := strconv.Atoi(v)
	return gen
}

func minPositiveInt(a, b int) int {
	switch {
	case a <= 0:
		return b
	case b <= 0 || a <= b:
		return a
	default:
		return b
	}
}

// classifyGPUFromVendorClass reports whether a device is a GPU die itself
// (PCI base class 0x03 — Display Controller — under NVIDIA/AMD's vendor
// ID), as opposed to a same-vendor companion device (NIC, storage
// controller, NVLink bridge) that shares the GPU's PCI vendor ID. Class-code
// based, not name-substring based, per the same reasoning as
// collector.IsGPUClass.
func classifyGPUFromVendorClass(vendorHex, classHex string) (isGPU bool, vendor string) {
	v := strings.TrimPrefix(strings.ToLower(strings.TrimSpace(vendorHex)), "0x")
	c := strings.TrimPrefix(strings.ToLower(strings.TrimSpace(classHex)), "0x")
	if len(c) < 2 || c[:2] != "03" {
		return false, ""
	}
	switch v {
	case "10de":
		return true, "nvidia"
	case "1002":
		return true, "amd"
	default:
		return false, ""
	}
}

func renderPCIeLinkCheckSummary(findings []pcieLinkFinding) string {
	var b strings.Builder
	fmt.Fprintf(&b, "run_at_utc=%s\n", time.Now().UTC().Format(time.RFC3339))
	fmt.Fprintf(&b, "devices_tested=%d\n", len(findings))

	gpuStatus := map[string]string{} // vendor -> OK/FAILED
	otherDegraded := 0
	otherTested := 0
	anyDegraded := false
	for _, f := range findings {
		if f.Skipped != "" && f.AfterSpeed == "" {
			continue
		}
		if f.IsGPU {
			if _, ok := gpuStatus[f.GPUVendor]; !ok {
				gpuStatus[f.GPUVendor] = "OK"
			}
			if f.Degraded {
				gpuStatus[f.GPUVendor] = "FAILED"
				anyDegraded = true
			}
			continue
		}
		otherTested++
		if f.Degraded {
			otherDegraded++
			anyDegraded = true
		}
	}
	for _, vendor := range []string{"nvidia", "amd"} {
		if status, ok := gpuStatus[vendor]; ok {
			fmt.Fprintf(&b, "gpu_%s_status=%s\n", vendor, status)
		}
	}
	fmt.Fprintf(&b, "other_devices_tested=%d\n", otherTested)
	fmt.Fprintf(&b, "other_devices_degraded=%d\n", otherDegraded)
	if otherTested > 0 {
		if otherDegraded > 0 {
			fmt.Fprintln(&b, "other_status=FAILED")
		} else {
			fmt.Fprintln(&b, "other_status=OK")
		}
	}

	if anyDegraded {
		fmt.Fprintln(&b, "overall_status=FAILED")
		var reasons []string
		for _, f := range findings {
			if !f.Degraded {
				continue
			}
			if f.NotPresent {
				reasons = append(reasons, fmt.Sprintf("%s (%s): no device detected downstream (link down / empty slot or riser, capable of %s)",
					f.BDF, nonEmptyOr(f.Description, "unknown device"), f.MaxSpeed))
				continue
			}
			reasons = append(reasons, fmt.Sprintf("%s (%s): retrained to %s, capable of %s",
				f.BDF, nonEmptyOr(f.Description, "unknown device"), f.AfterSpeed, f.MaxSpeed))
		}
		fmt.Fprintf(&b, "warnings=%s\n", strings.Join(reasons, "; "))
	} else {
		fmt.Fprintln(&b, "overall_status=OK")
	}
	return b.String()
}

func renderPCIeLinkCheckReport(findings []pcieLinkFinding) string {
	var b strings.Builder
	line := strings.Repeat("=", 80)
	b.WriteString(line + "\n")
	b.WriteString("PCIe Link Retrain Check\n")
	b.WriteString(line + "\n\n")
	for _, f := range findings {
		fmt.Fprintf(&b, "%s  %s\n", f.BDF, nonEmptyOr(f.Description, "(unknown device)"))
		if f.Skipped != "" && f.AfterSpeed == "" {
			fmt.Fprintf(&b, "    skipped: %s\n", f.Skipped)
			continue
		}
		verdict := "OK"
		switch {
		case f.NotPresent:
			verdict = "FELL OFF BUS"
		case f.Degraded:
			verdict = "DEGRADED"
		}
		fmt.Fprintf(&b, "    %s: before=%s after=%s max=%s width=%d/%d",
			verdict, f.BeforeSpeed, f.AfterSpeed, f.MaxSpeed, f.Width, f.MaxWidth)
		if f.PortMaxSpeed != "" && (f.PortMaxSpeed != f.MaxSpeed || f.PortMaxWidth != f.MaxWidth) {
			fmt.Fprintf(&b, " (port capability %s x%d, limited by downstream device)", f.PortMaxSpeed, f.PortMaxWidth)
		}
		fmt.Fprintln(&b)
		if f.Skipped != "" {
			fmt.Fprintf(&b, "    note: %s\n", f.Skipped)
		}
	}
	return b.String()
}