HostKeeper/pkg/tui/session.go

package tui

import (
	"context"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	cryptossh "golang.org/x/crypto/ssh"

	"git.tukangketik.id/swanadiva/hostkeeper/internal/models"
	sshclient "git.tukangketik.id/swanadiva/hostkeeper/pkg/ssh"
)

// SessionTab represents an interactive SSH session as a tab
type SessionTab struct {
	host    *models.Host
	dataDir string

	client *sshclient.Client
	sess   *cryptossh.Session

	width  int
	height int

	mu        sync.Mutex
	buffer    strings.Builder
	connected bool
	done      bool
	err       error

	stdinPipe   io.WriteCloser
	stdinMu     sync.Mutex
	outputCh    chan sessionOutputMsg
	doneCh      chan sessionDoneMsg
	windowCh    chan struct{}
	closeOnce   sync.Once
	closed      chan struct{}

	cancel context.CancelFunc
}

// NewSessionTab creates a new SSH session tab
func NewSessionTab(host *models.Host, dataDir string) *SessionTab {
	return &SessionTab{
		host:    host,
		dataDir: dataDir,
		width:   80,
		height:  24,

		outputCh: make(chan sessionOutputMsg, 256),
		doneCh:   make(chan sessionDoneMsg, 1),
		windowCh: make(chan struct{}, 8),
		closed:   make(chan struct{}),
	}
}

// Name returns the tab name
func (t *SessionTab) Name() string {
	return t.host.Name
}

// Init starts the SSH connection in the background
func (t *SessionTab) Init() tea.Cmd {
	ctx, cancel := context.WithCancel(context.Background())
	t.cancel = cancel

	go t.connectAndStream(ctx)
	return t.pollSession
}

// pollSession polls for session output and returns it as a message to trigger re-render
func (t *SessionTab) pollSession() tea.Msg {
	select {
	case msg := <-t.outputCh:
		t.mu.Lock()
		t.buffer.WriteString(string(msg))
		t.mu.Unlock()
		// Return the msg to trigger Update → View re-render
		return msg
	case done := <-t.doneCh:
		return done
	default:
		return nil
	}
}

// connectAndStream connects to the host and streams output
func (t *SessionTab) connectAndStream(ctx context.Context) {
	timeout := 30 * time.Second
	client := sshclient.NewClient(t.host, timeout)

	if err := client.Connect(ctx); err != nil {
		t.doneCh <- sessionDoneMsg{err: fmt.Errorf("connection failed: %w", err)}
		return
	}

	t.mu.Lock()
	t.client = client
	t.connected = true
	t.mu.Unlock()

	sshClient := client.GetClient()
	sess, err := sshClient.NewSession()
	if err != nil {
		client.Close()
		t.doneCh <- sessionDoneMsg{err: fmt.Errorf("session creation failed: %w", err)}
		return
	}

	t.mu.Lock()
	t.sess = sess
	t.mu.Unlock()

	width := t.width
	height := t.height

	modes := cryptossh.TerminalModes{
		cryptossh.ECHO:          1,
		cryptossh.TTY_OP_ISPEED: 14400,
		cryptossh.TTY_OP_OSPEED: 14400,
	}

	if err := sess.RequestPty("xterm-256color", height, width, modes); err != nil {
		client.Close()
		t.doneCh <- sessionDoneMsg{err: fmt.Errorf("PTY request failed: %w", err)}
		return
	}

	stdin, err := sess.StdinPipe()
	if err != nil {
		client.Close()
		t.doneCh <- sessionDoneMsg{err: fmt.Errorf("stdin pipe failed: %w", err)}
		return
	}
	t.stdinPipe = stdin

	stdout, err := sess.StdoutPipe()
	if err != nil {
		client.Close()
		t.doneCh <- sessionDoneMsg{err: fmt.Errorf("stdout pipe failed: %w", err)}
		return
	}

	stderr, err := sess.StderrPipe()
	if err != nil {
		client.Close()
		t.doneCh <- sessionDoneMsg{err: fmt.Errorf("stderr pipe failed: %w", err)}
		return
	}

	if err := sess.Shell(); err != nil {
		client.Close()
		t.doneCh <- sessionDoneMsg{err: fmt.Errorf("shell start failed: %w", err)}
		return
	}

	// Stream stdout
	go func() {
		buf := make([]byte, 4096)
		for {
			n, err := stdout.Read(buf)
			if n > 0 {
				select {
				case t.outputCh <- sessionOutputMsg(string(buf[:n])):
				default:
				}
			}
			if err != nil {
				return
			}
		}
	}()

	// Stream stderr
	go func() {
		buf := make([]byte, 4096)
		for {
			n, err := stderr.Read(buf)
			if n > 0 {
				select {
				case t.outputCh <- sessionOutputMsg(string(buf[:n])):
				default:
				}
			}
			if err != nil {
				return
		}
		}
	}()

	// Handle window resize
	go func() {
		for {
			select {
			case <-t.windowCh:
				t.mu.Lock()
				w, h := t.width, t.height
				t.mu.Unlock()
				sess.WindowChange(h, w)
			case <-t.closed:
				return
			case <-ctx.Done():
				return
			}
		}
	}()

	sess.Wait()
	client.Close()
	t.doneCh <- sessionDoneMsg{err: nil}
}

