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