bee/audit/internal/tui/screen_nvidia_sat.go

package tui

import (
	"context"
	"fmt"
	"os/exec"
	"strings"

	"bee/audit/internal/platform"

	tea "github.com/charmbracelet/bubbletea"
)

var nvidiaDurationOptions = []struct {
	label   string
	seconds int
}{
	{"10 minutes", 600},
	{"1 hour", 3600},
	{"8 hours", 28800},
	{"24 hours", 86400},
}

// enterNvidiaSATSetup resets the setup screen and starts loading GPU list.
func (m model) enterNvidiaSATSetup() (tea.Model, tea.Cmd) {
	m.screen = screenNvidiaSATSetup
	m.nvidiaGPUs = nil
	m.nvidiaGPUSel = nil
	m.nvidiaDurIdx = 0
	m.nvidiaSATCursor = 0
	m.busy = true
	m.busyTitle = "NVIDIA SAT"
	return m, func() tea.Msg {
		gpus, err := m.app.ListNvidiaGPUs()
		return nvidiaGPUsMsg{gpus: gpus, err: err}
	}
}

// handleNvidiaGPUsMsg processes the GPU list response.
func (m model) handleNvidiaGPUsMsg(msg nvidiaGPUsMsg) (tea.Model, tea.Cmd) {
	m.busy = false
	m.busyTitle = ""
	if msg.err != nil {
		m.title = "NVIDIA SAT"
		m.body = fmt.Sprintf("Failed to list GPUs: %v", msg.err)
		m.prevScreen = screenAcceptance
		m.screen = screenOutput
		return m, nil
	}
	m.nvidiaGPUs = msg.gpus
	m.nvidiaGPUSel = make([]bool, len(msg.gpus))
	for i := range m.nvidiaGPUSel {
		m.nvidiaGPUSel[i] = true // all selected by default
	}
	m.nvidiaSATCursor = 0
	return m, nil
}

// updateNvidiaSATSetup handles keys on the setup screen.
func (m model) updateNvidiaSATSetup(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	numDur := len(nvidiaDurationOptions)
	numGPU := len(m.nvidiaGPUs)
	totalItems := numDur + numGPU + 2 // +2: Start, Cancel
	switch msg.String() {
	case "up", "k":
		if m.nvidiaSATCursor > 0 {
			m.nvidiaSATCursor--
		}
	case "down", "j":
		if m.nvidiaSATCursor < totalItems-1 {
			m.nvidiaSATCursor++
		}
	case " ":
		switch {
		case m.nvidiaSATCursor < numDur:
			m.nvidiaDurIdx = m.nvidiaSATCursor
		case m.nvidiaSATCursor < numDur+numGPU:
			i := m.nvidiaSATCursor - numDur
			m.nvidiaGPUSel[i] = !m.nvidiaGPUSel[i]
		}
	case "enter":
		startIdx := numDur + numGPU
		cancelIdx := startIdx + 1
		switch {
		case m.nvidiaSATCursor < numDur:
			m.nvidiaDurIdx = m.nvidiaSATCursor
		case m.nvidiaSATCursor < startIdx:
			i := m.nvidiaSATCursor - numDur
			m.nvidiaGPUSel[i] = !m.nvidiaGPUSel[i]
		case m.nvidiaSATCursor == startIdx:
			return m.startNvidiaSAT()
		case m.nvidiaSATCursor == cancelIdx:
			m.screen = screenAcceptance
			m.cursor = 0
		}
	case "esc":
		m.screen = screenAcceptance
		m.cursor = 0
	case "ctrl+c", "q":
		return m, tea.Quit
	}
	return m, nil
}

