const std = @import("std");
const ati = @import("ati.zig");

const Allocator = std.mem.Allocator;
const ArrayList = std.ArrayList;
const ArenaAllocator = std.heap.ArenaAllocator;
const t = std.testing;

pub const OpCode = enum(u8) {
    lambda = 0xAD,
    delay_us = 0xA0,
    delay_ms = 0x94,
    set_icsp_pins = 0xB2,
    get_icsp_pins = 0x32,
    set_vpp_on = 0xB0,
    set_vpp_off = 0xB1,
    write_literal_32_lsb = 0xB4,
    write_bits_literal = 0xB5,
    p16env3_write_payload_param = 0x86,
    p16env3_write_payload_literal = 0x87,
    p16env3_write_buffer = 0x88,
    p16env3_write_buffer_dfm = 0x49,
    p16env3_read_payload_pfm = 0x84,
    p16env3_read_payload_dfm = 0x48,
    write_bits_literal_msb = 0x83,
    set_clk_hi = 0xBA,
    set_clk_lo = 0xBB,
    p16f_read_loc_buffer = 0x80,
    p16f_write_loc_buffer = 0x81,
    de_command = 0x31,
    coreinst24 = 0xE0,
    visi24 = 0xE2,
    p24_send_pe_word = 0xE7,
    p24_send_pe_word_buf = 0xE8,
    p24_pe_handshake = 0xE9,
    p24_receive_pe_word = 0xEA,
    set_speed = 0xEC,
    get_speed = 0xED,
};

pub const Builder = struct {
    arena: ArenaAllocator,
    elems: ArrayList([]u8) = .empty,

    pub fn init(child: Allocator) Builder {
        return .{ .arena = .init(child) };
    }

    pub fn deinit(self: *Builder) void {
        self.arena.deinit();
    }

    pub fn reset(self: *Builder) void {
        _ = self.arena.reset(.retain_capacity);
        self.elems = .empty;
    }

    fn a(self: *Builder) Allocator {
        return self.arena.allocator();
    }

    pub fn newElement(self: *Builder, opcode: OpCode) Allocator.Error!void {
        const el = try self.a().alloc(u8, 1);
        el[0] = @intFromEnum(opcode);
        try self.elems.append(self.a(), el);
    }

    pub fn appendByte(self: *Builder, value: u8) Allocator.Error!void {
        const last = &self.elems.items[self.elems.items.len - 1];
        const grown = try self.a().realloc(last.*, last.*.len + 1);
        grown[grown.len - 1] = value;
        last.* = grown;
    }

    pub fn appendLe16(self: *Builder, value: u16) Allocator.Error!void {
        try self.appendByte(@truncate(value));
        try self.appendByte(@truncate(value >> 8));
    }

    pub fn appendLe32(self: *Builder, value: u32) Allocator.Error!void {
        try self.appendByte(@truncate(value));
        try self.appendByte(@truncate(value >> 8));
        try self.appendByte(@truncate(value >> 16));
        try self.appendByte(@truncate(value >> 24));
    }

    pub fn elements(self: *const Builder) []const []u8 {
        return self.elems.items;
    }
};

const Node = union(enum) {
    b: u8,
    seq: []const Node,
};

fn nodeEql(x: Node, y: Node) bool {
    return switch (x) {
        .b => |xb| switch (y) {
            .b => |yb| xb == yb,
            .seq => false,
        },
        .seq => |xs| switch (y) {
            .b => false,
            .seq => |ys| seqEql(xs, ys),
        },
    };
}

fn seqEql(a: []const Node, b: []const Node) bool {
    if (a.len != b.len) return false;
    for (a, b) |xi, yi| {
        if (!nodeEql(xi, yi)) return false;
    }
    return true;
}

