pxmon/internal/cli/monitor_ssh.go

package cli

import (
	"context"
	"fmt"
	"strings"
	"time"

	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/hinshun/vt10x"

	"pxmon/internal/cluster"
)

// sshStartedMsg is dispatched when an embedded SSH session finishes dialing.
type sshStartedMsg struct {
	session *cluster.InteractiveSession
	cluster cluster.Cluster
	cols    int
	rows    int
	err     error
}

// sshChunkMsg delivers a chunk of remote PTY output to the model.
type sshChunkMsg struct {
	data []byte
	err  error
}

// sshClosedMsg signals that the embedded session was torn down.
type sshClosedMsg struct {
	err error
}

// sshTickMsg throttles vt10x → view repaints while output is flowing.
type sshTickMsg struct{}

const (
	sshMinCols         = 20
	sshMinRows         = 5
	sshChunkBuffer     = 16384
	sshRepaintInterval = 33 * time.Millisecond
)

// sshEmbeddedDims returns the usable grid size for the SSH panel.
func (m monitorModel) sshEmbeddedDims() (int, int) {
	w := m.width
	if w <= 0 {
		w = 120
	}
	h := m.height
	if h <= 0 {
		h = 32
	}
	cols := w - 4
	rows := h - 4
	if cols < sshMinCols {
		cols = sshMinCols
	}
	if rows < sshMinRows {
		rows = sshMinRows
	}
	return cols, rows
}

// startEmbeddedSSHCmd kicks off an SSH dial in a goroutine and returns
// the started session through an sshStartedMsg.
func (m monitorModel) startEmbeddedSSHCmd(selector string) tea.Cmd {
	svc := m.svc
	cols, rows := m.sshEmbeddedDims()
	c, err := svc.Get(selector)
	if err != nil {
		return func() tea.Msg { return sshStartedMsg{err: err} }
	}
	return func() tea.Msg {
		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
		defer cancel()
		sess, err := svc.StartInteractiveShell(ctx, c.ID, cols, rows, "xterm-256color")
		if err != nil {
			return sshStartedMsg{err: err, cluster: c}
		}
		return sshStartedMsg{
			session: sess,
			cluster: c,
			cols:    cols,
			rows:    rows,
		}
	}
}

// readSSHChunkCmd reads the next chunk from the SSH session in a goroutine.
func readSSHChunkCmd(sess *cluster.InteractiveSession) tea.Cmd {
	if sess == nil {
		return nil
	}
	return func() tea.Msg {
		buf := make([]byte, sshChunkBuffer)
		n, err := sess.Read(buf)
		if n > 0 {
			data := make([]byte, n)
			copy(data, buf[:n])
			return sshChunkMsg{data: data, err: err}
		}
		return sshChunkMsg{err: err}
	}
}

// sshRepaintTickCmd schedules a throttled repaint while output flows.
func sshRepaintTickCmd() tea.Cmd {
	return tea.Tick(sshRepaintInterval, func(_ time.Time) tea.Msg {
		return sshTickMsg{}
	})
}

// enterEmbeddedSSH switches the model into embedded SSH view for the
// freshly dialed session.
func (m *monitorModel) enterEmbeddedSSH(msg sshStartedMsg) {
	cols := msg.cols
	rows := msg.rows
	if cols < sshMinCols {
		cols = sshMinCols
	}
	if rows < sshMinRows {
		rows = sshMinRows
	}
	vt := vt10x.New(vt10x.WithSize(cols, rows))
	m.sshMode = true
	m.sshSess = msg.session
	m.sshCluster = msg.cluster
	m.sshVT = vt
	m.sshCols = cols
	m.sshRows = rows
	m.sshClosed = false
	m.sshErr = ""
	m.sshPendingRepaint = false
}

// exitEmbeddedSSH tears down the embedded session and clears state.
func (m *monitorModel) exitEmbeddedSSH(reason string) {
	if m.sshSess != nil {
		_ = m.sshSess.Close()
	}
	m.sshSess = nil
	m.sshVT = nil
	m.sshMode = false
	m.sshClosed = true
	m.sshPendingRepaint = false
	// Return to pxmon console so the user lands where they launched from.
	m.termMode = true
	m.termFull = false
	if reason != "" {
		m.setStatus(reason)
	}
}

// writeSSHInput pushes a byte slice to the SSH session's stdin.
func (m monitorModel) writeSSHInput(p []byte) {
	if m.sshSess == nil || len(p) == 0 {
		return
	}
	_, _ = m.sshSess.Write(p)
}

