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

//! Picking sixel images out of a stream, and passing everything else on.

const std = @import("std");
const Allocator = std.mem.Allocator;
const Io = std.Io;
const Canvas = @import("Canvas.zig");
const Image = @import("Image.zig");
const Parser = @import("Parser.zig");
const introducer = @import("introducer.zig");
const palette = @import("palette.zig");

pub const Options = struct {
    limits: Canvas.Limits = .{},
    /// See `Parser.Options.registers`.
    registers: u16 = palette.max_registers,
    /// Whether to recognize the 8-bit C1 forms of DCS (`0x90`) and ST
    /// (`0x9C`). Off by default, because in UTF-8 text those bytes are
    /// continuation bytes, and taking one for a DCS would swallow text.
    c1: bool = false,
};

/// How large the reader's buffer must be: enough to hold a whole introducer
/// while deciding whether it is one, and one byte more to read into.
pub const min_buffer_len = introducer.max_len + 1;

pub const Error = Allocator.Error || Io.Reader.Error || Io.Writer.Error || error{
    /// The reader's buffer is smaller than `min_buffer_len`.
    ReaderBufferTooSmall,
};

/// Copies `reader` to `writer` until the reader ends, except for sixel
/// images, which are decoded and handed to `handler` instead of being
/// written.
///
/// `handler` is anything with a method
///
///     fn image(self, image: sixel.Image) !void
///
/// which takes ownership of the image, even if it returns an error, and must
/// `deinit` it with `gpa`. Its errors are returned from `relay` as they are.
///
/// Everything that is not a sixel image is written through untouched and in
/// order, including other device control strings and escape sequences: only
/// `ESC P`, then digits and semicolons, then `q` starts an image. The image
/// ends at ST, which is consumed with it; at CAN or SUB, which are consumed
/// too, since they abandon the string; or at any other escape sequence or the
/// end of the stream, which are not, and which a handler can tell apart by
/// `Image.termination`. An image that draws nothing and gives no size is not
/// handed over at all, having no pixels.
///
/// Nothing is held back that does not need to be. Text goes to the writer as
/// soon as it has been read, and the writer is flushed before every read
/// that might block, so what has arrived reaches the writer's destination
/// without waiting for anything that has not. The only bytes ever held are
/// those of an escape sequence that has started arriving but not yet arrived
/// far enough to tell whether it is a sixel introducer.
///
/// `reader` must have a buffer of at least `min_buffer_len` bytes.
pub fn relay(
    gpa: Allocator,
    reader: *Io.Reader,
    writer: *Io.Writer,
    handler: anytype,
    options: Options,
) !void {
    if (reader.buffer.len < min_buffer_len) return error.ReaderBufferTooSmall;

    while (true) {
        const buf = reader.buffered();
        if (buf.len == 0) {
            if (!try fillMore(reader, writer)) return writer.flush();
            continue;
        }

        const start = for (buf, 0..) |b, i| {
            if (introducer.isStart(b, options.c1)) break i;
        } else buf.len;
        if (start > 0) {
            try writer.writeAll(buf[0..start]);
            reader.toss(start);
            continue;
        }

        switch (introducer.match(buf, options.c1)) {
            .not_sixel => {
                // Only the first byte is known not to start an image. The
                // rest is examined afresh, since `ESC ESC P q` is an escape
                // and then a sixel.
                try writer.writeAll(buf[0..1]);
                reader.toss(1);
            },
            .need_more => if (!try fillMore(reader, writer)) {
                // The stream ended partway into an escape sequence, which
                // is now known not to be an image.
                try writer.writeAll(reader.buffered());
                reader.toss(reader.bufferedLen());
                return writer.flush();
            },
            .sixel => |m| {
                reader.toss(m.len);
                if (try image(gpa, reader, writer, m.params, options)) |img| {
                    try handler.image(img);
                }
            },
        }
    }
}

/// Reads more after flushing what is already written, so that nothing
/// waits on a read that may block. Returns false at the end of the stream.
fn fillMore(reader: *Io.Reader, writer: *Io.Writer) Error!bool {
    try writer.flush();
    reader.fillMore() catch |err| switch (err) {
        error.EndOfStream => return false,
        else => |e| return e,
    };
    return true;
}

/// Decodes the image whose introducer has just been consumed, up to and
/// including its terminator if that is ST, CAN or SUB.
fn image(
    gpa: Allocator,
    reader: *Io.Reader,
    writer: *Io.Writer,
    params: introducer.Params,
    options: Options,
) Error!?Image {
    var parser: Parser = .init(gpa, params, .{
        .limits = options.limits,
        .registers = options.registers,
        .c1 = options.c1,
    });
    defer parser.deinit();

    const termination: Image.Termination = while (true) {
        const buf = reader.buffered();
        if (buf.len == 0) {
            if (!try fillMore(reader, writer)) break .end_of_data;
            continue;
        }
        const n = try parser.feed(buf);
        reader.toss(n);
        if (n == buf.len) continue;

        switch (buf[n]) {
            introducer.esc => {
                // Whether this is ST or the start of something else depends
                // on the next byte.
                if (reader.bufferedLen() < 2 and !try fillMore(reader, writer)) break .interrupted;
                if (reader.bufferedLen() < 2) continue;
                if (reader.buffered()[1] == '\\') {
                    reader.toss(2);
                    break .st;
                }
                break .interrupted;
            },
            introducer.st_c1 => {
                reader.toss(1);
                break .st;
            },
            0x18, 0x1a => {
                reader.toss(1);
                break .aborted;
            },
            else => break .interrupted,
        }
    };
    return parser.finish(termination);
}

