const std = @import("std");

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

const SerialiseError = error{
    OutOfMemory,
    StringTooLong,
    BinaryTooLong,
};

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, i: i64) SerialiseError![]u8 {
    if (i < 0 and i >= std.math.minInt(i6)) {
        // neg fixint
        const bytes = try alloc.alloc(u8, 1);

        std.mem.writeInt(i8, bytes[0..1], @intCast(i), .big);

        return bytes;
    }

    if (intInRange(i8, i)) { // int8
        const bytes = try alloc.alloc(u8, 1 + 1);

        bytes[0] = 0xd0;

        std.mem.writeInt(i8, bytes[1..2], @intCast(i), .big);

        return bytes;
    }

    if (intInRange(i16, i)) { // int16
        const bytes = try alloc.alloc(u8, 1 + 2);

        bytes[0] = 0xd1;

        std.mem.writeInt(i16, bytes[1..3], @intCast(i), .big);

        return bytes;
    }

    if (intInRange(i32, i)) { // int32
        const bytes = try alloc.alloc(u8, 1 + 4);

        bytes[0] = 0xd2;

        std.mem.writeInt(i32, bytes[1..5], @intCast(i), .big);

        return bytes;
    }

    // int64
    const bytes = try alloc.alloc(u8, 1 + 8);

    bytes[0] = 0xd3;

    std.mem.writeInt(i64, bytes[1..9], i, .big);

    return bytes;
}

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

fn serialise_uint(alloc: std.mem.Allocator, u: u64) SerialiseError![]u8 {
    if (uintInRange(u7, u)) { // pos fixint
        const bytes = try alloc.alloc(u8, 1);

        std.mem.writeInt(u8, bytes[0..1], @intCast(u), .big);

        return bytes;
    }

    if (uintInRange(u8, u)) { // uint8
        const bytes = try alloc.alloc(u8, 1 + 1);

        bytes[0] = 0xcc;

        std.mem.writeInt(u8, bytes[1..2], @intCast(u), .big);

        return bytes;
    }

    if (uintInRange(u16, u)) { // uint16
        const bytes = try alloc.alloc(u8, 1 + 2);

        bytes[0] = 0xcd;

        std.mem.writeInt(u16, bytes[1..3], @intCast(u), .big);

        return bytes;
    }

    if (uintInRange(u32, u)) { // uint32
        const bytes = try alloc.alloc(u8, 1 + 4);

        bytes[0] = 0xce;

        std.mem.writeInt(u32, bytes[1..5], @intCast(u), .big);

        return bytes;
    }

    // uint64
    const bytes = try alloc.alloc(u8, 1 + 8);

    bytes[0] = 0xcf;

    std.mem.writeInt(u64, bytes[1..9], u, .big);

    return bytes;
}

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

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

        bytes[0] = 0b101_00000 | @as(u8, @intCast(s.len));

        @memcpy(bytes[1..], s);

        return bytes;
    }

    if (s.len <= std.math.maxInt(u8)) { // str8 / bin8
        const bytes = try alloc.alloc(u8, 1 + 1 + s.len);

        bytes[0] = switch (kind) {
            .string => 0xd9,
            .binary => 0xc4,
        };

        bytes[1] = @as(u8, @intCast(s.len));

        @memcpy(bytes[2..], s);

        return bytes;
    }

    if (s.len <= std.math.maxInt(u16)) { // str16 / bin16
        const bytes = try alloc.alloc(u8, 1 + 2 + s.len);

        bytes[0] = switch (kind) {
            .string => 0xda,
            .binary => 0xc5,
        };

        std.mem.writeInt(u16, bytes[1..3], @intCast(s.len), .big);

        @memcpy(bytes[3..], s);

        return bytes;
    }

    if (s.len <= std.math.maxInt(u32)) { // str16 / bin16
        const bytes = try alloc.alloc(u8, 1 + 4 + s.len);

        bytes[0] = switch (kind) {
            .string => 0xdb,
            .binary => 0xc6,
        };

        std.mem.writeInt(u32, bytes[1..5], @intCast(s.len), .big);

        @memcpy(bytes[5..], s);

        return bytes;
    }

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

pub fn serialise(alloc: std.mem.Allocator, obj: Object) SerialiseError![]u8 {
    switch (obj) {
        .nil => return try alloc.dupe(u8, &[_]u8{0xc0}),
        .bool => |b| if (b) {
            return try alloc.dupe(u8, &[_]u8{0xc3});
        } else {
            return try alloc.dupe(u8, &[_]u8{0xc2});
        },
        .integer => |i| return try serialise_int(alloc, i),
        .uinteger => |u| return try serialise_uint(alloc, u),
        .float32 => |f| {
            const bytes = try alloc.alloc(u8, 1 + 4);

            bytes[0] = 0xca;

            std.mem.writeInt(u32, bytes[1..5], @bitCast(f), .big);

            return bytes;
        },
        .float64 => |f| {
            const bytes = try alloc.alloc(u8, 1 + 8);

            bytes[0] = 0xcb;

            std.mem.writeInt(u64, bytes[1..9], @bitCast(f), .big);

            return bytes;
        },
        .string => |s| return try serialise_raw(alloc, s, .string),
        .binary => |b| return try serialise_raw(alloc, b, .binary),
        else => unreachable,
    }

    // const bytes = try alloc.alloc(u8, 0);

    // return bytes;
}