//! 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, buffer: []u8) !NetWriter {
    std.debug.assert(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 = buffer,
        },
    };
}

fn drain(io_w: *Writer, data: []const []const u8, splat: usize) Writer.Error!usize {
    const w: *NetWriter = @alignCast(@fieldParentPtr("interface", io_w));
    const headers_byte_len = comptime (EthernetHeaders.byte_len + IpHeaders.byte_len + UdpHeaders.byte_len);
    const headers: [headers_byte_len]u8 = blk: {
        const ether_headers: EthernetHeaders = .{
            .dest_mac = .{ 0xff, 0xff, 0xff, 0xff, 0xfe, 0xff },
            .src_mac = src_blk: {
                var output_bytes: [6]u8 = undefined;
                output_bytes[0] = 0xee;
                w.rand.bytes(output_bytes[1..]);
                break :src_blk output_bytes;
            },
            .ether_type = 0x0800,
        };

        const total_len = Writer.countSplat(data, splat) + w.interface.end;

        const ip_headers: IpHeaders = .{
            // length of the packet minus eth header
            .total_length = @intCast(headers_byte_len + total_len - EthernetHeaders.byte_len), //@intCast(total_len),
            .ttl = 64,
            .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(total_len + UdpHeaders.byte_len),
        };

        var buf: [headers_byte_len]u8 = undefined;
        var buf_w = Writer.fixed(&buf);

        _ = try ether_headers.write(&buf_w);
        _ = try ip_headers.write(&buf_w);
        std.debug.print("after ip: {x}\n", .{buf_w.buffered()});
        _ = try udp_headers.write(&buf_w);
        std.debug.print("after udp: {x}\n", .{buf_w.buffered()});

        break :blk buf;
    };

    _ = try w.wrapped.write(&headers);
    const total_len = try w.wrapped.writeSplatHeader(w.interface.buffered(), data, splat);

    std.debug.print("total splat: {}\theaders.len: {}\tsplat: {}\n", .{ total_len, headers.len, splat });

    try w.wrapped.flush();
    return total_len - w.interface.consumeAll();
}

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

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 {
        try writer.writeInt(u48, @bitCast(hdr.dest_mac), .big);
        try writer.writeInt(u48, @bitCast(hdr.src_mac), .big);
        try writer.writeInt(u16, hdr.ether_type, .big);
        return byte_len;
    }

    const byte_len = blk: {
        var res: usize = 0;
        res += @bitSizeOf(u48) / 8;
        res += @bitSizeOf(u48) / 8;
        res += @bitSizeOf(u16) / 8;
        break :blk res;
    };

    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 {
        try writer.writeInt(u8, 0x45, .big); // ip version and header length
        try writer.writeByte(hdr.type_of_service);
        try writer.writeInt(u16, hdr.total_length, .big);
        try writer.writeInt(u16, hdr.identification, .big);
        try writer.writeInt(u16, 0x00, .big); // ethernet flags and fragment offset
        try writer.writeByte(hdr.ttl);
        try writer.writeByte(hdr.protocol);
        try writer.writeInt(u16, @bitCast(hdr.header_checksum), .big);
        try writer.writeInt(u32, @bitCast(hdr.src_ip), .big);
        try writer.writeInt(u32, @bitCast(hdr.dest_ip), .big);
        return byte_len;
    }

    const byte_len = blk: {
        var res: usize = 0;
        res += @sizeOf(u8);
        res += @sizeOf(u8);
        res += @sizeOf(u16);
        res += @sizeOf(u16);
        res += @sizeOf(u16);
        res += @sizeOf(u8);
        res += @sizeOf(u8);
        res += @sizeOf(u16);
        res += @sizeOf(u32);
        res += @sizeOf(u32);
        break :blk res;
    };

    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 {
        try writer.writeInt(u16, hdr.src_port, .big);
        try writer.writeInt(u16, hdr.dest_port, .big);
        try writer.writeInt(u16, hdr.length, .big);
        try writer.writeInt(u16, @bitCast(hdr.checksum), .big);
        return byte_len;
    }

    const byte_len = blk: {
        var res: usize = 0;
        res += @sizeOf(u16);
        res += @sizeOf(u16);
        res += @sizeOf(u16);
        res += @sizeOf(u16);
        break :blk res;
    };

    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();