const is_debug = builtin.mode == .Debug;
const base64Enc = std.base64.Base64Encoder.init(std.base64.standard_alphabet_chars, '=');
const base64Dec = std.base64.Base64Decoder.init(std.base64.standard_alphabet_chars, '=');
var dba: DebugAllocator =
    if (is_debug)
        DebugAllocator.init
    else
        @compileError("Should not use debug allocator in release mode");

/// Type tag for SaprusMessage union.
/// This is the first value in the actual packet sent over the network.
const SaprusPacketType = enum(u16) {
    relay = 0x003C,
    file_transfer = 0x8888,
    connection = 0x00E9,
    _,
};

/// Reserved option values.
/// Currently unused.
const SaprusConnectionOptions = packed struct(u8) {
    opt1: bool = false,
    opt2: bool = false,
    opt3: bool = false,
    opt4: bool = false,
    opt5: bool = false,
    opt6: bool = false,
    opt7: bool = false,
    opt8: bool = false,
};

const SaprusError = error{
    NotImplementedSaprusType,
    UnknownSaprusType,
};

/// All Saprus messages
const SaprusMessage = union(SaprusPacketType) {
    const Relay = struct {
        const Header = packed struct {
            dest: @Vector(4, u8),
        };
        header: Header,
        payload: []const u8,
    };
    const Connection = struct {
        const Header = packed struct {
            src_port: u16,
            dest_port: u16,
            seq_num: u32 = 0,
            msg_id: u32 = 0,
            reserved: u8 = 0,
            options: SaprusConnectionOptions = .{},
        };
        header: Header,
        payload: []const u8,
    };
    relay: Relay,
    file_transfer: void, // unimplemented
    connection: Connection,

    /// Should be called for any SaprusMessage that was declared using a function that you pass an allocator to.
    fn deinit(self: SaprusMessage, allocator: Allocator) void {
        switch (self) {
            .relay => |r| allocator.free(r.payload),
            .connection => |c| allocator.free(c.payload),
            else => unreachable,
        }
    }

    inline fn toBytesAux(
        Header: type,
        header: Header,
        payload: []const u8,
        w: std.ArrayList(u8).Writer,
        allocator: Allocator,
    ) !void {
        // Create a growable string to store the base64 bytes in.
        // Doing this first so I can use the length of the encoded bytes for the length field.
        var payload_list = std.ArrayList(u8).init(allocator);
        defer payload_list.deinit();
        const buf_w = payload_list.writer();

        // Write the payload bytes as base64 to the growable string.
        try base64Enc.encodeWriter(buf_w, payload);

        // Write the packet body to the output writer.
        try w.writeStructEndian(header, .big);
        try w.writeInt(u16, @intCast(payload_list.items.len), .big);
        try w.writeAll(payload_list.items);
    }

    /// Caller is responsible for freeing the returned bytes.
    fn toBytes(self: SaprusMessage, allocator: Allocator) ![]u8 {
        // Create a growable list of bytes to store the output in.
        var buf = std.ArrayList(u8).init(allocator);
        // Create a writer for an easy interface to append arbitrary bytes.
        const w = buf.writer();

        // Start with writing the message type, which is the first 16 bits of every Saprus message.
        try w.writeInt(u16, @intFromEnum(self), .big);

        // Write the proper header and payload for the given packet type.
        switch (self) {
            .relay => |r| try toBytesAux(Relay.Header, r.header, r.payload, w, allocator),
            .connection => |c| try toBytesAux(Connection.Header, c.header, c.payload, w, allocator),
            .file_transfer => return SaprusError.NotImplementedSaprusType,
        }

        // Collect the growable list as a slice and return it.
        return buf.toOwnedSlice();
    }

    inline fn fromBytesAux(
        packet: SaprusPacketType,
        Header: type,
        r: std.io.FixedBufferStream([]const u8).Reader,
        allocator: Allocator,
    ) !SaprusMessage {
        // Read the header for the current message type.
        const header = try r.readStructEndian(Header, .big);
        // Read the length of the base64 encoded payload.
        const len = try r.readInt(u16, .big);

        // Read the base64 bytes into a list to be able to call the decoder on it.
        var payload_buf = std.ArrayList(u8).init(allocator);
        defer payload_buf.deinit();
        try r.readAllArrayList(&payload_buf, len);

        // Create a buffer to store the payload in, and decode the base64 bytes into the payload field.
        const payload = try allocator.alloc(u8, try base64Dec.calcSizeForSlice(payload_buf.items));
        try base64Dec.decode(payload, payload_buf.items);

        // Return the type of SaprusMessage specified by the `packet` argument.
        return @unionInit(SaprusMessage, @tagName(packet), .{
            .header = header,
            .payload = payload,
        });
    }

    /// Caller is responsible for calling .deinit on the returned value.
    fn fromBytes(bytes: []const u8, allocator: Allocator) !SaprusMessage {
        var s = std.io.fixedBufferStream(bytes);
        const r = s.reader();

        switch (@as(SaprusPacketType, @enumFromInt(try r.readInt(u16, .big)))) {
            .relay => return fromBytesAux(.relay, Relay.Header, r, allocator),
            .connection => return fromBytesAux(.connection, Connection.Header, r, allocator),
            .file_transfer => return SaprusError.NotImplementedSaprusType,
            else => return SaprusError.UnknownSaprusType,
        }
    }
};

