const std = @import("std");

const Object = @import("root.zig").Object;

const MapEntry = @import("root.zig").MapEntry;

const SerialiseError = error{
    OutOfMemory,
    StringTooLong,
    BinaryTooLong,
    ArrayTooLong,
    MapTooLong,
    ExtTooLong,
};

fn intInRange(comptime T: type, i: i64) bool {
    return std.math.minInt(T) <= i and i <= std.math.maxInt(T);
}

fn serialise_int(alloc: std.mem.Allocator, buf: *std.ArrayListUnmanaged(u8), i: i64) SerialiseError!void {
    if (i < 0 and i >= std.math.minInt(i6)) { // neg fixint
        std.mem.writeInt(i8, try buf.addManyAsArray(alloc, 1), @intCast(i), .big);

        return;
    }

    inline for (0.., &[_]type{ i8, i16, i32, i64 }) |n, T| {
        if (intInRange(T, i)) {
            try buf.ensureUnusedCapacity(alloc, 1 + @sizeOf(T));

            try buf.append(alloc, 0xd0 + n);

            std.mem.writeInt(T, try buf.addManyAsArray(alloc, @sizeOf(T)), @intCast(i), .big);

            return;
        }
    }

    unreachable;
}

fn uintInRange(comptime T: type, u: u64) bool {
    return u <= std.math.maxInt(T);
}

fn serialise_uint(alloc: std.mem.Allocator, buf: *std.ArrayListUnmanaged(u8), u: u64) SerialiseError!void {
    if (uintInRange(u7, u)) { // pos fixint
        std.mem.writeInt(u8, try buf.addManyAsArray(alloc, 1), @intCast(u), .big);

        return;
    }

    inline for (0.., &[_]type{ u8, u16, u32, u64 }) |n, T| {
        if (uintInRange(T, u)) {
            try buf.ensureUnusedCapacity(alloc, 1 + @sizeOf(T));

            try buf.append(alloc, 0xcc + n);

            std.mem.writeInt(T, try buf.addManyAsArray(alloc, @sizeOf(T)), @intCast(u), .big);

            return;
        }
    }

    unreachable;
}

fn serialise_raw(alloc: std.mem.Allocator, buf: *std.ArrayListUnmanaged(u8), s: []const u8, comptime kind: enum { string, binary }) SerialiseError!void {
    // TODO: should we validate that `s` is UTF-8 here?

    if (kind == .string and s.len <= 31) { // fixstr
        try buf.ensureUnusedCapacity(alloc, 1 + s.len);

        try buf.append(alloc, 0b101_00000 | @as(u8, @intCast(s.len)));

        try buf.appendSlice(alloc, s);

        return;
    }

    inline for (0.., &[_]type{ u8, u16, u32 }) |i, T| {
        if (s.len <= std.math.maxInt(T)) {
            try buf.ensureUnusedCapacity(alloc, 1 + @sizeOf(T) + s.len);

            switch (kind) {
                .string => try buf.append(alloc, 0xd9 + i),
                .binary => try buf.append(alloc, 0xc4 + i),
            }

            std.mem.writeInt(T, try buf.addManyAsArray(alloc, @sizeOf(T)), @intCast(s.len), .big);

            try buf.appendSlice(alloc, s);

            return;
        }
    }

    return switch (kind) {
        .string => error.StringTooLong,
        .binary => error.BinaryTooLong,
    };
}

fn serialise_array(alloc: std.mem.Allocator, buf: *std.ArrayListUnmanaged(u8), array: []const Object) SerialiseError!void {
    if (array.len <= 15) {
        try buf.append(alloc, 0b1001_0000 | @as(u8, @intCast(array.len)));
    } else if (array.len <= std.math.maxInt(u16)) {
        try buf.append(alloc, 0xdc);

        std.mem.writeInt(u16, try buf.addManyAsArray(alloc, @sizeOf(u16)), @intCast(array.len), .big);
    } else if (array.len <= std.math.maxInt(u32)) {
        try buf.append(alloc, 0xdd);

        std.mem.writeInt(u32, try buf.addManyAsArray(alloc, @sizeOf(u32)), @intCast(array.len), .big);
    } else {
        return error.ArrayTooLong;
    }

    for (array) |obj| {
        try serialise_into_buf(alloc, buf, obj);
    }
}

