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