bee/audit/internal/tui/update.go

package tui

import (
	"fmt"
	"strings"

	tea "github.com/charmbracelet/bubbletea"
)

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.KeyMsg:
		if m.busy {
			switch msg.String() {
			case "ctrl+c":
				return m, tea.Quit
			default:
				return m, nil
			}
		}
		return m.updateKey(msg)
	case resultMsg:
		m.busy = false
		m.busyTitle = ""
		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, nil
	case servicesMsg:
		m.busy = false
		m.busyTitle = ""
		if msg.err != nil {
			m.title = "Services"
			m.body = msg.err.Error()
			m.prevScreen = screenMain
			m.screen = screenOutput
			return m, nil
		}
		m.services = msg.services
		m.screen = screenServices
		m.cursor = 0
		return m, nil
	case interfacesMsg:
		m.busy = false
		m.busyTitle = ""
		if msg.err != nil {
			m.title = "interfaces"
			m.body = msg.err.Error()
			m.prevScreen = screenMain
			m.screen = screenOutput
			return m, nil
		}
		m.interfaces = msg.ifaces
		m.screen = screenInterfacePick
		m.cursor = 0
		return m, nil
	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, nil
		}
		m.targets = msg.targets
		m.screen = screenExportTargets
		m.cursor = 0
		return m, nil
	case bannerMsg:
		m.banner = strings.TrimSpace(msg.text)
		return m, nil
	case nvidiaGPUsMsg:
		return m.handleNvidiaGPUsMsg(msg)
	case nvtopClosedMsg:
		// nvtop closed — stay on running screen (or result if SAT is already done)
		return m, nil
	case nvidiaSATDoneMsg:
		if m.nvidiaSATAborted {
			return m, nil
		}
		if m.nvidiaSATCancel != nil {
			m.nvidiaSATCancel()
			m.nvidiaSATCancel = nil
		}
		m.prevScreen = screenAcceptance
		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, nil
	}

	return m, nil
}

func (m model) updateKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch m.screen {
	case screenMain:
		return m.updateMenu(msg, len(m.mainMenu), m.handleMainMenu)
	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 screenAcceptance:
		return m.updateMenu(msg, 4, m.handleAcceptanceMenu)
	case screenNvidiaSATSetup:
		return m.updateNvidiaSATSetup(msg)
	case screenNvidiaSATRunning:
		return m.updateNvidiaSATRunning(msg)
	case screenExportTargets:
		return m.updateMenu(msg, len(m.targets), m.handleExportTargetsMenu)
	case screenInterfacePick:
		return m.updateMenu(msg, len(m.interfaces), m.handleInterfacePickMenu)
	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
}

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, screenAcceptance:
			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 "q", "ctrl+c":
		return m, tea.Quit
	}
	return m, nil
}