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
	Width       int
	MaxWidth    int
	Degraded    bool
}

// 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 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")
	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
}

// 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 {
				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"
		if f.Degraded {
			verdict = "DEGRADED"
		}
		fmt.Fprintf(&b, "    %s: before=%s after=%s max=%s width=%d/%d\n",
			verdict, f.BeforeSpeed, f.AfterSpeed, f.MaxSpeed, f.Width, f.MaxWidth)
		if f.Skipped != "" {
			fmt.Fprintf(&b, "    note: %s\n", f.Skipped)
		}
	}
	return b.String()
}