// SPDX-FileCopyrightText: © 2026 Jeffrey C. Ollie <jeff@ocjtech.us>
// SPDX-License-Identifier: MIT

//! Writing a z2d surface as sixel.

const std = @import("std");
const Allocator = std.mem.Allocator;
const Io = std.Io;
const z2d = @import("z2d");
const introducer = @import("introducer.zig");
const palette = @import("palette.zig");
const quantize = @import("quantize.zig");

pub const Dither = quantize.Dither;

pub const Options = struct {
    /// How many color registers to use, from 1 to `palette.max_registers`.
    /// 256 is what most terminals offer; a terminal with fewer registers
    /// than an image uses draws it in the wrong colors.
    max_colors: u16 = 256,
    /// How to make up for the colors lost when an image has more than
    /// `max_colors` of them. An image that fits is never dithered.
    dither: Dither = .floyd_steinberg,
    /// Pixels less opaque than this are not drawn, leaving whatever is
    /// underneath. Sixel has nothing between drawn and not drawn, so this is
    /// where a partly transparent pixel goes one way or the other. Those
    /// that are drawn are drawn in their own color, as if fully opaque.
    alpha_threshold: u8 = 128,
    /// Whether to write the raster attributes, `"1;1;W;H`, which give the
    /// pixel aspect ratio and the image's size. Without them a decoder
    /// knows the size only from what is drawn, so trailing transparent rows
    /// and columns are lost.
    raster_attributes: bool = true,
    /// Whether to wrap the data in its introducer and string terminator.
    /// Turn it off to get only the data between them, to wrap in something
    /// else.
    wrap: bool = true,
    /// Whether to write the 8-bit C1 forms of DCS and ST rather than
    /// `ESC P` and `ESC \`. Only for a channel known to be 8-bit clean and
    /// not UTF-8.
    c1: bool = false,
};

pub const Error = Allocator.Error || Io.Writer.Error;

/// Writes `surface` to `w` as sixel. `gpa` is for working memory only; the
/// output goes to `w`, which is not flushed.
///
/// What comes out is
///
///     ESC P 0 ; P2 ; 0 q  "1;1;W;H  #n;2;r;g;b ...  bands  ESC \
///
/// with `P2 = 1`, leave undrawn pixels alone, when the surface has
/// transparent pixels, and `0` otherwise; a register defined for each color
/// used; and then a band six pixels tall at a time, each color in the band
/// drawn left to right with runs of four or more run-length encoded, `$`
/// between colors and `-` between bands.
pub fn encode(gpa: Allocator, w: *Io.Writer, surface: *const z2d.Surface, options: Options) Error!void {
    const width: usize = @intCast(surface.getWidth());
    const height: usize = @intCast(surface.getHeight());
    const max_colors = std.math.clamp(options.max_colors, 1, palette.max_registers);

    const pixels = try gpa.alloc(quantize.Pixel, width * height);
    defer gpa.free(pixels);
    var any_transparent = false;
    for (0..height) |y| {
        for (0..width) |x| {
            const px = z2d.pixel.RGBA.fromPixel(surface.getPixel(@intCast(x), @intCast(y)).?).demultiply();
            const drawn = px.a >= @max(options.alpha_threshold, 1);
            any_transparent = any_transparent or !drawn;
            pixels[y * width + x] = .{ .color = .{ .r = px.r, .g = px.g, .b = px.b }, .opaque_ = drawn };
        }
    }

    var q = try quantize.quantize(gpa, pixels, width, max_colors, options.dither);
    defer q.deinit(gpa);

    if (options.wrap) {
        try w.writeAll(if (options.c1) "\x90" else "\x1bP");
        // P1 = 0 is DEC's advice for a new application, with the aspect
        // ratio in the raster attributes; without those, 9 asks for square
        // pixels itself.
        try w.print("{d};{d};0q", .{ @as(u8, if (options.raster_attributes) 0 else 9), @intFromBool(any_transparent) });
    }
    if (options.raster_attributes) try w.print("\"1;1;{d};{d}", .{ width, height });
    for (q.colors, 0..) |c, i| {
        try w.print("#{d};2;{d};{d};{d}", .{
            i,
            palette.channelToPercent(c.r),
            palette.channelToPercent(c.g),
            palette.channelToPercent(c.b),
        });
    }
    try bands(gpa, w, q.indices, width, height, q.colors.len);
    if (options.wrap) try w.writeAll(if (options.c1) "\x9c" else "\x1b\\");
}

