termd/internal/session/session.go

// Package session owns the PTY processes and their terminal emulators. All
// vt/ultraviolet types stay behind this package (and internal/keymap).
package session

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	uv "github.com/charmbracelet/ultraviolet"
	"github.com/charmbracelet/x/ansi"
	"github.com/charmbracelet/x/vt"
	"github.com/creack/pty"

	"termd/internal/api"
	"termd/internal/keymap"
)

// ErrExited is returned when input is sent to a session whose process has
// already exited.
var ErrExited = errors.New("session process has exited")

// MouseError is a mouse request rejected because the application's mode state
// can't deliver it; it carries the modes so the caller can see why.
type MouseError struct {
	Reason string
	Modes  api.Modes
}

func (e *MouseError) Error() string { return e.Reason }

// Options configures a new session.
type Options struct {
	Command []string
	Cwd     string
	Env     map[string]string
	Cols    int
	Rows    int
	Term    string
}

// Session is one PTY process plus the emulator that models its screen.
type Session struct {
	ID      string
	command []string
	pid     int

	// mu guards em, modes, size, exit state and the subscriber set. The
	// emulator's mode callbacks fire inside em.Write while mu is held and
	// write modeState fields directly — they must never re-lock.
	mu       sync.Mutex
	em       *vt.Emulator
	modes    modeState
	cols     int
	rows     int
	exited   bool
	exitCode int
	subs     map[int]chan []byte
	nextSub  int

	ptmx       *os.File
	cmd        *exec.Cmd
	inq        *byteQueue
	readerDone chan struct{}
	exitedCh   chan struct{}
}

// New starts the command in a fresh PTY and begins emulating its output.
func New(id string, opts Options) (*Session, error) {
	if len(opts.Command) == 0 {
		shell := os.Getenv("SHELL")
		if shell == "" {
			shell = "/bin/sh"
		}
		opts.Command = []string{shell}
	}
	if opts.Cols <= 0 {
		opts.Cols = 80
	}
	if opts.Rows <= 0 {
		opts.Rows = 24
	}
	if opts.Term == "" {
		opts.Term = "xterm-256color"
	}

	cmd := exec.Command(opts.Command[0], opts.Command[1:]...)
	cmd.Dir = opts.Cwd
	cmd.Env = append(os.Environ(), "TERM="+opts.Term)
	for k, v := range opts.Env {
		cmd.Env = append(cmd.Env, k+"="+v)
	}

	ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{
		Cols: uint16(opts.Cols),
		Rows: uint16(opts.Rows),
	})
	if err != nil {
		return nil, fmt.Errorf("starting %v in pty: %w", opts.Command, err)
	}

	s := &Session{
		ID:         id,
		command:    opts.Command,
		pid:        cmd.Process.Pid,
		em:         vt.NewEmulator(opts.Cols, opts.Rows),
		cols:       opts.Cols,
		rows:       opts.Rows,
		subs:       make(map[int]chan []byte),
		ptmx:       ptmx,
		cmd:        cmd,
		inq:        newByteQueue(),
		readerDone: make(chan struct{}),
		exitedCh:   make(chan struct{}),
	}
	s.modes.cursorVisible = true
	s.em.SetCallbacks(vt.Callbacks{
		EnableMode:       func(m ansi.Mode) { s.modes.set(m, true) },
		DisableMode:      func(m ansi.Mode) { s.modes.set(m, false) },
		AltScreen:        func(on bool) { s.modes.altScreen = on },
		CursorVisibility: func(v bool) { s.modes.cursorVisible = v },
		Title:            func(t string) { s.modes.title = t },
	})

	go s.readLoop()
	go s.pumpLoop()
	go s.writeLoop()
	go s.waitLoop()
	return s, nil
}

// readLoop feeds PTY output into the emulator and broadcasts the same raw
// bytes to attach subscribers. Exits when the PTY returns an error (EIO once
// the child side is closed).
func (s *Session) readLoop() {
	buf := make([]byte, 32*1024)
	for {
		n, err := s.ptmx.Read(buf)
		if n > 0 {
			chunk := append([]byte(nil), buf[:n]...)
			s.mu.Lock()
			_, _ = s.em.Write(chunk)
			for id, ch := range s.subs {
				select {
				case ch <- chunk:
				default:
					// Stalled attach client: drop it rather than the session.
					close(ch)
					delete(s.subs, id)
				}
			}
			s.mu.Unlock()
		}
		if err != nil {
			break
		}
	}
	// No more output will ever arrive: release attach subscribers.
	s.mu.Lock()
	for id, ch := range s.subs {
		close(ch)
		delete(s.subs, id)
	}
	s.mu.Unlock()
	close(s.readerDone)
}

