//! Wraps a writer with UDP headers.
//! This is useful for wrapping RawSocket Writer with appropriate headers.

rand: Random,
wrapped: *Writer,
interface: Writer,

pub fn init(w: *Writer) !NetWriter {
    std.debug.assert(w.buffer.len > @sizeOf(EthernetHeaders) + @sizeOf(IpHeaders) + @sizeOf(UdpHeaders));

    var prng = Random.DefaultPrng.init(blk: {
        var seed: u64 = undefined;
        try posix.getrandom(mem.asBytes(&seed));
        break :blk seed;
    });

    return .{
        .rand = prng.random(),
        .wrapped = w,
        .interface = .{
            .vtable = &.{
                .drain = drain,
                .flush = flush,
            },
            .buffer = &.{},
        },
    };
}

fn drain(io_w: *Writer, data: []const []const u8, splat: usize) Writer.Error!usize {
    const w: *NetWriter = @alignCast(@fieldParentPtr("interface", io_w));

    var res: usize = 0;

    if (io_w.end == 0) {
        const ether_headers: EthernetHeaders = .{
            .dest_mac = .{ 0xff, 0xff, 0xff, 0xff, 0xfe, 0xff },
            .src_mac = blk: {
                var output_bytes: [6]u8 = undefined;
                output_bytes[0] = 0xee;
                w.rand.bytes(output_bytes[1..]);
                break :blk output_bytes;
            },
            .ether_type = 0x0800,
        };

        const ip_headers: IpHeaders = .{
            .total_length = @intCast(res - 92),
            .ttl = 0x64,
            .protocol = 0x11,
            .src_ip = .{ 0xff, 0x02, 0x03, 0x04 },
            .dest_ip = .{ 0xff, 0xff, 0xff, 0xff },
        };

        const udp_headers: UdpHeaders = .{
            .src_port = 0xbbbb,
            .dest_port = 8888,
            .length = @intCast(res),
        };

        res += try ether_headers.write(w.wrapped);
        res += try ip_headers.write(w.wrapped);
        res += try udp_headers.write(w.wrapped);
    }

    res += try w.wrapped.writeSplat(data, splat);
    return res;
}

fn flush(io_w: *Writer) Writer.Error!void {
    const w: *NetWriter = @alignCast(@fieldParentPtr("interface", io_w));
    try w.wrapped.flush();
}

const EthernetHeaders = struct {
    dest_mac: @Vector(6, u8),

    src_mac: @Vector(6, u8),

    ether_type: u16,

    fn write(hdr: EthernetHeaders, writer: *std.Io.Writer) Writer.Error!usize {
        comptime var res: usize = 0;

        res += @sizeOf(u48);
        try writer.writeInt(u48, @bitCast(hdr.dest_mac), .big);

        res += @sizeOf(u48);
        try writer.writeInt(u48, @bitCast(hdr.src_mac), .big);

        res += @sizeOf(u16);
        try writer.writeInt(u16, hdr.ether_type, .big);

        return res;
    }

    const byte_len = @bitSizeOf(EthernetHeaders) / 8;

    fn bytes(hdr: EthernetHeaders) [byte_len]u8 {
        var res: [byte_len]u8 = undefined;
        hdr.write(Writer.fixed(&res)) catch unreachable;
    }
};

const IpHeaders = struct {
    // ip_version: u4,
    // header_length: u4 = 0,
    type_of_service: u8 = 0,
    total_length: u16 = 0x04,

    identification: u16 = 0,
    // ethernet_flags: u3 = 0,
    // fragment_offset: u13 = 0,
    ttl: u8 = 0,
    protocol: u8 = 0,

    header_checksum: @Vector(2, u8) = .{ 0, 0 },
    src_ip: @Vector(4, u8),

    dest_ip: @Vector(4, u8),

    fn write(hdr: @This(), writer: *std.Io.Writer) Writer.Error!usize {
        comptime var res: usize = 0;

        res += @sizeOf(u8);
        try writer.writeInt(u8, 0x45, .big); // ip version and header length

        res += @sizeOf(u8);
        try writer.writeByte(hdr.type_of_service);

        res += @sizeOf(u16);
        try writer.writeInt(u16, hdr.total_length, .big);

        res += @sizeOf(u16);
        try writer.writeInt(u16, hdr.identification, .big);

        res += @sizeOf(u16);
        try writer.writeInt(u16, 0x00, .big); // ethernet flags and fragment offset

        res += @sizeOf(u8);
        try writer.writeByte(hdr.ttl);

        res += @sizeOf(u8);
        try writer.writeByte(hdr.protocol);

        res += @sizeOf(u16);
        try writer.writeInt(u16, @bitCast(hdr.header_checksum), .big);

        res += @sizeOf(u32);
        try writer.writeInt(u32, @bitCast(hdr.src_ip), .big);

        res += @sizeOf(u32);
        try writer.writeInt(u32, @bitCast(hdr.dest_ip), .big);

        return res;
    }

    const byte_len = @bitSizeOf(IpHeaders) / 8;

    fn bytes(hdr: IpHeaders) [byte_len]u8 {
        var res: [byte_len]u8 = undefined;
        hdr.write(Writer.fixed(&res)) catch unreachable;
    }
};

const UdpHeaders = packed struct {
    src_port: u16,

    dest_port: u16,
    length: u16,
    checksum: @Vector(2, u8) = .{ 0, 0 },

    fn write(hdr: @This(), writer: *std.Io.Writer) Writer.Error!usize {
        comptime var res: usize = 0;

        res += @sizeOf(u16);
        try writer.writeInt(u16, hdr.src_port, .big);

        res += @sizeOf(u16);
        try writer.writeInt(u16, hdr.dest_port, .big);

        res += @sizeOf(u16);
        try writer.writeInt(u16, hdr.length, .big);

        res += @sizeOf(u16);
        try writer.writeInt(u16, @bitCast(hdr.checksum), .big);

        return res;
    }

    const byte_len = @bitSizeOf(UdpHeaders) / 8;

    fn bytes(hdr: UdpHeaders) [byte_len]u8 {
        var res: [byte_len]u8 = undefined;
        hdr.write(Writer.fixed(&res)) catch unreachable;
    }
};

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

const NetWriter = @This();