// handleSSHKey routes TUI key events into the remote PTY stdin.
func (m monitorModel) handleSSHKey(v tea.KeyMsg) (tea.Model, tea.Cmd) {
	// Escape hatch: Ctrl+] closes the embedded session.
	if v.Type == tea.KeyCtrlCloseBracket {
		m.exitEmbeddedSSH(fmt.Sprintf("ssh closed (%s)", m.sshCluster.Name))
		m.appendConsoleOutput(fmt.Sprintf("ssh session closed (%s)", m.sshCluster.Name))
		return m, nil
	}
	// Scrollback controls — mirror the pxmon console (alt+↑/↓, pgup/pgdn,
	// alt+home/end). Alt+Shift+↑/↓ jump 10 lines at a time for fast review.
	// These never reach the remote PTY.
	switch v.String() {
	case "alt+up":
		m.scrollSSH(1)
		return m, nil
	case "alt+down":
		m.scrollSSH(-1)
		return m, nil
	case "alt+shift+up":
		m.scrollSSH(10)
		return m, nil
	case "alt+shift+down":
		m.scrollSSH(-10)
		return m, nil
	case "pgup":
		m.scrollSSH(m.sshRows - 1)
		return m, nil
	case "pgdown":
		m.scrollSSH(-(m.sshRows - 1))
		return m, nil
	case "alt+home":
		m.sshScrollOffset = len(m.sshScrollback)
		m.clampSSHScroll()
		return m, nil
	case "alt+end":
		m.sshScrollOffset = 0
		return m, nil
	case "alt+p":
		m.privacyMode = !m.privacyMode
		if m.privacyMode {
			m.setStatus("privacy: on")
		} else {
			m.setStatus("privacy: off")
		}
		return m, nil
	}
	// Any other input resumes live view if the user was scrolling back.
	if m.sshScrollOffset != 0 {
		m.sshScrollOffset = 0
	}
	payload := keyMsgToPTY(v)
	if len(payload) == 0 {
		return m, nil
	}
	m.writeSSHInput(payload)
	return m, nil
}

// captureSSHScrollback appends remote output (with ANSI sequences stripped)
// to the scrollback buffer so the user can scroll through past output even
// though vt10x does not retain a scroll history of its own.
func (m *monitorModel) captureSSHScrollback(data []byte) {
	if len(data) == 0 {
		return
	}
	text := stripANSI(string(data))
	// Treat stand-alone CR as a rewrite of the current line; drop content
	// before the CR to avoid duplicating progress-bar style updates.
	combined := m.sshScrollPending + text
	combined = strings.ReplaceAll(combined, "\r\n", "\n")
	lines := strings.Split(combined, "\n")
	// Last element is either a trailing newline ("") or a partial line; stash.
	m.sshScrollPending = lines[len(lines)-1]
	lines = lines[:len(lines)-1]
	for _, ln := range lines {
		if idx := strings.LastIndex(ln, "\r"); idx >= 0 {
			ln = ln[idx+1:]
		}
		m.sshScrollback = append(m.sshScrollback, ln)
	}
	const maxScrollback = 5000
	if len(m.sshScrollback) > maxScrollback {
		m.sshScrollback = m.sshScrollback[len(m.sshScrollback)-maxScrollback:]
	}
	// A new chunk means more history arrived behind the current scroll
	// window — adjust offset so the user keeps looking at the same line.
	if m.sshScrollOffset > 0 {
		m.sshScrollOffset += len(lines)
		m.clampSSHScroll()
	}
}

func (m *monitorModel) scrollSSH(delta int) {
	m.sshScrollOffset += delta
	m.clampSSHScroll()
}

func (m *monitorModel) clampSSHScroll() {
	maxOffset := len(m.sshScrollback) - m.sshRows
	if maxOffset < 0 {
		maxOffset = 0
	}
	if m.sshScrollOffset > maxOffset {
		m.sshScrollOffset = maxOffset
	}
	if m.sshScrollOffset < 0 {
		m.sshScrollOffset = 0
	}
}

// handleSSHResize adjusts the vt10x grid and notifies the remote side.
func (m *monitorModel) handleSSHResize() {
	if !m.sshMode || m.sshVT == nil || m.sshSess == nil {
		return
	}
	cols, rows := m.sshEmbeddedDims()
	if cols == m.sshCols && rows == m.sshRows {
		return
	}
	m.sshCols = cols
	m.sshRows = rows
	m.sshVT.Resize(cols, rows)
	_ = m.sshSess.Resize(cols, rows)
}

// renderSSHView produces the TUI view for the embedded SSH session.
func (m monitorModel) renderSSHView() string {
	header := accentStyle.Bold(true).Render(fmt.Sprintf("  ssh://%s@%s:%d (%s) ", m.sshCluster.User, m.sshCluster.Host, m.sshCluster.Port, m.sshCluster.Name))
	hint := dimStyle.Render("Ctrl+] detach · alt+↑/↓ scroll · pgup/pgdn page · alt+end live")
	if m.sshScrollOffset > 0 {
		hint = warnStyle.Render(fmt.Sprintf("↑ scrolled +%d (alt+end to follow)", m.sshScrollOffset)) +
			dimStyle.Render("   Ctrl+] detach")
	}

	body := m.renderVTBody()

	panel := lipgloss.NewStyle().
		BorderStyle(thinBorder).
		BorderForeground(ccAccent).
		Padding(0, 1).
		Render(body)

	status := dimStyle.Render(strings.TrimSpace("status: " + m.statusMsg))
	return lipgloss.JoinVertical(lipgloss.Left, header+"  "+hint, panel, status)
}