// Update handles messages for the session tab
func (t *SessionTab) Update(msg tea.Msg) (Tab, tea.Cmd) {
	if t.done {
		return t, nil
	}

	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		t.mu.Lock()
		t.width = msg.Width
		t.height = msg.Height
		t.mu.Unlock()
		select {
		case t.windowCh <- struct{}{}:
		default:
		}

	case tea.KeyMsg:
		data := keyMsgToBytes(msg)
		if len(data) > 0 {
			t.stdinMu.Lock()
			if t.stdinPipe != nil {
				t.stdinPipe.Write(data)
			}
			t.stdinMu.Unlock()
		}

	case sessionOutputMsg:
		// Buffer already updated in pollSession, but we return pollSession
		// to keep the polling loop alive

	case sessionDoneMsg:
		t.mu.Lock()
		t.done = true
		if msg.err != nil {
			t.err = msg.err
		}
		t.mu.Unlock()
	}

	return t, t.pollSession
}

// View renders the session tab content
func (t *SessionTab) View() string {
	t.mu.Lock()
	defer t.mu.Unlock()

	if !t.connected {
		return lipgloss.NewStyle().Height(t.height - 2).Render(
			HighlightStyle.Render(fmt.Sprintf("Connecting to %s...", t.host.Name)))
	}

	if t.done && t.err != nil {
		return fmt.Sprintf("%s\n\n%s",
			ErrorStyle.Render(fmt.Sprintf("Session ended: %v", t.err)),
			SubtitleStyle.Render("Press Ctrl+Q to close this tab"))
	}

	if t.done {
		return fmt.Sprintf("%s\n\n%s",
			SubtitleStyle.Render("Session ended"),
			SubtitleStyle.Render("Press Ctrl+Q to close this tab"))
	}

	content := t.buffer.String()
	if content == "" {
		return lipgloss.NewStyle().Height(t.height - 2).Render(
			SubtitleStyle.Render("Connected. Waiting for output..."))
	}

	lines := strings.Split(content, "\n")
	const maxLines = 1000
	if len(lines) > maxLines {
		lines = lines[len(lines)-maxLines:]
	}
	visible := strings.Join(lines, "\n")

	return lipgloss.NewStyle().Height(t.height - 3).Render(visible)
}

// Close terminates the session
func (t *SessionTab) Close() {
	t.closeOnce.Do(func() {
		close(t.closed)
		if t.cancel != nil {
			t.cancel()
		}
		t.mu.Lock()
		if t.sess != nil {
			t.sess.Close()
		}
		if t.client != nil {
			t.client.Close()
		}
		t.mu.Unlock()
	})
}

// keyMsgToBytes converts a Bubble Tea key message to SSH terminal bytes
func keyMsgToBytes(msg tea.KeyMsg) []byte {
	switch msg.Type {
	case tea.KeyEnter:
		return []byte("\r")
	case tea.KeyBackspace:
		return []byte("\x7f")
	case tea.KeyTab:
		return []byte("\t")
	case tea.KeySpace:
		return []byte(" ")
	case tea.KeyUp:
		return []byte("\x1b[A")
	case tea.KeyDown:
		return []byte("\x1b[B")
	case tea.KeyRight:
		return []byte("\x1b[C")
	case tea.KeyLeft:
		return []byte("\x1b[D")
	case tea.KeyEscape:
		return []byte("\x1b")
	case tea.KeyDelete:
		return []byte("\x1b[3~")
	case tea.KeyHome:
		return []byte("\x1b[H")
	case tea.KeyEnd:
		return []byte("\x1b[F")
	case tea.KeyPgUp:
		return []byte("\x1b[5~")
	case tea.KeyPgDown:
		return []byte("\x1b[6~")
	default:
		if len(msg.Runes) > 0 {
			return []byte(string(msg.Runes))
		}
		return []byte(msg.String())
	}
}