image.go

  1package image
  2
  3import (
  4	"bytes"
  5	"fmt"
  6	"hash/fnv"
  7	"image"
  8	"image/color"
  9	"io"
 10	"log/slog"
 11	"strings"
 12	"sync"
 13
 14	tea "charm.land/bubbletea/v2"
 15	"github.com/charmbracelet/crush/internal/uiutil"
 16	"github.com/charmbracelet/x/ansi"
 17	"github.com/charmbracelet/x/ansi/kitty"
 18	"github.com/charmbracelet/x/mosaic"
 19	"github.com/disintegration/imaging"
 20)
 21
 22// Capabilities represents the capabilities of displaying images on the
 23// terminal.
 24type Capabilities struct {
 25	// Columns is the number of character columns in the terminal.
 26	Columns int
 27	// Rows is the number of character rows in the terminal.
 28	Rows int
 29	// PixelWidth is the width of the terminal in pixels.
 30	PixelWidth int
 31	// PixelHeight is the height of the terminal in pixels.
 32	PixelHeight int
 33	// SupportsKittyGraphics indicates whether the terminal supports the Kitty
 34	// graphics protocol.
 35	SupportsKittyGraphics bool
 36}
 37
 38// CellSize returns the size of a single terminal cell in pixels.
 39func (c Capabilities) CellSize() CellSize {
 40	return CalculateCellSize(c.PixelWidth, c.PixelHeight, c.Columns, c.Rows)
 41}
 42
 43// CalculateCellSize calculates the size of a single terminal cell in pixels
 44// based on the terminal's pixel dimensions and character dimensions.
 45func CalculateCellSize(pixelWidth, pixelHeight, charWidth, charHeight int) CellSize {
 46	if charWidth == 0 || charHeight == 0 {
 47		return CellSize{}
 48	}
 49
 50	return CellSize{
 51		Width:  pixelWidth / charWidth,
 52		Height: pixelHeight / charHeight,
 53	}
 54}
 55
 56// RequestCapabilities is a [tea.Cmd] that requests the terminal to report
 57// its image related capabilities to the program.
 58func RequestCapabilities() tea.Cmd {
 59	return tea.Raw(
 60		ansi.WindowOp(14) + // Window size in pixels
 61			// ID 31 is just a random ID used to detect Kitty graphics support.
 62			ansi.KittyGraphics([]byte("AAAA"), "i=31", "s=1", "v=1", "a=q", "t=d", "f=24"),
 63	)
 64}
 65
 66// TransmittedMsg is a message indicating that an image has been transmitted to
 67// the terminal.
 68type TransmittedMsg struct {
 69	ID string
 70}
 71
 72// Encoding represents the encoding format of the image.
 73type Encoding byte
 74
 75// Image encodings.
 76const (
 77	EncodingBlocks Encoding = iota
 78	EncodingKitty
 79)
 80
 81type imageKey struct {
 82	id   string
 83	cols int
 84	rows int
 85}
 86
 87// Hash returns a hash value for the image key.
 88// This uses FNV-32a for simplicity and speed.
 89func (k imageKey) Hash() uint32 {
 90	h := fnv.New32a()
 91	_, _ = io.WriteString(h, k.ID())
 92	return h.Sum32()
 93}
 94
 95// ID returns a unique string representation of the image key.
 96func (k imageKey) ID() string {
 97	return fmt.Sprintf("%s-%dx%d", k.id, k.cols, k.rows)
 98}
 99