fn serialise_map(alloc: std.mem.Allocator, buf: *std.ArrayListUnmanaged(u8), map: []const MapEntry) SerialiseError!void {
    if (map.len <= 15) {
        try buf.append(alloc, 0b1000_0000 | @as(u8, @intCast(map.len)));
    } else if (map.len <= std.math.maxInt(u16)) {
        try buf.append(alloc, 0xde);

        std.mem.writeInt(u16, try buf.addManyAsArray(alloc, @sizeOf(u16)), @intCast(map.len), .big);
    } else if (map.len <= std.math.maxInt(u32)) {
        try buf.append(alloc, 0xdf);

        std.mem.writeInt(u32, try buf.addManyAsArray(alloc, @sizeOf(u32)), @intCast(map.len), .big);
    } else {
        return error.MapTooLong;
    }

    for (map) |entry| {
        try serialise_into_buf(alloc, buf, entry.key);
        try serialise_into_buf(alloc, buf, entry.value);
    }
}

fn serialise_ext(alloc: std.mem.Allocator, buf: *std.ArrayListUnmanaged(u8), type_: u8, bytes: []u8) SerialiseError!void {
    // fixext
    for (0.., [_]usize{ 1, 2, 4, 8, 16 }) |n, len| {
        if (bytes.len == len) {
            try buf.ensureUnusedCapacity(alloc, 1 + 1 + len);

            try buf.append(alloc, 0xd4 + @as(u8, @intCast(n)));
            try buf.append(alloc, type_);

            try buf.appendSlice(alloc, bytes);

            return;
        }
    }

    // ext
    inline for (0.., [_]type{ u8, u16, u32 }) |n, T| {
        if (bytes.len <= std.math.maxInt(T)) {
            try buf.ensureUnusedCapacity(alloc, 1 + @sizeOf(T) + bytes.len);

            try buf.append(alloc, 0xc7 + n);

            std.mem.writeInt(T, try buf.addManyAsArray(alloc, @sizeOf(T)), @intCast(bytes.len), .big);

            try buf.append(alloc, type_);

            try buf.appendSlice(alloc, bytes);

            return;
        }
    }

    return error.ExtTooLong;
}

fn serialise_into_buf(alloc: std.mem.Allocator, buf: *std.ArrayListUnmanaged(u8), obj: Object) SerialiseError!void {
    switch (obj) {
        .nil => try buf.append(alloc, 0xc0),
        .bool => |b| if (b) {
            try buf.append(alloc, 0xc3);
        } else {
            try buf.append(alloc, 0xc2);
        },
        .integer => |i| try serialise_int(alloc, buf, i),
        .uinteger => |u| try serialise_uint(alloc, buf, u),
        .float32 => |f| {
            try buf.ensureUnusedCapacity(alloc, 1 + 4);

            try buf.append(alloc, 0xca);

            std.mem.writeInt(u32, try buf.addManyAsArray(alloc, 4), @bitCast(f), .big);
        },
        .float64 => |f| {
            try buf.ensureUnusedCapacity(alloc, 1 + 8);

            try buf.append(alloc, 0xcb);

            std.mem.writeInt(u64, try buf.addManyAsArray(alloc, 8), @bitCast(f), .big);
        },
        .string => |s| try serialise_raw(alloc, buf, s, .string),
        .binary => |b| try serialise_raw(alloc, buf, b, .binary),
        .array => |array| try serialise_array(alloc, buf, array),
        .map => |map| try serialise_map(alloc, buf, map),
        .extension => |ext| try serialise_ext(alloc, buf, ext.type, ext.bytes),
    }
}

pub fn serialise(alloc: std.mem.Allocator, obj: Object) SerialiseError![]u8 {
    var buf = std.ArrayListUnmanaged(u8){};
    defer buf.deinit(alloc);

    try serialise_into_buf(alloc, &buf, obj);

    return try buf.toOwnedSlice(alloc);
}