// renderVTBody iterates the vt10x grid and emits a styled string block.
// When the user has scrolled back it instead renders a window of the
// ANSI-stripped scrollback buffer.
func (m monitorModel) renderVTBody() string {
	if m.sshVT == nil {
		return dimStyle.Render("(session not ready)")
	}
	if m.sshScrollOffset > 0 {
		return m.renderScrollbackBody()
	}
	vt := m.sshVT
	vt.Lock()
	cols, rows := vt.Size()
	cur := vt.Cursor()
	cursorVisible := vt.CursorVisible()

	var lines []string
	for y := 0; y < rows; y++ {
		line := renderVTRow(vt, y, cols, cursorVisible, cur.X, cur.Y)
		lines = append(lines, line)
	}
	vt.Unlock()
	return strings.Join(lines, "\n")
}

// renderScrollbackBody paints a page of the captured scrollback buffer when
// the user is browsing history. Every row is padded to the full grid width
// so the panel keeps its live dimensions — otherwise lipgloss sizes the
// border to the longest line and the view visibly collapses.
func (m monitorModel) renderScrollbackBody() string {
	rows := m.sshRows
	cols := m.sshCols
	if rows <= 0 {
		rows = sshMinRows
	}
	if cols <= 0 {
		cols = sshMinCols
	}
	end := len(m.sshScrollback) - m.sshScrollOffset
	if end < 0 {
		end = 0
	}
	start := end - rows
	if start < 0 {
		start = 0
	}
	blank := strings.Repeat(" ", cols)
	out := make([]string, 0, rows)
	pad := func(line string) string {
		rs := []rune(line)
		if len(rs) > cols {
			rs = rs[:cols]
		}
		if len(rs) < cols {
			return string(rs) + strings.Repeat(" ", cols-len(rs))
		}
		return string(rs)
	}
	for i := start; i < end; i++ {
		out = append(out, pad(m.sshScrollback[i]))
	}
	for len(out) < rows {
		out = append(out, blank)
	}
	return strings.Join(out, "\n")
}

// renderVTRow renders a single grid row as a styled string.
func renderVTRow(vt vt10x.Terminal, row, cols int, cursorVisible bool, cursorX, cursorY int) string {
	var b strings.Builder
	var (
		runRunes strings.Builder
		runFG    vt10x.Color = vt10x.DefaultFG
		runBG    vt10x.Color = vt10x.DefaultBG
		runStart             = true
	)
	flush := func() {
		if runRunes.Len() == 0 {
			return
		}
		style := styleForColors(runFG, runBG)
		b.WriteString(style.Render(runRunes.String()))
		runRunes.Reset()
	}
	for x := 0; x < cols; x++ {
		cell := vt.Cell(x, row)
		ch := cell.Char
		if ch == 0 {
			ch = ' '
		}
		fg := cell.FG
		bg := cell.BG
		isCursor := cursorVisible && x == cursorX && row == cursorY
		if runStart {
			runFG = fg
			runBG = bg
			runStart = false
		}
		if isCursor {
			// Emit the current run first, then paint the cursor cell with
			// an explicit, always-visible color so the caret shows even on
			// empty cells with default FG/BG (where a plain invert would
			// collapse to the same color).
			flush()
			cursorStyle := lipgloss.NewStyle().
				Background(ccAccent).
				Foreground(lipgloss.Color("#101010")).
				Bold(true)
			b.WriteString(cursorStyle.Render(string(ch)))
			runFG = fg
			runBG = bg
			continue
		}
		if fg != runFG || bg != runBG {
			flush()
			runFG = fg
			runBG = bg
		}
		runRunes.WriteRune(ch)
	}
	flush()
	return b.String()
}

// styleForColors returns a lipgloss style for the given vt10x colors.
func styleForColors(fg, bg vt10x.Color) lipgloss.Style {
	style := lipgloss.NewStyle()
	if c, ok := vtColorToLipgloss(fg); ok {
		style = style.Foreground(c)
	}
	if c, ok := vtColorToLipgloss(bg); ok {
		style = style.Background(c)
	}
	return style
}

// vtColorToLipgloss maps a vt10x color to a lipgloss color. Returns ok=false
// for default so the caller leaves the attribute unset.
func vtColorToLipgloss(c vt10x.Color) (lipgloss.TerminalColor, bool) {
	if c == vt10x.DefaultFG || c == vt10x.DefaultBG || c == vt10x.DefaultCursor {
		return nil, false
	}
	// ANSI basic 16
	if c < 16 {
		return lipgloss.Color(fmt.Sprintf("%d", uint32(c))), true
	}
	// 256-color palette
	if c < 256 {
		return lipgloss.Color(fmt.Sprintf("%d", uint32(c))), true
	}
	// Truecolor (24-bit): stored in low 24 bits.
	r := (uint32(c) >> 16) & 0xff
	g := (uint32(c) >> 8) & 0xff
	bl := uint32(c) & 0xff
	return lipgloss.Color(fmt.Sprintf("#%02x%02x%02x", r, g, bl)), true
}