arr-tui/internal/ui/kitty.go

package ui

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"image"
	"image/png"
	"strings"

	"github.com/gdamore/tcell/v2"
	"github.com/rivo/tview"
)

// The Kitty graphics protocol with Unicode placeholders, as supported by
// kitty and Ghostty: the poster is uploaded to the terminal once, and then
// shown by writing special placeholder characters into ordinary text cells.
// Because the picture is made of text cells it behaves like text: popups and
// dialogs drawn over it cover it, it vanishes when the view changes, and the
// normal screen diffing keeps it in place. The poster is shown at the
// terminal's full resolution.

// placeholderRune marks a cell that shows a piece of an image. It is followed
// by combining characters that give the row and column in the image; the
// image id is the cell's foreground color.
const placeholderRune = '\U0010EEEE'

const kittyChunk = 4096 // bytes of base64 per escape sequence, the protocol maximum

// kittyGraphics uploads images to the terminal and keeps track of them.
type kittyGraphics struct {
	write  func([]byte) error // writes to the terminal, in step with the screen
	nextID uint32
	ids    map[cropKey]uint32
	latest map[*image.RGBA]cropKey // the crop of each image that is currently uploaded
	placed map[uint32][2]int       // columns and rows each image is currently placed with
}

// cropKey identifies an uploaded image: a poster and the part of it that was sent.
type cropKey struct {
	img  *image.RGBA
	rect image.Rectangle
}

func newKittyGraphics(write func([]byte) error) *kittyGraphics {
	return &kittyGraphics{write: write, nextID: 1, ids: map[cropKey]uint32{}, latest: map[*image.RGBA]cropKey{}, placed: map[uint32][2]int{}}
}

// kittyAPC builds one graphics command: ESC _ G <keys> ; <payload> ESC \.
func kittyAPC(keys string, payload []byte) []byte {
	var b bytes.Buffer
	b.WriteString("\x1b_G")
	b.WriteString(keys)
	if len(payload) > 0 {
		b.WriteByte(';')
		b.Write(payload)
	}
	b.WriteString("\x1b\\")
	return b.Bytes()
}

// imageID returns the terminal's id for the part rect of the image,
// uploading it on first use. When the same image is needed with another crop
// (after a resize) the earlier upload is removed.
func (k *kittyGraphics) imageID(img *image.RGBA, rect image.Rectangle) (uint32, error) {
	key := cropKey{img, rect}
	if id, ok := k.ids[key]; ok {
		return id, nil
	}
	var data bytes.Buffer
	enc := png.Encoder{CompressionLevel: png.BestSpeed}
	if err := enc.Encode(&data, img.SubImage(rect)); err != nil {
		return 0, err
	}
	encoded := base64.StdEncoding.EncodeToString(data.Bytes())
	id := k.nextID
	for first, rest := true, encoded; len(rest) > 0; first = false {
		n := min(len(rest), kittyChunk)
		more := 0
		if n < len(rest) {
			more = 1
		}
		keys := fmt.Sprintf("m=%d", more)
		if first {
			// Transmit only (a=t), PNG data (f=100) directly in the command (t=d),
			// no replies (q=2): a reply would arrive as keyboard input.
			keys = fmt.Sprintf("a=t,f=100,t=d,i=%d,q=2,m=%d", id, more)
		}
		if err := k.write(kittyAPC(keys, []byte(rest[:n]))); err != nil {
			return 0, err
		}
		rest = rest[n:]
	}
	k.nextID++
	if old, ok := k.latest[img]; ok {
		if oldID, ok := k.ids[old]; ok {
			_ = k.write(kittyAPC(fmt.Sprintf("a=d,d=I,i=%d,q=2", oldID), nil))
			delete(k.placed, oldID)
			delete(k.ids, old)
		}
	}
	k.ids[key] = id
	k.latest[img] = key
	return id, nil
}

// place makes the image cover cols×rows cells for the placeholders. A new size
// replaces the previous placement.
func (k *kittyGraphics) place(id uint32, cols, rows int) error {
	if p, ok := k.placed[id]; ok {
		if p == [2]int{cols, rows} {
			return nil
		}
		// Remove the old placement (the uploaded data stays) before making a new one.
		if err := k.write(kittyAPC(fmt.Sprintf("a=d,d=i,i=%d,q=2", id), nil)); err != nil {
			return err
		}
	}
	// U=1: a virtual placement, shown wherever placeholder cells refer to it.
	if err := k.write(kittyAPC(fmt.Sprintf("a=p,U=1,i=%d,c=%d,r=%d,q=2", id, cols, rows), nil)); err != nil {
		return err
	}
	k.placed[id] = [2]int{cols, rows}
	return nil
}

// invalidate drops every placement (the uploaded pictures stay) so the next
// draw places them again, e.g. after the terminal was resized.
func (k *kittyGraphics) invalidate() {
	for id := range k.placed {
		_ = k.write(kittyAPC(fmt.Sprintf("a=d,d=i,i=%d,q=2", id), nil))
	}
	k.placed = map[uint32][2]int{}
}

// deleteAll removes every image uploaded so far from the terminal.
func (k *kittyGraphics) deleteAll() {
	for _, id := range k.ids {
		_ = k.write(kittyAPC(fmt.Sprintf("a=d,d=I,i=%d,q=2", id), nil))
	}
	k.ids = map[cropKey]uint32{}
	k.latest = map[*image.RGBA]cropKey{}
	k.placed = map[uint32][2]int{}
}

// placeholderCell returns what to draw in the cell at (row, col) of the image
// with the given id: the placeholder character, its combining characters and
// the style that carries the id.
func placeholderCell(id uint32, row, col int) (main rune, comb []rune, style tcell.Style) {
	comb = []rune{kittyDiacritics[row], kittyDiacritics[col]}
	fg := tcell.NewRGBColor(int32(id>>16&0xFF), int32(id>>8&0xFF), int32(id&0xFF))
	// The background shows wherever the picture does not cover the cell (a
	// sliver at the edge): use the panes' own color so it is not visible.
	return placeholderRune, comb, tcell.StyleDefault.Foreground(fg).Background(tview.Styles.PrimitiveBackgroundColor)
}

// maxPlaceholders is how many rows and columns can be numbered.
const maxPlaceholders = len(kittyDiacritics)

// kittyTerminal reports whether the environment looks like a terminal that
// supports Kitty graphics with Unicode placeholders: kitty and Ghostty, not
// inside tmux or screen (which swallow the graphics commands).
func kittyTerminal(env func(string) string) bool {
	term := strings.ToLower(env("TERM"))
	if env("TMUX") != "" || env("STY") != "" || strings.HasPrefix(term, "screen") || strings.HasPrefix(term, "tmux") {
		return false
	}
	switch {
	case strings.Contains(term, "kitty"), strings.Contains(term, "ghostty"):
		return true
	case strings.EqualFold(env("TERM_PROGRAM"), "ghostty"), env("GHOSTTY_RESOURCES_DIR") != "", env("KITTY_WINDOW_ID") != "":
		return true
	}
	return false
}