pub fn main() !void {
    defer if (is_debug) {
        _ = dba.deinit();
    };

    var gpa = if (is_debug) dba.allocator() else std.heap.smp_allocator;

    const msg = SaprusMessage{
        .relay = .{
            .header = .{ .dest = .{ 255, 255, 255, 255 } },
            .payload = "Hello darkness my old friend",
        },
    };

    const msg_bytes = try msg.toBytes(gpa);
    defer gpa.free(msg_bytes);

    try network.init();
    defer network.deinit();

    var sock = try network.Socket.create(.ipv4, .udp);
    defer sock.close();

    try sock.setBroadcast(true);

    // Bind to 0.0.0.0:0
    const bind_addr = network.EndPoint{
        .address = network.Address{ .ipv4 = network.Address.IPv4.any },
        .port = 0,
    };

    const dest_addr = network.EndPoint{
        .address = network.Address{ .ipv4 = network.Address.IPv4.broadcast },
        .port = 8888,
    };

    try sock.bind(bind_addr);

    _ = try sock.sendTo(dest_addr, msg_bytes);
}

const builtin = @import("builtin");
const std = @import("std");
const Allocator = std.mem.Allocator;
const DebugAllocator = std.heap.DebugAllocator(.{});

const network = @import("network");

test "Round trip Relay toBytes and fromBytes" {
    const gpa = std.testing.allocator;
    const msg = SaprusMessage{
        .relay = .{
            .header = .{ .dest = .{ 255, 255, 255, 255 } },
            .payload = "Hello darkness my old friend",
        },
    };

    const to_bytes = try msg.toBytes(gpa);
    defer gpa.free(to_bytes);

    const from_bytes = try SaprusMessage.fromBytes(to_bytes, gpa);
    defer from_bytes.deinit(gpa);

    try std.testing.expectEqualDeep(msg, from_bytes);
}

test "Round trip Connection toBytes and fromBytes" {
    const gpa = std.testing.allocator;
    const msg = SaprusMessage{
        .connection = .{
            .header = .{
                .src_port = 0,
                .dest_port = 0,
            },
            .payload = "Hello darkness my old friend",
        },
    };

    const to_bytes = try msg.toBytes(gpa);
    defer gpa.free(to_bytes);

    const from_bytes = try SaprusMessage.fromBytes(to_bytes, gpa);
    defer from_bytes.deinit(gpa);

    try std.testing.expectEqualDeep(msg, from_bytes);
}