// pumpLoop drains the emulator's encoded-input pipe into an unbounded queue.
// This is what keeps SendKey/SendText (called under mu) from ever blocking on
// a child that is slow to read stdin — and with them the whole session.
func (s *Session) pumpLoop() {
	buf := make([]byte, 32*1024)
	for {
		n, err := s.em.Read(buf)
		if n > 0 {
			s.inq.push(append([]byte(nil), buf[:n]...))
		}
		if err != nil {
			s.inq.close()
			return
		}
	}
}

// writeLoop forwards queued input bytes to the PTY.
func (s *Session) writeLoop() {
	for {
		b, ok := s.inq.pop()
		if !ok {
			return
		}
		if _, err := s.ptmx.Write(b); err != nil {
			return
		}
	}
}

func (s *Session) waitLoop() {
	err := s.cmd.Wait()
	code := 0
	if err != nil {
		if exitErr, ok := errors.AsType[*exec.ExitError](err); ok {
			code = exitErr.ExitCode()
		} else {
			code = -1
		}
	}
	s.mu.Lock()
	s.exited = true
	s.exitCode = code
	s.mu.Unlock()
	close(s.exitedCh)
}

// Kill terminates the process (signal, then SIGKILL after 3s) and tears the
// session down. Safe to call on an already-exited session.
func (s *Session) Kill(sig os.Signal) error {
	s.mu.Lock()
	exited := s.exited
	s.mu.Unlock()
	if !exited {
		if err := s.cmd.Process.Signal(sig); err != nil && !errors.Is(err, os.ErrProcessDone) {
			return fmt.Errorf("signaling pid %d: %w", s.pid, err)
		}
		select {
		case <-s.exitedCh:
		case <-time.After(3 * time.Second):
			_ = s.cmd.Process.Kill()
			select {
			case <-s.exitedCh:
			case <-time.After(3 * time.Second):
				return fmt.Errorf("pid %d did not die after SIGKILL", s.pid)
			}
		}
	}
	// Unblock all loops: PTY reads/writes fail, and closing the emulator's
	// input pipe EOFs pumpLoop. Deliberately NOT em.Close(): it writes an
	// unsynchronized flag that races with the pump's concurrent em.Read.
	_ = s.ptmx.Close()
	if pw, ok := s.em.InputPipe().(io.Closer); ok {
		_ = pw.Close()
	}
	return nil
}

// SendText sends literal text (no key-name interpretation).
func (s *Session) SendText(text string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.exited {
		return ErrExited
	}
	s.em.SendText(text)
	return nil
}

// SendKey sends a parsed named key, either through the emulator's mode-aware
// encoder or as pre-encoded bytes. Both paths go through the same input pipe
// so ordering within a request is preserved.
func (s *Session) SendKey(k keymap.Key) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.exited {
		return ErrExited
	}
	if ev, ok := k.Event(); ok {
		s.em.SendKey(ev)
		return nil
	}
	raw, _ := k.Raw()
	if _, err := s.em.InputPipe().Write(raw); err != nil {
		return fmt.Errorf("writing key %s: %w", k.Name, err)
	}
	return nil
}

// SendRaw writes bytes to the terminal input verbatim.
func (s *Session) SendRaw(b []byte) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.exited {
		return ErrExited
	}
	if _, err := s.em.InputPipe().Write(b); err != nil {
		return fmt.Errorf("writing raw input: %w", err)
	}
	return nil
}

// WriteHuman writes an attached human's already-encoded terminal bytes
// straight to the PTY, bypassing the emulator's encoder.
func (s *Session) WriteHuman(b []byte) error {
	_, err := s.ptmx.Write(b)
	return err
}

var mouseButtons = map[string]uv.MouseButton{
	"left":       uv.MouseLeft,
	"middle":     uv.MouseMiddle,
	"right":      uv.MouseRight,
	"wheel-up":   uv.MouseWheelUp,
	"wheel-down": uv.MouseWheelDown,
}