100// CellSize represents the size of a single terminal cell in pixels.
101type CellSize struct {
102	Width, Height int
103}
104
105type cachedImage struct {
106	img        image.Image
107	cols, rows int
108}
109
110var (
111	cachedImages = map[imageKey]cachedImage{}
112	cachedMutex  sync.RWMutex
113)
114
115// fitImage resizes the image to fit within the specified dimensions in
116// terminal cells, maintaining the aspect ratio.
117func fitImage(id string, img image.Image, cs CellSize, cols, rows int) image.Image {
118	if img == nil {
119		return nil
120	}
121
122	key := imageKey{id: id, cols: cols, rows: rows}
123
124	cachedMutex.RLock()
125	cached, ok := cachedImages[key]
126	cachedMutex.RUnlock()
127	if ok {
128		return cached.img
129	}
130
131	if cs.Width == 0 || cs.Height == 0 {
132		return img
133	}
134
135	maxWidth := cols * cs.Width
136	maxHeight := rows * cs.Height
137
138	img = imaging.Fit(img, maxWidth, maxHeight, imaging.Lanczos)
139
140	cachedMutex.Lock()
141	cachedImages[key] = cachedImage{
142		img:  img,
143		cols: cols,
144		rows: rows,
145	}
146	cachedMutex.Unlock()
147
148	return img
149}
150
151// HasTransmitted checks if the image with the given ID has already been
152// transmitted to the terminal.
153func HasTransmitted(id string, cols, rows int) bool {
154	key := imageKey{id: id, cols: cols, rows: rows}
155
156	cachedMutex.RLock()
157	_, ok := cachedImages[key]
158	cachedMutex.RUnlock()
159	return ok
160}
161
162// Transmit transmits the image data to the terminal if needed. This is used to
163// cache the image on the terminal for later rendering.
164func (e Encoding) Transmit(id string, img image.Image, cs CellSize, cols, rows int) tea.Cmd {
165	if img == nil {
166		return nil
167	}
168
169	key := imageKey{id: id, cols: cols, rows: rows}
170
171	cachedMutex.RLock()
172	_, ok := cachedImages[key]
173	cachedMutex.RUnlock()
174	if ok {
175		return nil
176	}
177
178	cmd := func() tea.Msg {
179		if e != EncodingKitty {
180			cachedMutex.Lock()
181			cachedImages[key] = cachedImage{
182				img:  img,
183				cols: cols,
184				rows: rows,
185			}
186			cachedMutex.Unlock()
187			return TransmittedMsg{ID: key.ID()}
188		}
189
190		var buf bytes.Buffer
191		img := fitImage(id, img, cs, cols, rows)
192		bounds := img.Bounds()
193		imgWidth := bounds.Dx()
194		imgHeight := bounds.Dy()
195
196		imgID := int(key.Hash())
197		if err := kitty.EncodeGraphics(&buf, img, &kitty.Options{
198			ID:               imgID,
199			Action:           kitty.TransmitAndPut,
200			Transmission:     kitty.Direct,
201			Format:           kitty.RGBA,
202			ImageWidth:       imgWidth,
203			ImageHeight:      imgHeight,
204			Columns:          cols,
205			Rows:             rows,
206			VirtualPlacement: true,
207			Quite:            1,
208		}); err != nil {
209			slog.Error("failed to encode image for kitty graphics", "err", err)
210			return uiutil.ReportError(fmt.Errorf("failed to encode image"))
211		}
212
213		return tea.RawMsg{Msg: buf.String()}
214	}
215
216	return cmd
217}
218
219// Render renders the given image within the specified dimensions using the
220// specified encoding.
221func (e Encoding) Render(id string, cols, rows int) string {
222	key := imageKey{id: id, cols: cols, rows: rows}
223	cachedMutex.RLock()
224	cached, ok := cachedImages[key]
225	cachedMutex.RUnlock()
226	if !ok {
227		return ""
228	}
229
230	img := cached.img
231
232	switch e {
233	case EncodingBlocks:
234		m := mosaic.New().Width(cols).Height(rows).Scale(1)
235		return strings.TrimSpace(m.Render(img))
236	case EncodingKitty:
237		// Build Kitty graphics unicode place holders
238		var fg color.Color
239		var extra int
240		var r, g, b int
241		hashedID := key.Hash()
242		id := int(hashedID)
243		extra, r, g, b = id>>24&0xff, id>>16&0xff, id>>8&0xff, id&0xff
244
245		if id <= 255 {
246			fg = ansi.IndexedColor(b)
247		} else {
248			fg = color.RGBA{
249				R: uint8(r), //nolint:gosec
250				G: uint8(g), //nolint:gosec
251				B: uint8(b), //nolint:gosec
252				A: 0xff,
253			}
254		}
255
256		fgStyle := ansi.NewStyle().ForegroundColor(fg).String()
257
258		var buf bytes.Buffer
259		for y := range rows {
260			// As an optimization, we only write the fg color sequence id, and
261			// column-row data once on the first cell. The terminal will handle
262			// the rest.
263			buf.WriteString(fgStyle)
264			buf.WriteRune(kitty.Placeholder)
265			buf.WriteRune(kitty.Diacritic(y))
266			buf.WriteRune(kitty.Diacritic(0))
267			if extra > 0 {
268				buf.WriteRune(kitty.Diacritic(extra))
269			}
270			for x := 1; x < cols; x++ {
271				buf.WriteString(fgStyle)
272				buf.WriteRune(kitty.Placeholder)
273			}
274			if y < rows-1 {
275				buf.WriteByte('\n')
276			}
277		}
278
279		return buf.String()
280
281	default:
282		return ""
283	}
284}