patterm/internal/app/cursorshift.go

package app

import (
	"strconv"
	"strings"
)

// cursorShifter rewrites cursor-positioning ANSI escapes in a PTY byte
// stream so the child's "row 1" becomes the host's "row 1+offset".
// This lets patterm reserve top rows for chrome (SPEC §4 tab bar)
// while keeping the child unaware.
//
// Sequences rewritten:
//   - CSI H, CSI <r> H, CSI <r>;<c> H — CUP
//   - CSI <r>;<c> f                   — HVP
//   - CSI <n> d                       — VPA (line position absolute)
//   - CSI <t>;<b> r                   — DECSTBM (scrolling region)
//
// Other sequences (SGR, mode set, OSC titles, DCS, alt-screen toggles)
// are forwarded byte-for-byte. The parser tracks OSC/DCS/SOS/PM/APC
// state so byte sequences inside those wrappers are NOT misread as
// CSI commands.
type cursorShifter struct {
	rowOffset int
	childRows int // viewport height in child rows; used for DECSTBM resets

	state     shifterState
	buf       []byte // bytes accumulated in current escape sequence (incl. introducer)
	csiPrefix []byte // private prefix bytes (?, >, =) after CSI
	pending   strings.Builder
}

type shifterState int

const (
	stNormal shifterState = iota
	stEsc
	stCSI
	stCSIPrefix // CSI <private-prefix>... — private prefix means we DON'T rewrite
	stOSC
	stOSCEsc // we saw ESC inside OSC; expect '\' to close ST
	stDCS
	stDCSEsc
	stSOSPMAPC // SOS/PM/APC body — terminator is ESC \
	stSOSPMAPCEsc
)

func newCursorShifter(rowOffset, childRows int) *cursorShifter {
	return &cursorShifter{rowOffset: rowOffset, childRows: childRows}
}

func (cs *cursorShifter) SetGeometry(rowOffset, childRows int) {
	cs.rowOffset = rowOffset
	cs.childRows = childRows
}

// Shift consumes a chunk of PTY-master bytes, applies row offsets to
// any complete CUP/HVP/VPA/DECSTBM sequences, and returns the rewritten
// bytes. Partial sequences are buffered across calls so a CSI that
// straddles two PTY reads still gets rewritten.
func (cs *cursorShifter) Shift(in []byte) []byte {
	cs.pending.Reset()
	for _, b := range in {
		cs.feed(b)
	}
	out := cs.pending.String()
	return []byte(out)
}

func (cs *cursorShifter) feed(b byte) {
	switch cs.state {
	case stNormal:
		if b == 0x1b {
			cs.state = stEsc
			cs.buf = cs.buf[:0]
			cs.buf = append(cs.buf, b)
			return
		}
		cs.pending.WriteByte(b)

	case stEsc:
		cs.buf = append(cs.buf, b)
		switch b {
		case '[':
			cs.state = stCSI
			cs.csiPrefix = cs.csiPrefix[:0]
		case ']':
			cs.state = stOSC
		case 'P':
			cs.state = stDCS
		case 'X', '^', '_':
			cs.state = stSOSPMAPC
		default:
			// Two-byte ESC sequence: ESC <something>. Forward as-is.
			cs.pending.Write(cs.buf)
			cs.state = stNormal
			cs.buf = cs.buf[:0]
		}

	case stCSI:
		// First non-param byte after CSI might be a private prefix
		// (?, >, =, etc., 0x3c..0x3f). If so, switch to CSIPrefix and
		// don't rewrite this sequence.
		if len(cs.csiPrefix) == 0 && len(cs.buf) == 2 && b >= 0x3c && b <= 0x3f {
			cs.csiPrefix = append(cs.csiPrefix, b)
			cs.buf = append(cs.buf, b)
			cs.state = stCSIPrefix
			return
		}
		cs.buf = append(cs.buf, b)
		if isCSIFinal(b) {
			cs.emitCSI()
			cs.state = stNormal
			cs.buf = cs.buf[:0]
		}

	case stCSIPrefix:
		cs.buf = append(cs.buf, b)
		if isCSIFinal(b) {
			// Private CSI; forward unchanged.
			cs.pending.Write(cs.buf)
			cs.state = stNormal
			cs.buf = cs.buf[:0]
		}

	case stOSC:
		cs.buf = append(cs.buf, b)
		switch b {
		case 0x07: // BEL
			cs.pending.Write(cs.buf)
			cs.state = stNormal
			cs.buf = cs.buf[:0]
		case 0x1b:
			cs.state = stOSCEsc
		}

	case stOSCEsc:
		cs.buf = append(cs.buf, b)
		// ESC \ terminates ST.
		cs.pending.Write(cs.buf)
		cs.state = stNormal
		cs.buf = cs.buf[:0]

	case stDCS:
		cs.buf = append(cs.buf, b)
		if b == 0x1b {
			cs.state = stDCSEsc
		}

	case stDCSEsc:
		cs.buf = append(cs.buf, b)
		cs.pending.Write(cs.buf)
		cs.state = stNormal
		cs.buf = cs.buf[:0]

	case stSOSPMAPC:
		cs.buf = append(cs.buf, b)
		if b == 0x1b {
			cs.state = stSOSPMAPCEsc
		}

	case stSOSPMAPCEsc:
		cs.buf = append(cs.buf, b)
		cs.pending.Write(cs.buf)
		cs.state = stNormal
		cs.buf = cs.buf[:0]
	}
}