// startNvidiaSAT launches the SAT and nvtop.
func (m model) startNvidiaSAT() (tea.Model, tea.Cmd) {
	var selectedGPUs []platform.NvidiaGPU
	for i, sel := range m.nvidiaGPUSel {
		if sel {
			selectedGPUs = append(selectedGPUs, m.nvidiaGPUs[i])
		}
	}
	if len(selectedGPUs) == 0 {
		selectedGPUs = m.nvidiaGPUs // fallback: use all if none explicitly selected
	}

	sizeMB := 0
	for _, g := range selectedGPUs {
		if sizeMB == 0 || g.MemoryMB < sizeMB {
			sizeMB = g.MemoryMB
		}
	}
	if sizeMB == 0 {
		sizeMB = 64
	}

	var gpuIndices []int
	for _, g := range selectedGPUs {
		gpuIndices = append(gpuIndices, g.Index)
	}

	durationSec := nvidiaDurationOptions[m.nvidiaDurIdx].seconds

	ctx, cancel := context.WithCancel(context.Background())
	m.nvidiaSATCancel = cancel
	m.nvidiaSATAborted = false
	m.screen = screenNvidiaSATRunning
	m.nvidiaSATCursor = 0

	satCmd := func() tea.Msg {
		result, err := m.app.RunNvidiaAcceptancePackWithOptions(ctx, "", durationSec, sizeMB, gpuIndices)
		return nvidiaSATDoneMsg{title: result.Title, body: result.Body, err: err}
	}

	nvtopPath, lookErr := exec.LookPath("nvtop")
	if lookErr != nil {
		// nvtop not available: just run the SAT, show running screen
		return m, satCmd
	}

	return m, tea.Batch(
		satCmd,
		tea.ExecProcess(exec.Command(nvtopPath), func(_ error) tea.Msg {
			return nvtopClosedMsg{}
		}),
	)
}

// updateNvidiaSATRunning handles keys on the running screen.
func (m model) updateNvidiaSATRunning(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch msg.String() {
	case "o", "O":
		nvtopPath, err := exec.LookPath("nvtop")
		if err != nil {
			return m, nil
		}
		return m, tea.ExecProcess(exec.Command(nvtopPath), func(_ error) tea.Msg {
			return nvtopClosedMsg{}
		})
	case "a", "A":
		if m.nvidiaSATCancel != nil {
			m.nvidiaSATCancel()
			m.nvidiaSATCancel = nil
		}
		m.nvidiaSATAborted = true
		m.screen = screenAcceptance
		m.cursor = 0
	case "ctrl+c":
		return m, tea.Quit
	}
	return m, nil
}

// renderNvidiaSATSetup renders the setup screen.
func renderNvidiaSATSetup(m model) string {
	var b strings.Builder
	fmt.Fprintln(&b, "NVIDIA SAT")
	fmt.Fprintln(&b)
	fmt.Fprintln(&b, "Duration:")
	for i, opt := range nvidiaDurationOptions {
		radio := "( )"
		if i == m.nvidiaDurIdx {
			radio = "(*)"
		}
		prefix := "  "
		if m.nvidiaSATCursor == i {
			prefix = "> "
		}
		fmt.Fprintf(&b, "%s%s %s\n", prefix, radio, opt.label)
	}
	fmt.Fprintln(&b)
	if len(m.nvidiaGPUs) == 0 {
		fmt.Fprintln(&b, "GPUs: (none detected)")
	} else {
		fmt.Fprintln(&b, "GPUs:")
		for i, gpu := range m.nvidiaGPUs {
			check := "[ ]"
			if m.nvidiaGPUSel[i] {
				check = "[x]"
			}
			prefix := "  "
			if m.nvidiaSATCursor == len(nvidiaDurationOptions)+i {
				prefix = "> "
			}
			fmt.Fprintf(&b, "%s%s %d: %s (%d MB)\n", prefix, check, gpu.Index, gpu.Name, gpu.MemoryMB)
		}
	}
	fmt.Fprintln(&b)
	startIdx := len(nvidiaDurationOptions) + len(m.nvidiaGPUs)
	startPfx := "  "
	cancelPfx := "  "
	if m.nvidiaSATCursor == startIdx {
		startPfx = "> "
	}
	if m.nvidiaSATCursor == startIdx+1 {
		cancelPfx = "> "
	}
	fmt.Fprintf(&b, "%sStart\n", startPfx)
	fmt.Fprintf(&b, "%sCancel\n", cancelPfx)
	fmt.Fprintln(&b)
	b.WriteString("[↑/↓] move  [space] toggle  [enter] select  [esc] cancel\n")
	return b.String()
}

// renderNvidiaSATRunning renders the running screen.
func renderNvidiaSATRunning() string {
	return "NVIDIA SAT\n\nTest is running...\n\n[o] Open nvtop  [a] Abort test  [ctrl+c] quit\n"
}