const testing = std.testing;

/// Collects what `relay` hands over.
const Collector = struct {
    images: std.ArrayList(Image) = .empty,

    fn image(self: *Collector, img: Image) !void {
        try self.images.append(testing.allocator, img);
    }

    fn deinit(self: *Collector) void {
        for (self.images.items) |*img| img.deinit(testing.allocator);
        self.images.deinit(testing.allocator);
    }
};

fn run(input: []const u8, options: Options, collector: *Collector) ![]u8 {
    var reader: Io.Reader = .fixed(input);
    var out: Io.Writer.Allocating = .init(testing.allocator);
    errdefer out.deinit();
    // A fixed reader's buffer is the input; give a short one room to look
    // ahead by copying it into one that is large enough.
    var big: [256]u8 = undefined;
    if (input.len < min_buffer_len) {
        @memcpy(big[0..input.len], input);
        reader = .{ .vtable = reader.vtable, .buffer = &big, .seek = 0, .end = input.len };
    }
    try relay(testing.allocator, &reader, &out.writer, collector, options);
    return out.toOwnedSlice();
}

test "text passes through and images come out" {
    var c: Collector = .{};
    defer c.deinit();
    const out = try run(
        "before\x1bPq#1~~\x1b\\middle\x1bP0;1;0q\"1;1;3;3#2~\x1b\\after",
        .{},
        &c,
    );
    defer testing.allocator.free(out);
    try testing.expectEqualStrings("beforemiddleafter", out);
    try testing.expectEqual(@as(usize, 2), c.images.items.len);
    try testing.expectEqual(@as(u32, 2), c.images.items[0].width());
    try testing.expectEqual(Image.Termination.st, c.images.items[0].termination);
    try testing.expect(c.images.items[1].transparent);
    try testing.expectEqual(@as(u32, 3), c.images.items[1].width());
}

test "other sequences pass through untouched" {
    var c: Collector = .{};
    defer c.deinit();
    const input = "\x1b[1mbold\x1b[0m\x1bP$q\"p\x1b\\\x1bP+q544e\x1b\\\x1b\x1bPq~\x1b\\\x1bP1$r\x1b\\\x1b";
    const out = try run(input, .{}, &c);
    defer testing.allocator.free(out);
    try testing.expectEqualStrings("\x1b[1mbold\x1b[0m\x1bP$q\"p\x1b\\\x1bP+q544e\x1b\\\x1b\x1bP1$r\x1b\\\x1b", out);
    try testing.expectEqual(@as(usize, 1), c.images.items.len);
}

test "how images end" {
    var c: Collector = .{};
    defer c.deinit();
    const out = try run("\x1bPq~\x18x\x1bPq~\x1b[0my\x1bPq~", .{}, &c);
    defer testing.allocator.free(out);
    // CAN is swallowed with the image; the CSI that interrupts one is not.
    try testing.expectEqualStrings("x\x1b[0my", out);
    try testing.expectEqual(@as(usize, 3), c.images.items.len);
    try testing.expectEqual(Image.Termination.aborted, c.images.items[0].termination);
    try testing.expectEqual(Image.Termination.interrupted, c.images.items[1].termination);
    try testing.expectEqual(Image.Termination.end_of_data, c.images.items[2].termination);
}

test "an image interrupted by another image" {
    var c: Collector = .{};
    defer c.deinit();
    const out = try run("\x1bPq~\x1bPq~~\x1b\\", .{}, &c);
    defer testing.allocator.free(out);
    try testing.expectEqualStrings("", out);
    try testing.expectEqual(@as(usize, 2), c.images.items.len);
    try testing.expectEqual(@as(u32, 2), c.images.items[1].width());
}

test "C1 controls only when asked for" {
    var c: Collector = .{};
    defer c.deinit();
    const input = "é\x90q~\x9c!";
    const plain = try run(input, .{}, &c);
    defer testing.allocator.free(plain);
    try testing.expectEqualStrings(input, plain);
    try testing.expectEqual(@as(usize, 0), c.images.items.len);

    const eight = try run(input, .{ .c1 = true }, &c);
    defer testing.allocator.free(eight);
    try testing.expectEqualStrings("é!", eight);
    try testing.expectEqual(@as(usize, 1), c.images.items.len);
}

test "a reader with too small a buffer is refused" {
    var c: Collector = .{};
    defer c.deinit();
    var buf: [8]u8 = undefined;
    var reader: Io.Reader = .{ .vtable = Io.Reader.fixed("").vtable, .buffer = &buf, .seek = 0, .end = 0 };
    var discard: Io.Writer.Discarding = .init(&.{});
    try testing.expectError(error.ReaderBufferTooSmall, relay(testing.allocator, &reader, &discard.writer, &c, .{}));
}