// emitCSI writes the buffered CSI sequence to pending, rewriting row
// coordinates for CUP/HVP/VPA/DECSTBM.
func (cs *cursorShifter) emitCSI() {
	// cs.buf is ESC [ <params...> <final>. Slice out params + final.
	if len(cs.buf) < 3 {
		cs.pending.Write(cs.buf)
		return
	}
	final := cs.buf[len(cs.buf)-1]
	paramsRaw := cs.buf[2 : len(cs.buf)-1]
	// Intermediate bytes can appear before the final (rare). Skip
	// rewriting if any are present.
	for _, b := range paramsRaw {
		if b >= 0x20 && b <= 0x2f {
			cs.pending.Write(cs.buf)
			return
		}
	}
	switch final {
	case 'H', 'f':
		// CUP/HVP: r;c (both default 1).
		r, c, ok := parseTwoParams(paramsRaw)
		if !ok {
			cs.pending.Write(cs.buf)
			return
		}
		r += cs.rowOffset
		cs.pending.WriteString("\x1b[")
		cs.pending.WriteString(strconv.Itoa(r))
		cs.pending.WriteByte(';')
		cs.pending.WriteString(strconv.Itoa(c))
		cs.pending.WriteByte(final)
	case 'd':
		// VPA: row.
		r, ok := parseOneParam(paramsRaw, 1)
		if !ok {
			cs.pending.Write(cs.buf)
			return
		}
		r += cs.rowOffset
		cs.pending.WriteString("\x1b[")
		cs.pending.WriteString(strconv.Itoa(r))
		cs.pending.WriteByte(final)
	case 'r':
		// DECSTBM: top;bot. Empty params (\x1b[r) means "reset to the
		// full screen" from the child's point of view — for us that's
		// the viewport, not the host's full screen. Rewriting it as
		// (1,1)+offset would produce \x1b[4;4r, a one-row region that
		// causes catastrophic scroll-up of the replayed snapshot.
		if len(paramsRaw) == 0 && cs.childRows > 0 {
			cs.pending.WriteString("\x1b[")
			cs.pending.WriteString(strconv.Itoa(cs.rowOffset + 1))
			cs.pending.WriteByte(';')
			cs.pending.WriteString(strconv.Itoa(cs.rowOffset + cs.childRows))
			cs.pending.WriteByte(final)
			return
		}
		top, bot, ok := parseTwoParams(paramsRaw)
		if !ok {
			cs.pending.Write(cs.buf)
			return
		}
		top += cs.rowOffset
		bot += cs.rowOffset
		cs.pending.WriteString("\x1b[")
		cs.pending.WriteString(strconv.Itoa(top))
		cs.pending.WriteByte(';')
		cs.pending.WriteString(strconv.Itoa(bot))
		cs.pending.WriteByte(final)
	default:
		cs.pending.Write(cs.buf)
	}
}

func isCSIFinal(b byte) bool { return b >= 0x40 && b <= 0x7e }

func parseTwoParams(raw []byte) (int, int, bool) {
	parts := strings.Split(string(raw), ";")
	if len(parts) > 2 {
		return 0, 0, false
	}
	a := 1
	b := 1
	if len(parts) >= 1 && parts[0] != "" {
		n, err := strconv.Atoi(parts[0])
		if err != nil {
			return 0, 0, false
		}
		a = n
	}
	if len(parts) >= 2 && parts[1] != "" {
		n, err := strconv.Atoi(parts[1])
		if err != nil {
			return 0, 0, false
		}
		b = n
	}
	return a, b, true
}

func parseOneParam(raw []byte, def int) (int, bool) {
	s := string(raw)
	if s == "" {
		return def, true
	}
	n, err := strconv.Atoi(s)
	if err != nil {
		return 0, false
	}
	return n, true
}