fn rollOnce(gpa: Allocator, content: *ArrayList(Node)) Allocator.Error!bool {
    const n = content.items.len;
    const max_seq_len = n / 2;
    var seqlen: usize = 1;
    while (seqlen < max_seq_len) : (seqlen += 1) {
        var start: usize = 0;
        while (start < n - seqlen) : (start += 1) {
            const needle = content.items[start .. start + seqlen];
            var repeats: usize = 0;
            var pos: usize = start + seqlen;
            while (pos + seqlen <= n) : (pos += seqlen) {
                if (!seqEql(content.items[pos .. pos + seqlen], needle)) break;
                repeats += 1;
                if (repeats == 255) break;
            }
            if (repeats > 0 and seqlen * repeats > 4) {
                const prefix = content.items[0..start];
                const suffix_start = start + seqlen * (repeats + 1);
                const suffix = content.items[suffix_start..];

                var drop_i = start + seqlen;
                while (drop_i < suffix_start) : (drop_i += 1) {
                    freeTree(gpa, content.items[drop_i]);
                }

                var contraction: ArrayList(Node) = .empty;
                try contraction.append(gpa, .{ .b = @intFromEnum(OpCode.lambda) });
                try contraction.append(gpa, .{ .b = @intCast(repeats) });
                try contraction.append(gpa, .{ .b = @intCast(seqlen) });
                try contraction.appendSlice(gpa, needle);

                var new_content: ArrayList(Node) = .empty;
                try new_content.appendSlice(gpa, prefix);
                const contracted = try contraction.toOwnedSlice(gpa);
                try new_content.append(gpa, .{ .seq = contracted });
                try new_content.appendSlice(gpa, suffix);

                content.deinit(gpa);
                content.* = new_content;
                return true;
            }
        }
    }
    return false;
}

fn encloseAsLambda(gpa: Allocator, root: *ArrayList(Node)) (Allocator.Error || error{SequenceTooLong})!void {
    const length = root.items.len;
    if (length > 255) return error.SequenceTooLong;
    var outer: ArrayList(Node) = .empty;
    try outer.ensureTotalCapacity(gpa, 3 + length);
    outer.appendAssumeCapacity(.{ .b = @intFromEnum(OpCode.lambda) });
    outer.appendAssumeCapacity(.{ .b = 0 });
    outer.appendAssumeCapacity(.{ .b = @intCast(length) });
    outer.appendSliceAssumeCapacity(root.items);
    root.deinit(gpa);
    root.* = outer;
}

fn flatten(node: Node, out: *ArrayList(u8), gpa: Allocator) Allocator.Error!void {
    switch (node) {
        .b => |v| try out.append(gpa, v),
        .seq => |children| for (children) |child| {
            try flatten(child, out, gpa);
        },
    }
}

fn freeTree(gpa: Allocator, node: Node) void {
    switch (node) {
        .b => {},
        .seq => |children| {
            for (children) |child| freeTree(gpa, child);
            gpa.free(children);
        },
    }
}

pub fn processSequence(
    gpa: Allocator,
    elements: []const []const u8,
    use_lambda: bool,
) (Allocator.Error || error{SequenceTooLong})![]u8 {
    var root: ArrayList(Node) = .empty;
    defer root.deinit(gpa);
    for (elements) |element| {
        const children = try gpa.alloc(Node, element.len);
        for (element, 0..) |byte, i| children[i] = .{ .b = byte };
        try root.append(gpa, .{ .seq = children });
    }
    while (try rollOnce(gpa, &root)) {}
    if (use_lambda) try encloseAsLambda(gpa, &root);
    var out: ArrayList(u8) = .empty;
    errdefer out.deinit(gpa);
    for (root.items) |node| try flatten(node, &out, gpa);
    for (root.items) |node| freeTree(gpa, node);
    return out.toOwnedSlice(gpa);
}

const envelope_major: u8 = 1;
const envelope_minor: u8 = 0;
const data_source_undefined: u8 = 0xFF;
const data_dest_undefined: u8 = 0xFF;

const Command = struct {
    data_source: u8 = data_source_undefined,
    data_dest: u8 = data_dest_undefined,
    content: []const u8,

    pub fn bytestream(self: Command, gpa: Allocator) Allocator.Error![]u8 {
        var out: ArrayList(u8) = .empty;
        errdefer out.deinit(gpa);
        try out.ensureTotalCapacity(gpa, 4 + self.content.len);
        out.appendAssumeCapacity(self.data_source);
        out.appendAssumeCapacity(self.data_dest);
        out.appendAssumeCapacity(0);
        out.appendAssumeCapacity(0xFF);
        out.appendSliceAssumeCapacity(self.content);
        return out.toOwnedSlice(gpa);
    }
};

pub fn executionPayload(
    gpa: Allocator,
    blocks: []const []const u8,
) (Allocator.Error || error{BlockTooBig} || error{TooManyBlocks})![]u8 {
    if (blocks.len > 255) return error.TooManyBlocks;
    var out: ArrayList(u8) = .empty;
    errdefer out.deinit(gpa);
    const hdr = ati.header(.pic_primitive);
    try out.appendSlice(gpa, &hdr);
    try out.append(gpa, envelope_major);
    try out.append(gpa, envelope_minor);
    try out.append(gpa, @intCast(blocks.len));
    for (blocks) |block| {
        if (block.len > 255) return error.BlockTooBig;
        try out.append(gpa, @intCast(block.len));
        try out.appendSlice(gpa, block);
    }
    return out.toOwnedSlice(gpa);
}