var mouseMods = map[string]uv.KeyMod{
	"ctrl":  uv.ModCtrl,
	"alt":   uv.ModAlt,
	"shift": uv.ModShift,
}

// SendMouse validates a mouse request against the application's tracking
// modes and injects the event(s). Requests the app can't receive are
// rejected with a *MouseError rather than silently dropped.
func (s *Session) SendMouse(req api.MouseRequest) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.exited {
		return ErrExited
	}

	mods := uv.KeyMod(0)
	for _, m := range req.Modifiers {
		mod, ok := mouseMods[m]
		if !ok {
			return fmt.Errorf("unknown modifier %q (want ctrl, alt or shift)", m)
		}
		mods |= mod
	}

	button := uv.MouseNone
	if req.Button != "" {
		b, ok := mouseButtons[req.Button]
		if !ok {
			return fmt.Errorf("unknown button %q", req.Button)
		}
		button = b
	}

	mm := s.modes.mouseMode()
	reject := func(reason string) error {
		return &MouseError{Reason: reason, Modes: s.modes.api()}
	}
	if mm == "off" {
		return reject("application has not enabled mouse reporting")
	}

	mouse := func(b uv.MouseButton) uv.Mouse {
		return uv.Mouse{X: req.X, Y: req.Y, Button: b, Mod: mods}
	}

	switch req.Type {
	case "press", "release", "click":
		if button == uv.MouseNone || button == uv.MouseWheelUp || button == uv.MouseWheelDown {
			return fmt.Errorf("%s needs button left, middle or right", req.Type)
		}
		if mm == "x10" && req.Type != "press" {
			return reject("application only enabled X10 mouse mode (?9), which reports presses only")
		}
		if req.Type != "release" {
			s.em.SendMouse(uv.MouseClickEvent(mouse(button)))
		}
		if req.Type != "press" {
			s.em.SendMouse(uv.MouseReleaseEvent(mouse(button)))
		}
	case "scroll":
		if button != uv.MouseWheelUp && button != uv.MouseWheelDown {
			return fmt.Errorf("scroll needs button wheel-up or wheel-down")
		}
		if mm == "x10" {
			return reject("application only enabled X10 mouse mode (?9), which cannot report scrolling")
		}
		s.em.SendMouse(uv.MouseWheelEvent(mouse(button)))
	case "drag":
		if button == uv.MouseNone {
			return fmt.Errorf("drag needs a button")
		}
		if mm != "button_event" && mm != "any_event" {
			return reject("drag needs mouse mode ?1002 or ?1003; application enabled " + mm)
		}
		s.em.SendMouse(uv.MouseMotionEvent(mouse(button)))
	case "move":
		if mm != "any_event" {
			return reject("bare motion needs mouse mode ?1003; application enabled " + mm)
		}
		s.em.SendMouse(uv.MouseMotionEvent(mouse(uv.MouseNone)))
	default:
		return fmt.Errorf("unknown mouse event type %q", req.Type)
	}
	return nil
}

