bubbles/image/image.go

package image

import (
	"bytes"
	"fmt"
	"image"
	"image/color"
	"log"
	"os"
	"strings"
	"sync/atomic"

	tea "github.com/charmbracelet/bubbletea/v2"
	"github.com/charmbracelet/x/ansi"
	"github.com/charmbracelet/x/ansi/kitty"
	"github.com/charmbracelet/x/input"
)

// number is a global number used to generate unique image numbers.
var number int32

// nextImageNumber returns the next unique image number.
func nextNumber() int32 {
	return atomic.AddInt32(&number, 1)
}

// Protocol is the terminal graphics protocol used to render the image.
type Protocol byte

// Graphic protocol constants.
const (
	HalfBlocks Protocol = iota + 1
	Sixel
	ITerm2
	Kitty
)

// Model represents a terminal graphics image.
type Model struct {
	// The protocol used
	Protocol Protocol
	// The area covering the image in cells
	area image.Rectangle
	// The image data (exclusive with file)
	m image.Image
	// The image file path (exclusive with m)
	file string

	// The image options
	opts kitty.Options

	// The image unique id. A non-zero indicates the image was transmitted successfully.
	id int
	// The image number
	num int

	// Whether the image was transmitted
	didTransmit bool

	// The terminal width and height
	w, h int
}

func newModel(area image.Rectangle) (m Model) {
	// We always use virtual placement for images
	m.opts.VirtualPlacement = true
	// Always chunk the image
	m.opts.Chunk = true
	// Transmit and put/display the image
	m.opts.Action = kitty.TransmitAndPut

	num := int(nextNumber())
	m.num = num
	m.opts.Number = num

	m.SetArea(area)

	return
}

// NewLocal creates a new image model from a local file.
func NewLocal(file string, area image.Rectangle) (m Model, err error) {
	m = newModel(area)
	m.file = file
	m.area = area
	// TODO: Fix me! This currently doesn't work with PNG
	// ext := filepath.Ext(file)
	// if strings.Contains(ext, "png") {
	// 	// We're done here, there's no need to decode the image.
	// 	m.opts.Format = kitty.PNG
	// 	m.opts.File = file
	// 	return
	// }

	f, err := os.Open(file)
	if err != nil {
		return m, fmt.Errorf("could not open image file: %w", err)
	}

	defer f.Close() //nolint:errcheck

	im, mtyp, err := image.Decode(f)
	if err != nil {
		return m, fmt.Errorf("could not decode image: %w", err)
	}

	m.m = im

	// Use a temporary file to store the image data
	m.opts.Transmission = kitty.TempFile

	// TODO: Enable compression
	// m.opts.Compression = kitty.Zlib

	// Set the image size
	bounds := im.Bounds()
	m.opts.ImageWidth = bounds.Dx()
	m.opts.ImageHeight = bounds.Dy()

	// Optimize for JPEG images and alpha transparency
	switch mtyp {
	case "jpeg":
		m.opts.Format = kitty.RGB
	default:
		m.opts.Format = kitty.RGBA
	}

	return
}

// New creates a new image model given an image and an area in cells
func New(im image.Image, area image.Rectangle) Model {
	m := newModel(area)
	m.opts.Transmission = kitty.Direct
	// m.opts.Compression = kitty.Zlib
	m.opts.Format = kitty.RGBA
	m.m = im
	// Set the image size
	bounds := im.Bounds()
	m.opts.ImageWidth = bounds.Dx()
	m.opts.ImageHeight = bounds.Dy()
	m.SetArea(area)
	return m
}

// ID returns the image id unique with respect to the terminal.
func (m Model) ID() int {
	return m.id
}

// Number returns the image number unique with respect to the library.
func (m Model) Number() int {
	return m.num
}

// SetArea sets the image area in cells.
func (m *Model) SetArea(area image.Rectangle) {
	m.area = area
	m.opts.Columns = m.area.Dx()
	m.opts.Rows = m.area.Dy()
	m.didTransmit = false
}

// Area returns the image area in cells.
func (m Model) Area() image.Rectangle {
	return m.area
}

// transmit is a command that transmits the image to the terminal.
func (m *Model) transmit() tea.Msg {
	var seq bytes.Buffer
	if err := ansi.WriteKittyGraphics(&seq, m.m, &m.opts); err != nil {
		// TODO: Error handling
		return nil
	}

	m.didTransmit = true
	return tea.RawMsg{Msg: seq.String()}
}

// Init initializes the image model.
func (m Model) Init() (tea.Model, tea.Cmd) {
	return m, tea.Batch(
	// TODO: Query support
	)
}

// Update updates the image model.
func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	var cmds []tea.Cmd
	switch msg := msg.(type) {
	case tea.KeyMsg:
		m.didTransmit = false
	case tea.WindowSizeMsg:
		m.w, m.h = msg.Width, msg.Height
	case input.KittyGraphicsEvent:
		if msg.Options.Number == m.num &&
			msg.Options.ID > 0 &&
			bytes.Equal(msg.Payload, []byte("OK")) {
			// Store the actual image id
			m.id = msg.Options.ID
		}
	}

	if !m.didTransmit {
		cmds = append(cmds, m.transmit)
		m.didTransmit = true
	}

	return m, tea.Batch(cmds...)
}

// View returns a string representation to render the image.
func (m Model) View() string {
	if m.id == 0 {
		// TODO: Maybe use a spinner?
		return "Loading image..."
	}

	log.Printf("area: %v", m.area)

	// Build Kitty graphics unicode place holders
	var fgSeq string
	var extra int
	var r, g, b int
	extra, r, g, b = m.id>>24&0xff, m.id>>16&0xff, m.id>>8&0xff, m.id&0xff

	if r == 0 && g == 0 {
		fgSeq = ansi.Style{}.ForegroundColor(ansi.ExtendedColor(b)).String() //nolint:gosec
	} else {
		fgSeq = ansi.Style{}.ForegroundColor(color.RGBA{
			R: uint8(r), //nolint:gosec
			G: uint8(g), //nolint:gosec
			B: uint8(b), //nolint:gosec
			A: 0xff,
		}).String()
	}

	var s strings.Builder
	width := min(m.area.Dx(), m.w)
	height := min(m.area.Dy(), m.h)
	s.WriteString(ansi.ResetStyle)

	for y := 0; y < height; y++ {
		// As an optimization, we only write the fg color sequence id, and
		// column-row data once on the first cell. The terminal will handle
		// the rest.
		s.WriteString(fgSeq)
		s.WriteRune(kitty.Placeholder)
		s.WriteRune(kitty.Diacritic(y))
		s.WriteRune(kitty.Diacritic(0))
		if extra > 0 {
			s.WriteRune(kitty.Diacritic(extra))
		}

		for x := 1; x < width; x++ {
			s.WriteRune(kitty.Placeholder)
		}

		s.WriteString(ansi.ResetStyle)
		if y != m.area.Dy()-1 {
			s.WriteByte('\n')
		}
	}

	return s.String()
}

// Rect returns a rectangle from the given x, y, width, and height.
func Rect(x, y, w, h int) image.Rectangle {
	return image.Rect(x, y, x+w, y+h)
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}