1package kitty
 2
 3import (
 4	"compress/zlib"
 5	"fmt"
 6	"image"
 7	"image/png"
 8	"io"
 9)
10
11// Encoder is an encoder for the Kitty graphics protocol. It supports encoding
12// images in the 24-bit [RGB], 32-bit [RGBA], and [PNG] formats, and
13// compressing the data using zlib.
14// The default format is 32-bit [RGBA].
15type Encoder struct {
16	// Uses zlib compression.
17	Compress bool
18
19	// Can be one of [RGBA], [RGB], or [PNG].
20	Format int
21}
22
23// Encode encodes the image data in the specified format and writes it to w.
24func (e *Encoder) Encode(w io.Writer, m image.Image) error {
25	if m == nil {
26		return nil
27	}
28
29	if e.Compress {
30		zw := zlib.NewWriter(w)
31		defer zw.Close() //nolint:errcheck
32		w = zw
33	}
34
35	if e.Format == 0 {
36		e.Format = RGBA
37	}
38
39	switch e.Format {
40	case RGBA, RGB:
41		bounds := m.Bounds()
42		for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
43			for x := bounds.Min.X; x < bounds.Max.X; x++ {
44				r, g, b, a := m.At(x, y).RGBA()
45				switch e.Format {
46				case RGBA:
47					w.Write([]byte{byte(r >> 8), byte(g >> 8), byte(b >> 8), byte(a >> 8)}) //nolint:errcheck
48				case RGB:
49					w.Write([]byte{byte(r >> 8), byte(g >> 8), byte(b >> 8)}) //nolint:errcheck
50				}
51			}
52		}
53
54	case PNG:
55		if err := png.Encode(w, m); err != nil {
56			return fmt.Errorf("failed to encode PNG: %w", err)
57		}
58
59	default:
60		return fmt.Errorf("unsupported format: %d", e.Format)
61	}
62
63	return nil
64}