//! The Espressif ROM loader protocol, enough of it to flash a chip: SLIP framing, SYNC, flash
//! attach, and uncompressed block writes. No software stub is uploaded - the ROM can do all of
//! this by itself, and for a ~1 KB image the stub's compression and 16 KB blocks buy nothing.
//!
//! Frame format (esptool loader.py:526-534, 577):
//!     request:  C0 | 00 op len16 chk32 | payload | C0
//!     response: C0 | 01 op len16 val32 | data    | C0
//! with C0 -> DB DC and DB -> DB DD inside the frame.
//!
//! The ESP32 ROM loaders (unlike the ESP8266's, and unlike the software stub) append FOUR trailing
//! bytes to every response: status, reason, and two reserved bytes (esptool loader.py:653-655,
//! 676). Reading the status at data[len - 2] therefore reads a reserved byte and turns every ROM
//! error into a success - which is exactly the bug an adversarial review of this file found, after
//! driving it over a pty with a rejected FLASH_DATA block.

const std = @import("std");
const Port = @import("serial.zig").Port;

pub const Cmd = enum(u8) {
    flash_begin = 0x02,
    flash_data = 0x03,
    flash_end = 0x04,
    sync = 0x08,
    read_reg = 0x0A,
    spi_set_params = 0x0B,
    spi_attach = 0x0D,
    change_baud = 0x0F,
    spi_flash_md5 = 0x13,
    get_security_info = 0x14,
};

pub const Error = error{
    SyncFailed,
    CommandFailed,
    ShortResponse,
    Timeout,
    BadFrame,
};