// Resize changes both the PTY and the emulator size, PTY first so the
// SIGWINCH the child receives matches what the emulator models.
func (s *Session) Resize(cols, rows int) error {
	if cols <= 0 || rows <= 0 {
		return fmt.Errorf("invalid size %dx%d", cols, rows)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := pty.Setsize(s.ptmx, &pty.Winsize{Cols: uint16(cols), Rows: uint16(rows)}); err != nil {
		return fmt.Errorf("resizing pty: %w", err)
	}
	s.em.Resize(cols, rows)
	s.cols, s.rows = cols, rows
	return nil
}

// Snapshot returns the current screen state. withRaw adds the styled ANSI
// rendering.
func (s *Session) Snapshot(withRaw bool) api.Screen {
	s.mu.Lock()
	defer s.mu.Unlock()

	lines := strings.Split(s.em.String(), "\n")
	if len(lines) > s.rows {
		lines = lines[:s.rows]
	}
	for len(lines) < s.rows {
		lines = append(lines, "")
	}
	for i, l := range lines {
		lines[i] = strings.TrimRight(l, " \t")
	}

	cur := s.em.CursorPosition()
	scr := api.Screen{
		Lines:     lines,
		Cols:      s.cols,
		Rows:      s.rows,
		Cursor:    api.Cursor{X: cur.X, Y: cur.Y, Visible: s.modes.cursorVisible},
		AltScreen: s.modes.altScreen,
		Title:     s.modes.title,
		Modes:     s.modes.api(),
		Exited:    s.exited,
	}
	if s.exited {
		code := s.exitCode
		scr.ExitCode = &code
	}
	if withRaw {
		scr.Raw = s.em.Render()
	}
	return scr
}

// RenderFrame builds a full repaint of the emulated grid using absolute
// positioning only, for viewers whose terminal is larger than the session:
// the raw byte stream encodes the session's geometry (scroll margins, wrap
// column) and breaks on a bigger screen, while a frame just paints the
// session box top-left and clears everything outside it to spaces. Each line
// resets styling before clearing so no background color bleeds to the edge.
func (s *Session) RenderFrame(clearAll bool) []byte {
	s.mu.Lock()
	defer s.mu.Unlock()
	var b strings.Builder
	if clearAll {
		b.WriteString("\x1b[2J")
	}
	b.WriteString("\x1b[?25l\x1b[H")
	lines := strings.Split(s.em.Render(), "\n")
	if len(lines) > s.rows {
		lines = lines[:s.rows]
	}
	for i, l := range lines {
		if i > 0 {
			b.WriteString("\r\n")
		}
		b.WriteString(l)
		b.WriteString("\x1b[0m\x1b[K")
	}
	b.WriteString("\x1b[J")
	cur := s.em.CursorPosition()
	fmt.Fprintf(&b, "\x1b[%d;%dH", cur.Y+1, cur.X+1)
	if s.modes.cursorVisible {
		b.WriteString("\x1b[?25h")
	}
	return []byte(b.String())
}

// Info returns the session's metadata.
func (s *Session) Info() api.SessionInfo {
	s.mu.Lock()
	defer s.mu.Unlock()
	info := api.SessionInfo{
		ID:      s.ID,
		PID:     s.pid,
		Command: s.command,
		Cols:    s.cols,
		Rows:    s.rows,
		Title:   s.modes.title,
		Exited:  s.exited,
	}
	if s.exited {
		code := s.exitCode
		info.ExitCode = &code
	}
	return info
}

// Subscribe registers an attach client. It returns the initial repaint frame
// (clear + full styled render), the live channel, and an unsubscribe func.
// The channel is closed if the client stalls or the session is torn down.
func (s *Session) Subscribe() (snapshot []byte, ch <-chan []byte, cancel func()) {
	s.mu.Lock()
	defer s.mu.Unlock()
	c := make(chan []byte, 256)
	id := s.nextSub
	s.nextSub++
	s.subs[id] = c
	snap := []byte("\x1b[2J\x1b[H" + s.em.Render())
	return snap, c, func() {
		s.mu.Lock()
		defer s.mu.Unlock()
		if _, ok := s.subs[id]; ok {
			close(c)
			delete(s.subs, id)
		}
	}
}

// Exited returns a channel closed when the process exits, and the exit code
// once it has.
func (s *Session) Exited() (<-chan struct{}, int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.exitedCh, s.exitCode
}

// byteQueue is an unbounded FIFO of byte chunks: pushes never block, pops
// block until data or close.
type byteQueue struct {
	mu     sync.Mutex
	cond   *sync.Cond
	chunks [][]byte
	closed bool
}

func newByteQueue() *byteQueue {
	q := &byteQueue{}
	q.cond = sync.NewCond(&q.mu)
	return q
}

func (q *byteQueue) push(b []byte) {
	q.mu.Lock()
	defer q.mu.Unlock()
	if q.closed {
		return
	}
	q.chunks = append(q.chunks, b)
	q.cond.Signal()
}

func (q *byteQueue) pop() ([]byte, bool) {
	q.mu.Lock()
	defer q.mu.Unlock()
	for len(q.chunks) == 0 && !q.closed {
		q.cond.Wait()
	}
	if len(q.chunks) == 0 {
		return nil, false
	}
	b := q.chunks[0]
	q.chunks = q.chunks[1:]
	return b, true
}

func (q *byteQueue) close() {
	q.mu.Lock()
	defer q.mu.Unlock()
	q.closed = true
	q.cond.Broadcast()
}