package platform

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

// nvidiaPCIeBandwidthFinding is one GPU's post-load PCIe link-speed result.
type nvidiaPCIeBandwidthFinding struct {
	Index       int
	BDF         string
	BeforeSpeed string
	AfterSpeed  string
	MaxSpeed    string
	Width       int
	MaxWidth    int
	Degraded    bool
}

// RunNvidiaPCIeBandwidthPack drives real host<->device traffic across the
// GPUs' PCIe links (via `dcgmi diag -r nvbandwidth`, the same tool the
// existing nvidia-bandwidth SAT uses for GPU-to-GPU throughput) and then
// resamples each GPU's PCIe link speed from sysfs immediately afterward.
//
// This is deliberately a separate, narrower check from "nvidia-bandwidth":
// that SAT's overall_status reflects nvbandwidth's own pass/fail (did the
// measured GB/s clear its internal threshold), never PCIe link speed
// itself. This pack exists purely to answer "did sustained real traffic
// make the link train up to its negotiated maximum" — the load-bearing
// counterpart to the idle-time collector reading that used to (falsely)
// drive pcie:gpu:nvidia's status. See
// bible-local/decisions/2026-08-24-pcie-gpu-gen1-idle-warning-unresolved.md.
func (s *System) RunNvidiaPCIeBandwidthPack(ctx context.Context, baseDir string, gpuIndices []int, 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, "nvidia-pcie-bandwidth-"+ts)
	if err := os.MkdirAll(runDir, 0755); err != nil {
		return "", err
	}
	verboseLog := filepath.Join(runDir, "verbose.log")

	bdfByIndex, err := gpuIndexToBDF(gpuIndices)
	if err != nil {
		return "", fmt.Errorf("resolve GPU BDFs: %w", err)
	}
	indices := make([]int, 0, len(bdfByIndex))
	for idx := range bdfByIndex {
		indices = append(indices, idx)
	}
	sort.Ints(indices)

	findings := make([]nvidiaPCIeBandwidthFinding, 0, len(indices))
	for _, idx := range indices {
		bdf := bdfByIndex[idx]
		before, _ := readPCIeSysfsString(bdf, "current_link_speed")
		max, _ := readPCIeSysfsString(bdf, "max_link_speed")
		maxWidth, _ := readPCIeSysfsInt(bdf, "max_link_width")
		findings = append(findings, nvidiaPCIeBandwidthFinding{
			Index: idx, BDF: bdf, BeforeSpeed: before, MaxSpeed: max, MaxWidth: maxWidth,
		})
	}

	cmd := []string{"dcgmi", "diag", "-r", "nvbandwidth"}
	if len(indices) > 0 {
		cmd = append(cmd, "-i", joinIndexList(indices))
	}
	out, runErr := runSATCommandCtx(ctx, verboseLog, "dcgmi-nvbandwidth", cmd, nil, logFunc)
	_ = os.WriteFile(filepath.Join(runDir, "01-dcgmi-nvbandwidth.log"), out, 0644)

	for i := range findings {
		bdf := findings[i].BDF
		after, _ := readPCIeSysfsString(bdf, "current_link_speed")
		width, _ := readPCIeSysfsInt(bdf, "current_link_width")
		findings[i].AfterSpeed = after
		findings[i].Width = width
		findings[i].Degraded = after != findings[i].MaxSpeed
	}

	summary := renderNvidiaPCIeBandwidthSummary(findings, runErr)
	if err := os.WriteFile(filepath.Join(runDir, "summary.txt"), []byte(summary), 0644); err != nil {
		return "", err
	}
	report := renderNvidiaPCIeBandwidthReport(findings, runErr)
	if err := os.WriteFile(filepath.Join(runDir, "nvidia-pcie-bandwidth-report.txt"), []byte(report), 0644); err != nil {
		return "", err
	}
	return runDir, nil
}

// gpuIndexToBDF resolves each of gpuIndices to its PCI BDF via nvidia-smi.
// An empty gpuIndices resolves every GPU nvidia-smi reports.
func gpuIndexToBDF(gpuIndices []int) (map[int]string, error) {
	out, err := satExecCommand("nvidia-smi", "--query-gpu=index,pci.bus_id", "--format=csv,noheader,nounits").Output()
	if err != nil {
		return nil, fmt.Errorf("nvidia-smi: %w", err)
	}
	want := make(map[int]struct{}, len(gpuIndices))
	for _, idx := range gpuIndices {
		want[idx] = struct{}{}
	}
	filterByIndex := len(want) > 0

	result := make(map[int]string)
	for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
		fields := strings.SplitN(line, ",", 2)
		if len(fields) != 2 {
			continue
		}
		idx, err := strconv.Atoi(strings.TrimSpace(fields[0]))
		if err != nil {
			continue
		}
		if filterByIndex {
			if _, ok := want[idx]; !ok {
				continue
			}
		}
		result[idx] = normalizeNvidiaBDF(strings.TrimSpace(fields[1]))
	}
	return result, nil
}

func renderNvidiaPCIeBandwidthSummary(findings []nvidiaPCIeBandwidthFinding, runErr error) string {
	var b strings.Builder
	fmt.Fprintf(&b, "run_at_utc=%s\n", time.Now().UTC().Format(time.RFC3339))
	fmt.Fprintf(&b, "gpu_count=%d\n", len(findings))
	degraded := 0
	for _, f := range findings {
		fmt.Fprintf(&b, "gpu%d_status=%s\n", f.Index, statusLabel(!f.Degraded))
		if f.Degraded {
			degraded++
		}
	}
	fmt.Fprintf(&b, "degraded=%d\n", degraded)
	if runErr != nil {
		fmt.Fprintf(&b, "dcgmi_nvbandwidth_error=%s\n", runErr.Error())
	}
	if degraded > 0 {
		fmt.Fprintln(&b, "overall_status=FAILED")
		var reasons []string
		for _, f := range findings {
			if f.Degraded {
				reasons = append(reasons, fmt.Sprintf("GPU%d (%s): still at %s under load, capable of %s",
					f.Index, f.BDF, f.AfterSpeed, f.MaxSpeed))
			}
		}
		fmt.Fprintf(&b, "warnings=%s\n", strings.Join(reasons, "; "))
	} else {
		fmt.Fprintln(&b, "overall_status=OK")
	}
	return b.String()
}

func statusLabel(ok bool) string {
	if ok {
		return "OK"
	}
	return "FAILED"
}

func renderNvidiaPCIeBandwidthReport(findings []nvidiaPCIeBandwidthFinding, runErr error) string {
	var b strings.Builder
	line := strings.Repeat("=", 80)
	b.WriteString(line + "\n")
	b.WriteString("NVIDIA GPU PCIe Bandwidth / Link-Under-Load Check\n")
	b.WriteString(line + "\n\n")
	if runErr != nil {
		fmt.Fprintf(&b, "dcgmi diag -r nvbandwidth: %s\n\n", runErr.Error())
	}
	for _, f := range findings {
		verdict := "OK"
		if f.Degraded {
			verdict = "DEGRADED"
		}
		fmt.Fprintf(&b, "GPU%d (%s): %s\n", f.Index, f.BDF, verdict)
		fmt.Fprintf(&b, "    before=%s after=%s max=%s width=%d/%d\n",
			f.BeforeSpeed, f.AfterSpeed, f.MaxSpeed, f.Width, f.MaxWidth)
	}
	return b.String()
}