pub const Loader = struct {
    port: *Port,
    /// Bytes already read from the port but not yet consumed by the frame parser. Reading a byte
    /// at a time costs a poll+read syscall pair each, which turned a 1.5 KB flash into a 600 ms
    /// affair; refilling in bursts brings it under 60 ms.
    rx: [1024]u8 = undefined,
    rx_len: usize = 0,
    rx_pos: usize = 0,

    /// The ROM's own block size. The stub raises this to 0x4000; we do not use the stub.
    pub const block_size = 0x400;

    fn nextByte(l: *Loader, deadline_ms: i64) !?u8 {
        while (l.rx_pos == l.rx_len) {
            const remaining = deadline_ms - l.port.nowMs();
            if (remaining <= 0) return null;
            const n = try l.port.readTimeout(&l.rx, @intCast(@min(remaining, 50)));
            if (n == 0) continue;
            l.rx_len = n;
            l.rx_pos = 0;
        }
        defer l.rx_pos += 1;
        return l.rx[l.rx_pos];
    }

    /// Consume input until the line has been quiet for `quiet_ms`, but never for longer than
    /// twenty such windows: a board stuck in a brownout-reset loop re-prints its ROM banner
    /// forever, and an unbounded version of this loop hangs the flash with no output at all.
    fn drainUntilQuiet(l: *Loader, quiet_ms: i64) void {
        l.rx_pos = 0;
        l.rx_len = 0;
        const deadline = l.port.nowMs() + 20 * quiet_ms;
        while (l.port.nowMs() < deadline) {
            const n = l.port.readTimeout(&l.rx, @intCast(quiet_ms)) catch return;
            if (n == 0) return;
        }
    }

    fn frame(gpa: std.mem.Allocator, cmd: Cmd, payload: []const u8, checksum: u32) ![]u8 {
        var out: std.ArrayList(u8) = .empty;
        errdefer out.deinit(gpa);
        var head: [8]u8 = undefined;
        head[0] = 0x00;
        head[1] = @intFromEnum(cmd);
        std.mem.writeInt(u16, head[2..4], @intCast(payload.len), .little);
        std.mem.writeInt(u32, head[4..8], checksum, .little);

        try out.append(gpa, 0xC0);
        for (head) |b| try escape(gpa, &out, b);
        for (payload) |b| try escape(gpa, &out, b);
        try out.append(gpa, 0xC0);
        return out.toOwnedSlice(gpa);
    }

    fn escape(gpa: std.mem.Allocator, out: *std.ArrayList(u8), b: u8) !void {
        switch (b) {
            0xC0 => try out.appendSlice(gpa, &.{ 0xDB, 0xDC }),
            0xDB => try out.appendSlice(gpa, &.{ 0xDB, 0xDD }),
            else => try out.append(gpa, b),
        }
    }

    /// Read one SLIP frame, un-escaping as it goes.
    fn readFrame(l: *Loader, gpa: std.mem.Allocator, timeout_ms: u32) ![]u8 {
        var out: std.ArrayList(u8) = .empty;
        errdefer out.deinit(gpa);
        var started = false;
        var escaping = false;
        const deadline = l.port.nowMs() + @as(i64, timeout_ms);
        while (try l.nextByte(deadline)) |b| {
            if (!started) {
                if (b == 0xC0) started = true;
                continue;
            }
            if (escaping) {
                try out.append(gpa, switch (b) {
                    0xDC => 0xC0,
                    0xDD => 0xDB,
                    else => return Error.BadFrame,
                });
                escaping = false;
                continue;
            }
            switch (b) {
                0xDB => escaping = true,
                0xC0 => {
                    if (out.items.len == 0) continue; // empty frame, keep looking
                    return out.toOwnedSlice(gpa);
                },
                else => try out.append(gpa, b),
            }
        }
        return Error.Timeout;
    }

    /// Send a command and wait for its matching response. Returns the response `val` field, and
    /// copies any leading response data into `out` when one is given.
    pub fn command(
        l: *Loader,
        gpa: std.mem.Allocator,
        cmd: Cmd,
        payload: []const u8,
        checksum: u32,
        timeout_ms: u32,
        out: ?[]u8,
    ) !u32 {
        const pkt = try frame(gpa, cmd, payload, checksum);
        defer gpa.free(pkt);
        try l.port.write(pkt);

        const want_data: usize = if (out) |o| o.len else 0;
        var tries: usize = 0;
        while (tries < 100) : (tries += 1) {
            const resp = l.readFrame(gpa, timeout_ms) catch |err| return err;
            defer gpa.free(resp);
            // Skip anything that is not this command's reply: stale frames from a previous
            // session, or the ROM's repeated SYNC echoes.
            if (resp.len < 8) continue;
            if (resp[0] != 0x01 or resp[1] != @intFromEnum(cmd)) continue;

            const val = std.mem.readInt(u32, resp[4..8], .little);
            const data = resp[8..];
            // Status sits after the expected payload. The ROM appends four bytes (status, reason,
            // two reserved) where the stub appends two; esptool tolerates either, so gate on two
            // and read the status at the payload end - reading it at len-2 is what made every ROM
            // error look like success.
            if (data.len < want_data + 2) return Error.ShortResponse;
            if (data[want_data] != 0) return Error.CommandFailed;
            if (out) |o| @memcpy(o, data[0..want_data]);
            return val;
        }
        return Error.Timeout;
    }

    pub fn sync(l: *Loader, gpa: std.mem.Allocator) !void {
        var payload: [36]u8 = undefined;
        payload[0..4].* = .{ 0x07, 0x07, 0x12, 0x20 };
        @memset(payload[4..], 0x55);
        var attempt: usize = 0;
        // Short per-attempt timeout: the first SYNC after a reset usually lands before the ROM is
        // listening, and waiting 200 ms for that is most of the flash time on a small image.
        while (attempt < 20) : (attempt += 1) {
            // A reset that did not take is the usual reason SYNC never answers, so re-run it
            // periodically rather than failing the build - esptool retries the whole connect
            // seven times for the same reason (loader.py:891-899).
            if (attempt > 0 and attempt % 5 == 0) l.port.resetToDownload(.{}) catch {};
            if (l.command(gpa, .sync, &payload, 0, 40, null)) |_| {
                // The ROM answers SYNC eight times. Swallow the echoes, but stop as soon as the
                // line goes quiet instead of burning a fixed 350 ms.
                l.drainUntilQuiet(15);
                return;
            } else |_| {}
        }
        return Error.SyncFailed;
    }

    pub fn attachFlash(l: *Loader, gpa: std.mem.Allocator) !void {
        var payload: [8]u8 = @splat(0); // default SPI pins, not legacy
        _ = try l.command(gpa, .spi_attach, &payload, 0, 3000, null);
    }

    pub fn setFlashParams(l: *Loader, gpa: std.mem.Allocator, total_size: u32) !void {
        var payload: [24]u8 = undefined;
        std.mem.writeInt(u32, payload[0..4], 0, .little); // fl_id, ignored by the ROM
        std.mem.writeInt(u32, payload[4..8], total_size, .little);
        std.mem.writeInt(u32, payload[8..12], 64 * 1024, .little); // block
        std.mem.writeInt(u32, payload[12..16], 4 * 1024, .little); // sector
        std.mem.writeInt(u32, payload[16..20], 256, .little); // page
        std.mem.writeInt(u32, payload[20..24], 0xFFFF, .little); // status mask
        _ = try l.command(gpa, .spi_set_params, &payload, 0, 3000, null);
    }

    /// Write `data` at `offset`. The ROM erases synchronously inside FLASH_BEGIN.
    pub fn writeFlash(l: *Loader, gpa: std.mem.Allocator, offset: u32, data: []const u8) !void {
        const blocks: u32 = @intCast(std.math.divCeil(usize, data.len, block_size) catch unreachable);

        var begin: [20]u8 = undefined;
        std.mem.writeInt(u32, begin[0..4], @intCast(data.len), .little); // erase size
        std.mem.writeInt(u32, begin[4..8], blocks, .little);
        std.mem.writeInt(u32, begin[8..12], block_size, .little);
        std.mem.writeInt(u32, begin[12..16], offset, .little);
        std.mem.writeInt(u32, begin[16..20], 0, .little); // not encrypted
        const erase_timeout: u32 = @intCast(@max(@as(usize, 3000), data.len / 1024 * 30));
        _ = try l.command(gpa, .flash_begin, &begin, 0, erase_timeout, null);

        var seq: u32 = 0;
        var sent: usize = 0;
        var block_buf: [16 + block_size]u8 = undefined;
        while (sent < data.len) : (seq += 1) {
            const take = @min(block_size, data.len - sent);
            const chunk = data[sent .. sent + take];
            std.mem.writeInt(u32, block_buf[0..4], block_size, .little);
            std.mem.writeInt(u32, block_buf[4..8], seq, .little);
            std.mem.writeInt(u32, block_buf[8..12], 0, .little);
            std.mem.writeInt(u32, block_buf[12..16], 0, .little);
            @memcpy(block_buf[16 .. 16 + take], chunk);
            @memset(block_buf[16 + take ..], 0xFF); // ROM writes whole blocks; pad with erased value

            var checksum: u32 = 0xEF;
            for (block_buf[16..]) |b| checksum ^= b;

            var attempt: usize = 0;
            while (true) : (attempt += 1) {
                if (l.command(gpa, .flash_data, &block_buf, checksum, 3000, null)) |_| break else |err| {
                    if (attempt >= 2) return err;
                }
            }
            sent += take;
        }
    }

    /// MD5 over a flash range, as 32 ASCII hex characters from the ROM. Routed through
    /// `command()` so the direction byte, the opcode and the status are all checked - this is the
    /// only thing standing between a rejected block and a bricked image.
    pub fn flashMd5(l: *Loader, gpa: std.mem.Allocator, offset: u32, len: u32, out: *[32]u8) !void {
        var payload: [16]u8 = undefined;
        std.mem.writeInt(u32, payload[0..4], offset, .little);
        std.mem.writeInt(u32, payload[4..8], len, .little);
        std.mem.writeInt(u32, payload[8..12], 0, .little);
        std.mem.writeInt(u32, payload[12..16], 0, .little);
        _ = try l.command(gpa, .spi_flash_md5, &payload, 0, @max(1000, len / 1024 * 8), out);
    }
};