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/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// fitImage resizes the image to fit within the specified dimensions in
72// terminal cells, maintaining the aspect ratio.
73func fitImage(id string, img image.Image, cs CellSize, cols, rows int) image.Image {
74 if img == nil {
75 return nil
76 }
77
78 key := imageKey{id: id, cols: cols, rows: rows}
79
80 cachedMutex.RLock()
81 cached, ok := cachedImages[key]
82 cachedMutex.RUnlock()
83 if ok {
84 return cached.img
85 }
86
87 if cs.Width == 0 || cs.Height == 0 {
88 return img
89 }
90
91 maxWidth := cols * cs.Width
92 maxHeight := rows * cs.Height
93
94 img = imaging.Fit(img, maxWidth, maxHeight, imaging.Lanczos)
95
96 cachedMutex.Lock()
97 cachedImages[key] = cachedImage{
98 img: img,
99 cols: cols,
100 rows: rows,
101 }
102 cachedMutex.Unlock()
103
104 return img
105}
106
107// HasTransmitted checks if the image with the given ID has already been
108// transmitted to the terminal.
109func HasTransmitted(id string, cols, rows int) bool {
110 key := imageKey{id: id, cols: cols, rows: rows}
111
112 cachedMutex.RLock()
113 _, ok := cachedImages[key]
114 cachedMutex.RUnlock()
115 return ok
116}
117
118// Transmit transmits the image data to the terminal if needed. This is used to
119// cache the image on the terminal for later rendering.
120func (e Encoding) Transmit(id string, img image.Image, cs CellSize, cols, rows int, tmux bool) tea.Cmd {
121 if img == nil {
122 return nil
123 }
124
125 key := imageKey{id: id, cols: cols, rows: rows}
126
127 cachedMutex.RLock()
128 _, ok := cachedImages[key]
129 cachedMutex.RUnlock()
130 if ok {
131 return nil
132 }
133
134 cmd := func() tea.Msg {
135 if e != EncodingKitty {
136 cachedMutex.Lock()
137 cachedImages[key] = cachedImage{
138 img: img,
139 cols: cols,
140 rows: rows,
141 }
142 cachedMutex.Unlock()
143 return TransmittedMsg{ID: key.ID()}
144 }
145
146 var buf bytes.Buffer
147 img := fitImage(id, img, cs, cols, rows)
148 bounds := img.Bounds()
149 imgWidth := bounds.Dx()
150 imgHeight := bounds.Dy()
151 imgID := int(key.Hash())
152 if err := kitty.EncodeGraphics(&buf, img, &kitty.Options{
153 ID: imgID,
154 Action: kitty.TransmitAndPut,
155 Transmission: kitty.Direct,
156 Format: kitty.RGBA,
157 ImageWidth: imgWidth,
158 ImageHeight: imgHeight,
159 Columns: cols,
160 Rows: rows,
161 VirtualPlacement: true,
162 Quite: 1,
163 Chunk: true,
164 ChunkFormatter: func(chunk string) string {
165 if tmux {
166 return ansi.TmuxPassthrough(chunk)
167 }
168 return chunk
169 },
170 }); err != nil {
171 slog.Error("Failed to encode image for kitty graphics", "err", err)
172 return uiutil.InfoMsg{
173 Type: uiutil.InfoTypeError,
174 Msg: "failed to encode image",
175 }
176 }
177
178 return tea.RawMsg{Msg: buf.String()}
179 }
180
181 return cmd
182}
183
184// Render renders the given image within the specified dimensions using the
185// specified encoding.
186func (e Encoding) Render(id string, cols, rows int) string {
187 key := imageKey{id: id, cols: cols, rows: rows}
188 cachedMutex.RLock()
189 cached, ok := cachedImages[key]
190 cachedMutex.RUnlock()
191 if !ok {
192 return ""
193 }
194
195 img := cached.img
196
197 switch e {
198 case EncodingBlocks:
199 canvas := paintbrush.New()
200 canvas.SetImage(img)
201 canvas.SetWidth(cols)
202 canvas.SetHeight(rows)
203 canvas.Weights = map[rune]float64{
204 '': .95,
205 '': .95,
206 '▁': .9,
207 '▂': .9,
208 '▃': .9,
209 '▄': .9,
210 '▅': .9,
211 '▆': .85,
212 '█': .85,
213 '▊': .95,
214 '▋': .95,
215 '▌': .95,
216 '▍': .95,
217 '▎': .95,
218 '▏': .95,
219 '●': .95,
220 '◀': .95,
221 '▲': .95,
222 '▶': .95,
223 '▼': .9,
224 '○': .8,
225 '◉': .95,
226 '◧': .9,
227 '◨': .9,
228 '◩': .9,
229 '◪': .9,
230 }
231 canvas.Paint()
232 return strings.TrimSpace(canvas.GetResult())
233 case EncodingKitty:
234 // Build Kitty graphics unicode place holders
235 var fg color.Color
236 var extra int
237 var r, g, b int
238 hashedID := key.Hash()
239 id := int(hashedID)
240 extra, r, g, b = id>>24&0xff, id>>16&0xff, id>>8&0xff, id&0xff
241
242 if id <= 255 {
243 fg = ansi.IndexedColor(b)
244 } else {
245 fg = color.RGBA{
246 R: uint8(r), //nolint:gosec
247 G: uint8(g), //nolint:gosec
248 B: uint8(b), //nolint:gosec
249 A: 0xff,
250 }
251 }
252
253 fgStyle := ansi.NewStyle().ForegroundColor(fg).String()
254
255 var buf bytes.Buffer
256 for y := range rows {
257 // As an optimization, we only write the fg color sequence id, and
258 // column-row data once on the first cell. The terminal will handle
259 // the rest.
260 buf.WriteString(fgStyle)
261 buf.WriteRune(kitty.Placeholder)
262 buf.WriteRune(kitty.Diacritic(y))
263 buf.WriteRune(kitty.Diacritic(0))
264 if extra > 0 {
265 buf.WriteRune(kitty.Diacritic(extra))
266 }
267 for x := 1; x < cols; x++ {
268 buf.WriteString(fgStyle)
269 buf.WriteRune(kitty.Placeholder)
270 }
271 if y < rows-1 {
272 buf.WriteByte('\n')
273 }
274 }
275
276 return buf.String()
277
278 default:
279 return ""
280 }
281}