const base64Enc = std.base64.Base64Encoder.init(std.base64.standard_alphabet_chars, '=');
const base64Dec = std.base64.Base64Decoder.init(std.base64.standard_alphabet_chars, '=');

rand: Random,
writer: *std.Io.Writer,

const Self = @This();

const max_message_size = 2048;

pub fn init(writer: *std.Io.Writer) !Self {
    var prng = Random.DefaultPrng.init(blk: {
        var seed: u64 = undefined;
        try posix.getrandom(mem.asBytes(&seed));
        break :blk seed;
    });
    const rand = prng.random();

    return .{
        .rand = rand,
        .writer = writer,
    };
}

pub fn deinit(self: *Self) void {
    self.writer.flush() catch {};
}

/// Used for relay messages and connection handshake.
/// Assumes Client .init has been called.
fn broadcastInitialInterestMessage(self: *Self, msg_bytes: []align(@alignOf(SaprusMessage)) u8) !void {
    const writer = self.writer;

    const UdpHeaders = packed struct {
        dest_mac: u48, // [6]u8

        src_mac: u48, // [6]u8

        ether_type: u16,
        ip_version: u4,
        header_length: u4,
        type_of_service: u8,
        total_length: u16,

        identification: u16,
        ethernet_flags: u3,
        fragment_offset: u13,
        ttl: u8,
        protocol: u8,

        header_checksum: u16, // [2]u8
        src_ip: u16, // [2]u8
        dest_ip: u16, // [2]u8

        src_port: u16, // [2]u8
        dest_port: u16, // [2]u8
        length: u16,

        checksum: u16, // [2]u8

        fn init(
            in: struct {
                dest_mac: [6]u8,
                src_mac: [6]u8,
                ether_type: u16,
                ip_version: u4,
                header_length: u4 = 0,
                type_of_service: u4 = 0,
                total_length: usize,
                identification: u16 = 0,
                ethernet_flags: u3 = 0,
                fragment_offset: u13 = 0,
                ttl: u8 = 0,
                protocol: u8 = 0,
                header_checksum: [2]u8 = .{ 0, 0 },
                src_ip: [2]u8,
                dest_ip: [2]u8,
                src_port: [2]u8,
                dest_port: [2]u8,
                length: usize,
                checksum: [2]u8 = .{ 0, 0 },
            },
        ) @This() {
            return .{
                .dest_mac = @bitCast(in.dest_mac),
                .src_mac = @bitCast(in.src_mac),
                .ether_type = in.ether_type,
                .ip_version = in.ip_version,
                .header_length = in.header_length,
                .type_of_service = in.type_of_service,
                .total_length = @intCast(in.total_length),
                .identification = in.identification,
                .ethernet_flags = in.ethernet_flags,
                .fragment_offset = in.fragment_offset,
                .ttl = in.ttl,
                .protocol = in.protocol,
                .header_checksum = @bitCast(in.header_checksum),
                .src_ip = @bitCast(in.src_ip),
                .dest_ip = @bitCast(in.dest_ip),
                .src_port = @bitCast(in.src_port),
                .dest_port = @bitCast(in.dest_port),
                .length = @intCast(in.length),
                .checksum = @bitCast(in.checksum),
            };
        }
    };

    const total_len = (@bitSizeOf(UdpHeaders) / 8) + msg_bytes.len;
    std.debug.assert(writer.buffer.len >= total_len);
    _ = writer.consumeAll();

    const headers: UdpHeaders = .init(.{
        .dest_mac = .{ 0xff, 0xff, 0xff, 0xff, 0xff, 0xff },
        .src_mac = blk: {
            const r_bytes = try writer.writableArray(6);
            self.rand.bytes(r_bytes);
            break :blk r_bytes.*;
        },
        .ether_type = 0x0800,
        .ip_version = 0x4,
        .header_length = 0x5,
        .total_length = total_len - 8, // 8 is the ethernet frame length (macs + type)
        .protocol = 0x11,
        .src_ip = .{ 0, 0 },
        .dest_ip = .{ 0, 0 },
        .src_port = .{ 0, 0 },
        .dest_port = .{ 0xb8, 0x22 },
        .length = msg_bytes.len,
    });

    _ = try writer.write(&@as([@bitSizeOf(UdpHeaders) / 8]u8, @bitCast(headers)));

    // try writer.writeStruct(headers, .big);

    // // Ensure buffer is large enough
    // std.debug.assert(writer.buffer.len > 38 + msg_bytes.len);

    // // Ensure writer is clean
    // writer.consumeAll();

    // // Destination MAC addr to FF:FF:FF:FF:FF:FF
    // try writer.write();

    // // Source MAC to random bytes
    // {
    // }

    // // 96 bits

    // // Set EtherType to IPv4
    // try writer.write(.{ 0x08, 0x00 });

    // // 112 bits

    // // Set IPv4 version to 4
    // try writer.writeByte(0x45);

    // // 120 bits

    // // Unset section (Version, IHL, ToS)
    // writer.advance(2);

    // // 136 bits

    // // Total length
    // try write.writeInt(u16, 38 + msg_bytes.len);

    // // Destination broadcast
    // writer.splatByte(0xff, 0x22 - 0x1e);

    // var packet_bytes: [_]u8 = comptime blk: {
    //     var b: [max_message_size]u8 = @splat(0);

    //     for (0x1e..0x22) |i| {
    //         b[i] = 0xff;
    //     }

    //     // Set TTL
    //     b[0x16] = 0x40;

    //     // Set IPv4 protocol to UDP
    //     b[0x17] = 0x11;

    //     // Set interest filter value to 8888.
    //     b[0x24] = 0x22;
    //     b[0x25] = 0xb8;
    //     break :blk &b;
    // };
    var msg: *SaprusMessage = try .bytesAsValue(msg_bytes);
    try msg.networkFromNativeEndian();
    defer msg.nativeFromNetworkEndian() catch unreachable;

    std.debug.print("headers: {x}\n", .{@as([@bitSizeOf(UdpHeaders) / 8]u8, @bitCast(headers))});
    std.debug.print("bytes: {x}\n", .{writer.buffer[0..writer.end]});
    _ = try writer.write(msg_bytes);

    // The byte position within the packet that the saprus message starts at.
    // const saprus_start_byte = 42;
    // @memcpy(packet_bytes[saprus_start_byte .. saprus_start_byte + msg_bytes.len], msg_bytes);
    // _ = try writer.write(packet_bytes[0 .. saprus_start_byte + msg_bytes.len]);
    try writer.flush();
}