pub fn blockStatus(raw: []const u8, index: usize) u32 {
    return std.mem.readInt(u32, raw[index * 4 ..][0..4], .little);
}

test "builder accumulates elements" {
    var b: Builder = .init(t.allocator);
    defer b.deinit();
    try b.newElement(.set_vpp_on);
    try b.newElement(.delay_ms);
    try b.appendLe16(100);
    try b.newElement(.set_icsp_pins);
    try b.appendByte(0);

    const els = b.elements();
    try t.expectEqual(@as(usize, 3), els.len);
    try t.expectEqualSlices(u8, &.{@intFromEnum(OpCode.set_vpp_on)}, els[0]);
    try t.expectEqualSlices(u8, &.{ @intFromEnum(OpCode.delay_ms), 100, 0 }, els[1]);
    try t.expectEqualSlices(u8, &.{ @intFromEnum(OpCode.set_icsp_pins), 0 }, els[2]);

    b.reset();
    try t.expectEqual(@as(usize, 0), b.elements().len);
}

test "processSequence rolls repeated elements" {
    const e = [_]u8{ @intFromEnum(OpCode.write_literal_32_lsb), 0xAA };
    var els: [6][]const u8 = undefined;
    for (&els) |*el| el.* = &e;

    const seq = try processSequence(t.allocator, &els, true);
    defer t.allocator.free(seq);

    try t.expectEqualSlices(u8, &.{ 0xAD, 0x00, 0x01, 0xAD, 0x05, 0x01, 0xB4, 0xAA }, seq);
}

test "enterTmod sequence exact bytes" {
    const elements = [_][]const u8{
        &.{@intFromEnum(OpCode.set_vpp_on)},
        &.{ @intFromEnum(OpCode.delay_ms), 100, 0 },
        &.{ @intFromEnum(OpCode.set_icsp_pins), 0 },
        &.{@intFromEnum(OpCode.set_vpp_off)},
        &.{ @intFromEnum(OpCode.delay_ms), 100, 0 },
        &.{ @intFromEnum(OpCode.delay_us), 250, 0 },
        &.{ @intFromEnum(OpCode.write_literal_32_lsb), 0x50, 0x48, 0x43, 0x4D },
        &.{@intFromEnum(OpCode.set_clk_hi)},
        &.{@intFromEnum(OpCode.set_clk_lo)},
    };

    const content = try processSequence(t.allocator, &elements, true);
    defer t.allocator.free(content);

    const expected_content = [_]u8{
        0xAD,                             0x00,                                      0x09,
        @intFromEnum(OpCode.set_vpp_on),  @intFromEnum(OpCode.delay_ms),             100,
        0,                                @intFromEnum(OpCode.set_icsp_pins),        0,
        @intFromEnum(OpCode.set_vpp_off), @intFromEnum(OpCode.delay_ms),             100,
        0,                                @intFromEnum(OpCode.delay_us),             250,
        0,                                @intFromEnum(OpCode.write_literal_32_lsb), 0x50,
        0x48,                             0x43,                                      0x4D,
        @intFromEnum(OpCode.set_clk_hi),  @intFromEnum(OpCode.set_clk_lo),
    };
    try t.expectEqualSlices(u8, &expected_content, content);

    const cmd: Command = .{ .content = content };
    const stream = try cmd.bytestream(t.allocator);
    defer t.allocator.free(stream);
    try t.expectEqual(27, stream.len);
    try t.expectEqualSlices(u8, &.{ 0xFF, 0xFF, 0x00, 0xFF }, stream[0..4]);

    const payload = try executionPayload(t.allocator, &.{stream});
    defer t.allocator.free(payload);
    try t.expectEqual(35, payload.len);
    try t.expectEqualSlices(u8, &.{ 1, 0, 0x30, 0, 1, 0, 1, 27 }, payload[0..8]);
}

test "command framing with data dest" {
    const cmd: Command = .{ .data_dest = 1, .content = "AB" };
    const stream = try cmd.bytestream(t.allocator);
    defer t.allocator.free(stream);
    try t.expectEqualSlices(u8, &.{ 0xFF, 0x01, 0x00, 0xFF, 'A', 'B' }, stream);
}

test "blockStatus parses le32 results" {
    const raw = [_]u8{ 0x2A, 0x00, 0x00, 0x00, 0x01, 0x02, 0x03, 0x04 };
    try t.expectEqual(42, blockStatus(&raw, 0));
    try t.expectEqual(0x04030201, blockStatus(&raw, 1));
}