Mainline/engine/display/backends/sixel.py

"""
Sixel graphics display backend - renders to sixel graphics in terminal.
"""

import time

from engine.display.renderer import get_default_font_path, parse_ansi


def _encode_sixel(image) -> str:
    """Encode a PIL Image to sixel format (pure Python)."""
    img = image.convert("RGBA")
    width, height = img.size
    pixels = img.load()

    palette = []
    pixel_palette_idx = {}

    def get_color_idx(r, g, b, a):
        if a < 128:
            return -1
        key = (r // 32, g // 32, b // 32)
        if key not in pixel_palette_idx:
            idx = len(palette)
            if idx < 256:
                palette.append((r, g, b))
                pixel_palette_idx[key] = idx
        return pixel_palette_idx.get(key, 0)

    for y in range(height):
        for x in range(width):
            r, g, b, a = pixels[x, y]
            get_color_idx(r, g, b, a)

    if not palette:
        return ""

    if len(palette) == 1:
        palette = [palette[0], (0, 0, 0)]

    sixel_data = []
    sixel_data.append(
        f'"{"".join(f"#{i};2;{r};{g};{b}" for i, (r, g, b) in enumerate(palette))}'
    )

    for x in range(width):
        col_data = []
        for y in range(0, height, 6):
            bits = 0
            color_idx = -1
            for dy in range(6):
                if y + dy < height:
                    r, g, b, a = pixels[x, y + dy]
                    if a >= 128:
                        bits |= 1 << dy
                        idx = get_color_idx(r, g, b, a)
                        if color_idx == -1:
                            color_idx = idx
                        elif color_idx != idx:
                            color_idx = -2

            if color_idx >= 0:
                col_data.append(
                    chr(63 + color_idx) + chr(63 + bits)
                    if bits
                    else chr(63 + color_idx) + "?"
                )
            elif color_idx == -2:
                pass

        if col_data:
            sixel_data.append("".join(col_data) + "$")
        else:
            sixel_data.append("-" if x < width - 1 else "$")

    sixel_data.append("\x1b\\")

    return "\x1bPq" + "".join(sixel_data)


class SixelDisplay:
    """Sixel graphics display backend - renders to sixel graphics in terminal."""

    width: int = 80
    height: int = 24

    def __init__(self, cell_width: int = 9, cell_height: int = 16):
        self.width = 80
        self.height = 24
        self.cell_width = cell_width
        self.cell_height = cell_height
        self._initialized = False
        self._font_path = None

    def _get_font_path(self) -> str | None:
        """Get font path from env or detect common locations."""
        import os

        if self._font_path:
            return self._font_path

        env_font = os.environ.get("MAINLINE_SIXEL_FONT")
        if env_font and os.path.exists(env_font):
            self._font_path = env_font
            return env_font

        font_path = get_default_font_path()
        if font_path:
            self._font_path = font_path

        return self._font_path

    def init(self, width: int, height: int, reuse: bool = False) -> None:
        """Initialize display with dimensions.

        Args:
            width: Terminal width in characters
            height: Terminal height in rows
            reuse: Ignored for SixelDisplay
        """
        self.width = width
        self.height = height
        self._initialized = True

    def show(self, buffer: list[str]) -> None:
        import sys

        t0 = time.perf_counter()

        img_width = self.width * self.cell_width
        img_height = self.height * self.cell_height

        try:
            from PIL import Image, ImageDraw, ImageFont
        except ImportError:
            return

        img = Image.new("RGBA", (img_width, img_height), (0, 0, 0, 255))
        draw = ImageDraw.Draw(img)

        font_path = self._get_font_path()
        font = None
        if font_path:
            try:
                font = ImageFont.truetype(font_path, self.cell_height - 2)
            except Exception:
                font = None

        if font is None:
            try:
                font = ImageFont.load_default()
            except Exception:
                font = None

        for row_idx, line in enumerate(buffer[: self.height]):
            if row_idx >= self.height:
                break

            tokens = parse_ansi(line)
            x_pos = 0
            y_pos = row_idx * self.cell_height

            for text, fg, bg, bold in tokens:
                if not text:
                    continue

                if bg != (0, 0, 0):
                    bbox = draw.textbbox((x_pos, y_pos), text, font=font)
                    draw.rectangle(bbox, fill=(*bg, 255))

                if bold and font:
                    draw.text((x_pos - 1, y_pos - 1), text, fill=(*fg, 255), font=font)

                draw.text((x_pos, y_pos), text, fill=(*fg, 255), font=font)

                if font:
                    x_pos += draw.textlength(text, font=font)

        sixel = _encode_sixel(img)

        sys.stdout.buffer.write(sixel.encode("utf-8"))
        sys.stdout.flush()

        elapsed_ms = (time.perf_counter() - t0) * 1000

        from engine.display import get_monitor

        monitor = get_monitor()
        if monitor:
            chars_in = sum(len(line) for line in buffer)
            monitor.record_effect("sixel_display", elapsed_ms, chars_in, chars_in)

    def clear(self) -> None:
        import sys

        sys.stdout.buffer.write(b"\x1b[2J\x1b[H")
        sys.stdout.flush()

    def cleanup(self) -> None:
        pass