bee/audit/internal/tui/update.go

package tui

import (
	"fmt"
	"strings"
	"time"

	tea "github.com/charmbracelet/bubbletea"
)

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		m.width = msg.Width
		return m, nil
	case tea.KeyMsg:
		if m.busy {
			if msg.String() == "ctrl+c" {
				return m, tea.Quit
			}
			return m, nil
		}
		next, cmd := m.updateKey(msg)
		nextModel := next.(model)
		if shouldRefreshSnapshot(m, nextModel) {
			return nextModel, tea.Batch(cmd, nextModel.refreshSnapshotCmd())
		}
		return nextModel, cmd
	case satProgressMsg:
		if m.busy && m.progressPrefix != "" {
			if len(msg.lines) > 0 {
				m.progressLines = msg.lines
			}
			return m, pollSATProgress(m.progressPrefix, m.progressSince)
		}
		if m.busy && m.installCancel != nil {
			if len(msg.lines) > 0 {
				m.progressLines = msg.lines
			}
			return m, pollInstallProgress(DefaultInstallLogFile)
		}
		return m, nil
	case snapshotMsg:
		m.banner = msg.banner
		m.panel = msg.panel
		return m, nil
	case resultMsg:
		m.busy = false
		m.busyTitle = ""
		m.progressLines = nil
		m.progressPrefix = ""
		m.title = msg.title
		if msg.err != nil {
			body := strings.TrimSpace(msg.body)
			if body == "" {
				m.body = fmt.Sprintf("ERROR: %v", msg.err)
			} else {
				m.body = fmt.Sprintf("%s\n\nERROR: %v", body, msg.err)
			}
		} else {
			m.body = msg.body
		}
		m.pendingAction = actionNone
		if msg.back != "" {
			m.prevScreen = msg.back
		} else {
			m.prevScreen = m.screen
		}
		m.screen = screenOutput
		m.cursor = 0
		return m, m.refreshSnapshotCmd()
	case servicesMsg:
		m.busy = false
		m.busyTitle = ""
		if msg.err != nil {
			m.title = "Services"
			m.body = msg.err.Error()
			m.prevScreen = screenSettings
			m.screen = screenOutput
			return m, m.refreshSnapshotCmd()
		}
		m.services = msg.services
		m.screen = screenServices
		m.cursor = 0
		return m, m.refreshSnapshotCmd()
	case interfacesMsg:
		m.busy = false
		m.busyTitle = ""
		if msg.err != nil {
			m.title = "interfaces"
			m.body = msg.err.Error()
			m.prevScreen = screenNetwork
			m.screen = screenOutput
			return m, m.refreshSnapshotCmd()
		}
		m.interfaces = msg.ifaces
		m.screen = screenInterfacePick
		m.cursor = 0
		return m, m.refreshSnapshotCmd()
	case exportTargetsMsg:
		m.busy = false
		m.busyTitle = ""
		if msg.err != nil {
			m.title = "export"
			m.body = msg.err.Error()
			m.prevScreen = screenMain
			m.screen = screenOutput
			return m, m.refreshSnapshotCmd()
		}
		if len(msg.targets) == 0 {
			m.title = "Export support bundle"
			m.body = "No writable removable filesystems found.\n\nRead-only or boot media are hidden from this list."
			m.prevScreen = screenMain
			m.screen = screenOutput
			return m, m.refreshSnapshotCmd()
		}
		m.targets = msg.targets
		m.screen = screenExportTargets
		m.cursor = 0
		return m, m.refreshSnapshotCmd()
	case installDisksMsg:
		m.busy = false
		m.busyTitle = ""
		if msg.err != nil {
			m.title = "Install to disk"
			m.body = msg.err.Error()
			m.prevScreen = screenTools
			m.screen = screenOutput
			return m, m.refreshSnapshotCmd()
		}
		if len(msg.disks) == 0 {
			m.title = "Install to disk"
			m.body = "No suitable disks found.\n\nOnly non-USB, non-boot disks are shown.\nAttach a target disk and try again."
			m.prevScreen = screenTools
			m.screen = screenOutput
			return m, m.refreshSnapshotCmd()
		}
		m.installDisks = msg.disks
		m.screen = screenInstallDiskPick
		m.cursor = 0
		return m, m.refreshSnapshotCmd()
	case installDoneMsg:
		if m.installCancel != nil {
			m.installCancel()
			m.installCancel = nil
		}
		m.busy = false
		m.busyTitle = ""
		m.progressLines = nil
		m.prevScreen = screenTools
		m.screen = screenOutput
		m.title = "Install to disk"
		if msg.err != nil {
			m.body = fmt.Sprintf("Installation FAILED.\n\nLog: %s\n\nERROR: %v", DefaultInstallLogFile, msg.err)
		} else {
			m.body = fmt.Sprintf("Installation complete.\n\nRemove the ISO and reboot to start the installed system.\n\nLog: %s", DefaultInstallLogFile)
		}
		return m, m.refreshSnapshotCmd()
	case nvtopClosedMsg:
		return m, nil
	case gpuStressDoneMsg:
		if m.gpuStressAborted {
			return m, nil
		}
		if m.gpuStressCancel != nil {
			m.gpuStressCancel()
			m.gpuStressCancel = nil
		}
		m.prevScreen = screenBurnInTests
		m.screen = screenOutput
		m.title = msg.title
		if msg.err != nil {
			body := strings.TrimSpace(msg.body)
			if body == "" {
				m.body = fmt.Sprintf("ERROR: %v", msg.err)
			} else {
				m.body = fmt.Sprintf("%s\n\nERROR: %v", body, msg.err)
			}
		} else {
			m.body = msg.body
		}
		return m, m.refreshSnapshotCmd()
	case gpuLiveTickMsg:
		if m.screen == screenGPUStressRunning {
			if len(msg.rows) > 0 {
				elapsed := time.Since(m.gpuLiveStart).Seconds()
				for i := range msg.rows {
					msg.rows[i].ElapsedSec = elapsed
				}
				m.gpuLiveRows = append(m.gpuLiveRows, msg.rows...)
				n := max(1, len(msg.indices))
				if len(m.gpuLiveRows) > 60*n {
					m.gpuLiveRows = m.gpuLiveRows[len(m.gpuLiveRows)-60*n:]
				}
			}
			return m, pollGPULive(msg.indices)
		}
		return m, nil
	case nvidiaSATDoneMsg:
		if m.nvidiaSATAborted {
			return m, nil
		}
		if m.nvidiaSATCancel != nil {
			m.nvidiaSATCancel()
			m.nvidiaSATCancel = nil
		}
		m.prevScreen = screenHealthCheck
		m.screen = screenOutput
		m.title = msg.title
		if msg.err != nil {
			body := strings.TrimSpace(msg.body)
			if body == "" {
				m.body = fmt.Sprintf("ERROR: %v", msg.err)
			} else {
				m.body = fmt.Sprintf("%s\n\nERROR: %v", body, msg.err)
			}
		} else {
			m.body = msg.body
		}
		return m, m.refreshSnapshotCmd()
	}
	return m, nil
}