// fn broadcastSaprusMessage(msg_bytes: []align(@alignOf(SaprusMessage)) u8) !void {}

fn broadcastSaprusMessage(msg_bytes: []align(@alignOf(SaprusMessage)) u8, udp_port: u16) !void {
    const msg: *SaprusMessage = try .bytesAsValue(msg_bytes);
    try msg.networkFromNativeEndian();
    defer msg.nativeFromNetworkEndian() catch unreachable;

    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 = udp_port,
    };

    try sock.bind(bind_addr);

    std.debug.print("{x}\n", .{msg_bytes});

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

pub fn sendRelay(self: *Self, payload: []const u8, dest: [4]u8) !void {
    var buf: [max_message_size]u8 align(@alignOf(SaprusMessage)) = undefined;
    const msg_bytes = buf[0..try SaprusMessage.calcSize(
        .relay,
        base64Enc.calcSize(payload.len),
    )];
    const msg: *SaprusMessage = .init(.relay, msg_bytes);

    const relay = (try msg.getSaprusTypePayload()).relay;
    relay.dest = dest;
    _ = base64Enc.encode(relay.getPayload(), payload);

    try self.broadcastInitialInterestMessage(msg_bytes);
}

fn randomPort(self: Self) u16 {
    return self.rand.intRangeAtMost(u16, 1024, 65000);
}

pub fn sendInitialConnection(
    self: Self,
    payload: []const u8,
    output_bytes: []align(@alignOf(SaprusMessage)) u8,
    initial_port: u16,
) !*SaprusMessage {
    const dest_port = self.randomPort();
    const msg_bytes = output_bytes[0..try SaprusMessage.calcSize(
        .connection,
        base64Enc.calcSize(payload.len),
    )];
    const msg: *SaprusMessage = .init(.connection, msg_bytes);

    const connection = (try msg.getSaprusTypePayload()).connection;
    connection.src_port = initial_port;
    connection.dest_port = dest_port;
    _ = base64Enc.encode(connection.getPayload(), payload);

    try broadcastSaprusMessage(msg_bytes, 8888);

    return msg;
}

pub fn connect(self: Self, payload: []const u8) !?SaprusConnection {
    const initial_port = self.randomPort();

    var initial_conn_res: ?*SaprusMessage = null;

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

    // Bind to 255.255.255.255:8888
    const bind_addr = network.EndPoint{
        .address = network.Address{ .ipv4 = network.Address.IPv4.broadcast },
        .port = 8888,
    };

    // timeout 1s
    try sock.setReadTimeout(1 * std.time.us_per_s);
    try sock.bind(bind_addr);

    var sent_msg_bytes: [max_message_size]u8 align(@alignOf(SaprusMessage)) = undefined;
    const msg = try self.sendInitialConnection(payload, &sent_msg_bytes, initial_port);

    var response_buf: [max_message_size]u8 align(@alignOf(SaprusMessage)) = undefined;
    _ = try sock.receive(&response_buf); // Ignore message that I sent.
    const len = try sock.receive(&response_buf);

    initial_conn_res = try .networkBytesAsValue(response_buf[0..len]);

    // Complete handshake after awaiting response
    try broadcastSaprusMessage(msg.asBytes(), self.randomPort());

    if (false) {
        return initial_conn_res.?;
    }
    return null;
}

const SaprusMessage = @import("message.zig").Message;
const SaprusConnection = @import("Connection.zig");

const std = @import("std");
const Random = std.Random;
const posix = std.posix;
const mem = std.mem;

const network = @import("network");