1package kitty
 2
 3import (
 4	"compress/zlib"
 5	"fmt"
 6	"image"
 7	"image/color"
 8	"image/png"
 9	"io"
10)
11
12// Decoder is a decoder for the Kitty graphics protocol. It supports decoding
13// images in the 24-bit [RGB], 32-bit [RGBA], and [PNG] formats. It can also
14// decompress data using zlib.
15// The default format is 32-bit [RGBA].
16type Decoder struct {
17	// Uses zlib decompression.
18	Decompress bool
19
20	// Can be one of [RGB], [RGBA], or [PNG].
21	Format int
22
23	// Width of the image in pixels. This can be omitted if the image is [PNG]
24	// formatted.
25	Width int
26
27	// Height of the image in pixels. This can be omitted if the image is [PNG]
28	// formatted.
29	Height int
30}
31
32// Decode decodes the image data from r in the specified format.
33func (d *Decoder) Decode(r io.Reader) (image.Image, error) {
34	if d.Decompress {
35		zr, err := zlib.NewReader(r)
36		if err != nil {
37			return nil, fmt.Errorf("failed to create zlib reader: %w", err)
38		}
39
40		defer zr.Close() //nolint:errcheck
41		r = zr
42	}
43
44	if d.Format == 0 {
45		d.Format = RGBA
46	}
47
48	switch d.Format {
49	case RGBA, RGB:
50		return d.decodeRGBA(r, d.Format == RGBA)
51
52	case PNG:
53		return png.Decode(r)
54
55	default:
56		return nil, fmt.Errorf("unsupported format: %d", d.Format)
57	}
58}
59
60// decodeRGBA decodes the image data in 32-bit RGBA or 24-bit RGB formats.
61func (d *Decoder) decodeRGBA(r io.Reader, alpha bool) (image.Image, error) {
62	m := image.NewRGBA(image.Rect(0, 0, d.Width, d.Height))
63
64	var buf []byte
65	if alpha {
66		buf = make([]byte, 4)
67	} else {
68		buf = make([]byte, 3)
69	}
70
71	for y := range d.Height {
72		for x := range d.Width {
73			if _, err := io.ReadFull(r, buf[:]); err != nil {
74				return nil, fmt.Errorf("failed to read pixel data: %w", err)
75			}
76			if alpha {
77				m.SetRGBA(x, y, color.RGBA{buf[0], buf[1], buf[2], buf[3]})
78			} else {
79				m.SetRGBA(x, y, color.RGBA{buf[0], buf[1], buf[2], 0xff})
80			}
81		}
82	}
83
84	return m, nil
85}