/// Writes the pixel data, a band of six rows at a time.
fn bands(gpa: Allocator, w: *Io.Writer, indices: []const u16, width: usize, height: usize, ncolors: usize) Error!void {
    if (ncolors == 0) return;
    // The sixel each color would draw in each column of the band, and which
    // colors appear in it at all and over what columns, so that a band with
    // a few colors costs no more than those few.
    const sixels = try gpa.alloc(u8, ncolors * width);
    defer gpa.free(sixels);
    @memset(sixels, 0);
    const span = try gpa.alloc([2]usize, ncolors);
    defer gpa.free(span);
    // The last band each color was seen in, counting from one.
    const seen = try gpa.alloc(usize, ncolors);
    defer gpa.free(seen);
    @memset(seen, 0);
    var present: std.ArrayList(u16) = .empty;
    defer present.deinit(gpa);

    var top: usize = 0;
    var band: usize = 1;
    while (top < height) : ({
        top += 6;
        band += 1;
    }) {
        present.clearRetainingCapacity();
        for (top..@min(top + 6, height)) |y| {
            const bit: u8 = @as(u8, 1) << @intCast(y - top);
            for (indices[y * width ..][0..width], 0..) |c, x| {
                if (c == quantize.transparent) continue;
                if (seen[c] != band) {
                    seen[c] = band;
                    try present.append(gpa, c);
                    span[c] = .{ x, x };
                }
                sixels[@as(usize, c) * width + x] |= bit;
                span[c] = .{ @min(span[c][0], x), @max(span[c][1], x) };
            }
        }
        // In register order, which keeps the output the same from run to
        // run and is as good an order as any.
        std.mem.sortUnstable(u16, present.items, {}, std.sort.asc(u16));

        for (present.items, 0..) |c, n| {
            if (n > 0) try w.writeByte('$');
            try w.print("#{d}", .{c});
            const row = sixels[@as(usize, c) * width ..][0..width];
            // Leading empty columns have to be written to get across them;
            // trailing ones need not be.
            try runs(w, row[0 .. span[c][1] + 1]);
            @memset(row[span[c][0] .. span[c][1] + 1], 0);
        }
        if (top + 6 < height) try w.writeByte('-');
    }
}

/// Writes one color's sixels for a band, run-length encoded.
fn runs(w: *Io.Writer, row: []const u8) Io.Writer.Error!void {
    var i: usize = 0;
    while (i < row.len) {
        var j = i + 1;
        while (j < row.len and row[j] == row[i]) j += 1;
        const ch = '?' + row[i];
        const n = j - i;
        // `!3x` is no shorter than `xxx`.
        if (n > 3) {
            try w.print("!{d}{c}", .{ n, ch });
        } else {
            try w.splatByteAll(ch, n);
        }
        i = j;
    }
}

/// Writes `surface` as sixel into `buf`, and returns the part of it used.
/// `gpa` is for working memory only. Fails with `error.NoSpaceLeft` if the
/// output does not fit.
pub fn encodeBuffer(gpa: Allocator, buf: []u8, surface: *const z2d.Surface, options: Options) (Allocator.Error || error{NoSpaceLeft})![]u8 {
    var w: Io.Writer = .fixed(buf);
    encode(gpa, &w, surface, options) catch |err| switch (err) {
        error.WriteFailed => return error.NoSpaceLeft,
        error.OutOfMemory => return error.OutOfMemory,
    };
    return w.buffered();
}

/// Writes `surface` as sixel into memory allocated with `gpa`, which the
/// caller owns.
pub fn encodeAlloc(gpa: Allocator, surface: *const z2d.Surface, options: Options) Allocator.Error![]u8 {
    var out: Io.Writer.Allocating = .init(gpa);
    defer out.deinit();
    encode(gpa, &out.writer, surface, options) catch |err| switch (err) {
        error.WriteFailed => return error.OutOfMemory,
        error.OutOfMemory => return error.OutOfMemory,
    };
    return out.toOwnedSlice();
}

const testing = std.testing;

fn testSurface(gpa: Allocator, w: i32, h: i32, pixels: []const z2d.pixel.RGBA) !z2d.Surface {
    const s = try z2d.Surface.init(.image_surface_rgba, gpa, w, h);
    @memcpy(s.image_surface_rgba.buf, pixels);
    return s;
}

const R: z2d.pixel.RGBA = .{ .r = 255, .g = 0, .b = 0, .a = 255 };
const G: z2d.pixel.RGBA = .{ .r = 0, .g = 255, .b = 0, .a = 255 };
const T: z2d.pixel.RGBA = .{ .r = 0, .g = 0, .b = 0, .a = 0 };

test "a small image, byte for byte" {
    var s = try testSurface(testing.allocator, 5, 2, &.{
        R, R, R, R, R,
        G, G, G, G, T,
    });
    defer s.deinit(testing.allocator);
    const out = try encodeAlloc(testing.allocator, &s, .{});
    defer testing.allocator.free(out);
    try testing.expectEqualStrings(
        "\x1bP0;1;0q\"1;1;5;2#0;2;100;0;0#1;2;0;100;0#0!5@$#1!4A\x1b\\",
        out,
    );
}

test "bands, gaps and options" {
    var px: [3 * 8]z2d.pixel.RGBA = @splat(T);
    px[0] = R; // row 0
    px[2] = G;
    px[7 * 3 + 1] = R; // row 7, in the second band
    var s = try testSurface(testing.allocator, 3, 8, &px);
    defer s.deinit(testing.allocator);
    const out = try encodeAlloc(testing.allocator, &s, .{ .raster_attributes = false, .c1 = true });
    defer testing.allocator.free(out);
    try testing.expectEqualStrings("\x909;1;0q#0;2;100;0;0#1;2;0;100;0#0@$#1??@-#0?A\x9c", out);

    const bare = try encodeAlloc(testing.allocator, &s, .{ .wrap = false });
    defer testing.allocator.free(bare);
    try testing.expect(bare[0] == '"');
}

test "a buffer too small" {
    var s = try testSurface(testing.allocator, 1, 1, &.{R});
    defer s.deinit(testing.allocator);
    var buf: [8]u8 = undefined;
    try testing.expectError(error.NoSpaceLeft, encodeBuffer(testing.allocator, &buf, &s, .{}));
    var big: [128]u8 = undefined;
    const out = try encodeBuffer(testing.allocator, &big, &s, .{});
    try testing.expectEqualStrings("\x1bP0;0;0q\"1;1;1;1#0;2;100;0;0#0@\x1b\\", out);
}