func (m model) updateKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch m.screen {
	case screenMain:
		return m.updateMain(msg)
	case screenHealthCheck:
		return m.updateHealthCheck(msg)
	case screenBurnInTests:
		return m.updateBurnInTests(msg)
	case screenSettings:
		return m.updateMenu(msg, len(m.settingsMenu), m.handleSettingsMenu)
	case screenNetwork:
		return m.updateMenu(msg, len(m.networkMenu), m.handleNetworkMenu)
	case screenServices:
		return m.updateMenu(msg, len(m.services), m.handleServicesMenu)
	case screenServiceAction:
		return m.updateMenu(msg, len(m.serviceMenu), m.handleServiceActionMenu)
	case screenNvidiaSATSetup:
		return m.updateNvidiaSATSetup(msg)
	case screenNvidiaSATRunning:
		return m.updateNvidiaSATRunning(msg)
	case screenGPUStressRunning:
		return m.updateGPUStressRunning(msg)
	case screenExportTargets:
		return m.updateMenu(msg, len(m.targets), m.handleExportTargetsMenu)
	case screenInterfacePick:
		return m.updateMenu(msg, len(m.interfaces), m.handleInterfacePickMenu)
	case screenTools:
		return m.updateMenu(msg, len(m.toolsMenu), m.handleToolsMenu)
	case screenInstallDiskPick:
		return m.updateMenu(msg, len(m.installDisks), m.handleInstallDiskPickMenu)
	case screenOutput:
		switch msg.String() {
		case "esc", "enter", "q":
			m.screen = m.prevScreen
			m.body = ""
			m.title = ""
			m.pendingAction = actionNone
			return m, nil
		case "ctrl+c":
			return m, tea.Quit
		}
	case screenStaticForm:
		return m.updateStaticForm(msg)
	case screenConfirm:
		return m.updateConfirm(msg)
	}
	if msg.String() == "ctrl+c" {
		return m, tea.Quit
	}
	return m, nil
}

// updateMain handles keys on the main (two-column) screen.
func (m model) updateMain(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	if m.panelFocus {
		return m.updateMainPanel(msg)
	}
	// Switch focus to right panel.
	if (msg.String() == "tab" || msg.String() == "right" || msg.String() == "l") && len(m.panel.Rows) > 0 {
		m.panelFocus = true
		return m, nil
	}
	return m.updateMenu(msg, len(m.mainMenu), m.handleMainMenu)
}

// updateMainPanel handles keys when right panel has focus.
func (m model) updateMainPanel(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch msg.String() {
	case "up", "k":
		if m.panelCursor > 0 {
			m.panelCursor--
		}
	case "down", "j":
		if m.panelCursor < len(m.panel.Rows)-1 {
			m.panelCursor++
		}
	case "enter":
		if m.panelCursor < len(m.panel.Rows) {
			key := m.panel.Rows[m.panelCursor].Key
			m.busy = true
			m.busyTitle = key
			return m, func() tea.Msg {
				r := m.app.ComponentDetailResult(key)
				return resultMsg{title: r.Title, body: r.Body, back: screenMain}
			}
		}
	case "tab", "left", "h", "esc":
		m.panelFocus = false
	case "q", "ctrl+c":
		return m, tea.Quit
	}
	return m, nil
}

func (m model) updateMenu(msg tea.KeyMsg, size int, onEnter func() (tea.Model, tea.Cmd)) (tea.Model, tea.Cmd) {
	if size == 0 {
		size = 1
	}
	switch msg.String() {
	case "up", "k":
		if m.cursor > 0 {
			m.cursor--
		}
	case "down", "j":
		if m.cursor < size-1 {
			m.cursor++
		}
	case "enter":
		return onEnter()
	case "esc":
		switch m.screen {
		case screenNetwork, screenServices:
			m.screen = screenSettings
			m.cursor = 0
		case screenSettings:
			m.screen = screenMain
			m.cursor = 0
		case screenServiceAction:
			m.screen = screenServices
			m.cursor = 0
		case screenExportTargets:
			m.screen = screenMain
			m.cursor = 0
		case screenInterfacePick:
			m.screen = screenNetwork
			m.cursor = 0
		case screenTools:
			m.screen = screenSettings
			m.cursor = 0
		case screenInstallDiskPick:
			m.screen = screenTools
			m.cursor = 0
		}
	case "q", "ctrl+c":
		return m, tea.Quit
	}
	return m, nil
}