const is_debug = builtin.mode == .Debug;

const help =
    \\-h, --help           Display this help and exit.
    \\-r, --relay <str>    A relay message to send.
    \\-d, --dest <str>     An IPv4 or <= 4 ASCII byte string.
    \\-c, --connect <str>  A connection message to send.
    \\
;

const Option = enum { help, relay, dest, connect };
const to_option: StaticStringMap(Option) = .initComptime(.{
    .{ "-h", .help },
    .{ "--help", .help },
    .{ "-r", .relay },
    .{ "--relay", .relay },
    .{ "-d", .dest },
    .{ "--dest", .dest },
    .{ "-c", .connect },
    .{ "--connect", .connect },
});

pub fn main(init: std.process.Init) !void {
    // CLI parsing adapted from the example here
    // https://codeberg.org/ziglang/zig/pulls/30644

    const args = try init.minimal.args.toSlice(init.arena.allocator());

    if (args.len == 1) {
        std.debug.print("{s}", .{help});
        return;
    }

    var flags: struct {
        relay: ?[]const u8 = null,
        dest: ?[]const u8 = null,
        connect: ?[]const u8 = null,
    } = .{};

    {
        var i: usize = 1;
        while (i < args.len) : (i += 1) {
            if (to_option.get(args[i])) |opt| {
                switch (opt) {
                    .help => {
                        std.debug.print("{s}\n", .{help});
                        return;
                    },
                    .relay => {
                        i += 1;
                        if (i < args.len) {
                            flags.relay = args[i];
                        } else {
                            std.debug.print("-r/--relay requires a string\n", .{});
                            return error.InvalidArguments;
                        }
                    },
                    .dest => {
                        i += 1;
                        if (i < args.len) {
                            flags.dest = args[i];
                        } else {
                            std.debug.print("-d/--dest requires a string\n", .{});
                            return error.InvalidArguments;
                        }
                    },
                    .connect => {
                        i += 1;
                        if (i < args.len) {
                            flags.connect = args[i];
                        } else {
                            std.debug.print("-c/--connect requires a string\n", .{});
                            return error.InvalidArguments;
                        }
                    },
                }
            } else {
                std.debug.print("Unknown argument: {s}\n", .{args[i]});
                return error.InvalidArguments;
            }
        }
    }

    if (flags.connect != null and (flags.relay != null or flags.dest != null)) {
        std.debug.print("Incompatible arguments.\nCannot use --connect/-c with dest or relay.\n", .{});
        return error.InvalidArguments;
    }

    // std.debug.print("relay: {s}\n", .{flags.relay orelse "<null>"});
    // std.debug.print("dest: {s}\n", .{flags.dest orelse "<null>"});
    // std.debug.print("connect: {s}\n", .{flags.connect orelse "<null>"});

    const rand = blk: {
        const io_source: std.Random.IoSource = .{ .io = init.io };
        break :blk io_source.interface();
    };

    // const net_interface: std.Io.net.Interface = .{ .index = 1 };
    // std.debug.print("Interface: {s}\n", .{(try net_interface.name(init.io)).toSlice()});
    const EthIpUdp = packed struct(u336) { // 42 bytes * 8 bits = 336
        // --- UDP (Last in memory, defined first for LSB->MSB) ---
        udp: packed struct {
            checksum: u16 = 0,
            len: u16,
            dst_port: u16,
            src_port: u16,
        },

        // --- IP ---
        ip: packed struct {
            dst_addr: u32,
            src_addr: u32,
            header_checksum: u16 = 0,
            protocol: u8 = 17, // udp
            ttl: u8 = 0x40,

            // fragment_offset (13 bits) + flags (3 bits) = 16 bits
            // In Big Endian, flags are the high bits of the first byte.
            // To have flags appear first in the stream, define them last here.
            fragment_offset: u13 = 0,
            flags: packed struct(u3) {
                reserved: u1 = 0,
                dont_fragment: u1 = 0,
                more_fragments: u1 = 0,
            } = .{},

            id: u16,
            len: u16,
            tos: u8 = undefined,

            // ip_version (4 bits) + ihl (4 bits) = 8 bits
            // To have version appear first (high nibble), define it last.
            ihl: u4 = 5,
            ip_version: u4 = 4,
        },

        // --- Ethernet ---
        eth_type: u16 = std.os.linux.ETH.P.IP,
        src_mac: @Vector(6, u8),
        dst_mac: @Vector(6, u8) = @splat(0xff),

        fn toBytes(self: @This()) [336 / 8]u8 {
            var res: [336 / 8]u8 = undefined;
            var w: Writer = .fixed(&res);
            w.writeStruct(self, .big) catch unreachable;
            return res;
        }

        fn setPayloadLen(self: *@This(), len: usize) void {
            self.ip.len = @intCast(len + (@bitSizeOf(@TypeOf(self.udp)) / 8) + (@bitSizeOf(@TypeOf(self.ip)) / 8));
            self.udp.len = @intCast(len + (@bitSizeOf(@TypeOf(self.udp)) / 8));
        }
    };

    var socket: RawSocket = try .init();
    defer socket.deinit();

    var headers: EthIpUdp = .{
        .src_mac = socket.mac,
        .ip = .{
            .id = 0,
            .src_addr = rand.int(u32),
            .dst_addr = @bitCast([_]u8{ 255, 255, 255, 255 }),
            .len = undefined,
        },
        .udp = .{
            .src_port = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16)),
            .dst_port = 8888,
            .len = undefined,
        },
    };

    if (flags.relay != null) {
        const relay: SaprusMessage = .{
            .relay = .{
                .dest = .fromBytes(&parseDest(flags.dest)),
                .payload = flags.relay.?,
            },
        };

        var relay_buf: [2048]u8 = undefined;
        const relay_bytes = relay.toBytes(&relay_buf);
        headers.setPayloadLen(relay_bytes.len);

        const full_msg = blk: {
            var msg_buf: [2048]u8 = undefined;
            var msg_w: Writer = .fixed(&msg_buf);
            msg_w.writeAll(&headers.toBytes()) catch unreachable;
            msg_w.writeAll(relay_bytes) catch unreachable;
            break :blk msg_w.buffered();
        };

        try socket.send(full_msg);
        return;
    }

    if (flags.connect != null) {
        const dest = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16));
        const src = rand.intRangeAtMost(u16, 1025, std.math.maxInt(u16));
        // udp dest port should not be 8888 after first
        const udp_dest_port = rand.intRangeAtMost(u16, 9000, std.math.maxInt(u16));
        const connection: SaprusMessage = .{
            .connection = .{
                .src = src,
                .dest = dest,
                .seq = undefined,
                .id = undefined,
                .payload = flags.connect.?,
            },
        };

        try socket.attachSaprusPortFilter(src);

        var connection_buf: [2048]u8 = undefined;
        const connection_bytes = connection.toBytes(&connection_buf);
        headers.setPayloadLen(connection_bytes.len);

        var full_msg = blk: {
            var msg_buf: [2048]u8 = undefined;
            var msg_w: Writer = .fixed(&msg_buf);
            msg_w.writeAll(&headers.toBytes()) catch unreachable;
            msg_w.writeAll(connection_bytes) catch unreachable;
            break :blk msg_w.buffered();
        };

        try socket.send(full_msg);
        var res_buf: [4096]u8 = undefined;

        const res = try socket.receive(&res_buf);
        std.debug.print("response: {any}\n", .{res});

        headers.udp.dst_port = udp_dest_port;

        full_msg = blk: {
            var msg_buf: [2048]u8 = undefined;
            var msg_w: Writer = .fixed(&msg_buf);
            msg_w.writeAll(&headers.toBytes()) catch unreachable;
            msg_w.writeAll(connection_bytes) catch unreachable;
            break :blk msg_w.buffered();
        };
        try socket.send(full_msg);

        return;
    }

    unreachable;
}

fn parseDest(in: ?[]const u8) [4]u8 {
    if (in) |dest| {
        if (dest.len <= 4) {
            var res: [4]u8 = @splat(0);
            @memcpy(res[0..dest.len], dest);
            return res;
        }

        const addr = std.Io.net.Ip4Address.parse(dest, 0) catch return "FAIL".*;
        return addr.bytes;
    }
    return "disc".*;
}

const builtin = @import("builtin");
const std = @import("std");
const ArrayList = std.ArrayList;
const StaticStringMap = std.StaticStringMap;

const zaprus = @import("zaprus");
const SaprusClient = zaprus.Client;
const SaprusMessage = zaprus.Message;
const RawSocketWriter = zaprus.RawSocketWriter;

const AF = std.os.linux.AF;
const SOCK = std.os.linux.SOCK;

const RawSocket = struct {
    fd: i32,
    sockaddr_ll: std.posix.sockaddr.ll,
    mac: [6]u8,

    const IFF_LOOPBACK = 0x8;

    const ifconf = extern struct {
        ifc_len: i32,
        ifc_ifcu: extern union {
            ifcu_buf: ?[*]u8,
            ifcu_req: ?[*]std.os.linux.ifreq,
        },
    };

    fn init() !RawSocket {
        const socket: i32 = @intCast(std.os.linux.socket(std.posix.AF.PACKET, std.posix.SOCK.RAW, 0));
        if (socket < 0) return error.SocketError;

        // 1. Enumerate interfaces
        var ifreq_storage: [16]std.os.linux.ifreq = undefined;
        var ifc = ifconf{
            .ifc_len = @sizeOf(@TypeOf(ifreq_storage)),
            .ifc_ifcu = .{ .ifcu_req = &ifreq_storage },
        };

        if (std.os.linux.ioctl(socket, std.os.linux.SIOCGIFCONF, @intFromPtr(&ifc)) != 0) {
            return error.NicError;
        }

        const count = @divExact(ifc.ifc_len, @sizeOf(std.os.linux.ifreq));
        var target_ifr: ?std.os.linux.ifreq = null;

        for (ifreq_storage[0..@intCast(count)]) |ifr| {
            var temp_ifr = ifr;
            if (std.os.linux.ioctl(socket, std.os.linux.SIOCGIFFLAGS, @intFromPtr(&temp_ifr)) == 0) {
                // Cast the packed flags to u16 to match the kernel's ifr_flags size
                const flags: u16 = @bitCast(temp_ifr.ifru.flags);
                if (flags & IFF_LOOPBACK != 0) continue;

                target_ifr = ifr;
                break;
            }
        }

        var ifr = target_ifr orelse return error.NoInterfaceFound;

        // 2. Get Interface Index
        if (std.os.linux.ioctl(socket, std.os.linux.SIOCGIFINDEX, @intFromPtr(&ifr)) != 0) {
            return error.NicError;
        }
        const ifindex: i32 = ifr.ifru.ivalue;

        // 3. Get Real MAC Address
        if (std.os.linux.ioctl(socket, std.os.linux.SIOCGIFHWADDR, @intFromPtr(&ifr)) != 0) {
            return error.NicError;
        }
        var mac: [6]u8 = ifr.ifru.hwaddr.data[0..6].*;
        std.mem.reverse(u8, &mac);

        // 4. Set Flags (Promiscuous/Broadcast)
        if (std.os.linux.ioctl(socket, std.os.linux.SIOCGIFFLAGS, @intFromPtr(&ifr)) != 0) {
            return error.NicError;
        }
        ifr.ifru.flags.BROADCAST = true;
        ifr.ifru.flags.PROMISC = true;
        if (std.os.linux.ioctl(socket, std.os.linux.SIOCSIFFLAGS, @intFromPtr(&ifr)) != 0) {
            return error.NicError;
        }

        const sockaddr_ll = std.posix.sockaddr.ll{
            .family = std.posix.AF.PACKET,
            .ifindex = ifindex,
            .protocol = std.mem.nativeToBig(u16, @as(u16, std.os.linux.ETH.P.IP)),
            .halen = 0,
            .addr = .{ 0, 0, 0, 0, 0, 0, 0, 0 },
            .pkttype = 0,
            .hatype = 0,
        };

        const bind_ret = std.os.linux.bind(socket, @ptrCast(&sockaddr_ll), @sizeOf(@TypeOf(sockaddr_ll)));
        if (bind_ret != 0) return error.BindError;

        return .{
            .fd = socket,
            .sockaddr_ll = sockaddr_ll,
            .mac = mac,
        };
    }

    fn deinit(self: *RawSocket) void {
        _ = std.os.linux.close(self.fd);
        self.* = undefined;
    }

    fn send(self: RawSocket, payload: []const u8) !void {
        const sent_bytes = std.os.linux.sendto(
            self.fd,
            payload.ptr,
            payload.len,
            0,
            @ptrCast(&self.sockaddr_ll),
            @sizeOf(@TypeOf(self.sockaddr_ll)),
        );
        std.debug.assert(sent_bytes == payload.len);
    }

    fn receive(self: RawSocket, buf: []u8) ![]u8 {
        const len = std.os.linux.recvfrom(
            self.fd,
            buf.ptr,
            buf.len,
            0,
            null,
            null,
        );
        return buf[0..len];
    }

    fn attachSaprusPortFilter(self: RawSocket, port: u16) !void {
        const linux = std.os.linux;
        // BPF instruction structure for classic BPF
        const sock_filter = extern struct {
            code: u16,
            jt: u8,
            jf: u8,
            k: u32,
        };

        const sock_fprog = extern struct {
            len: u16,
            filter: [*]const sock_filter,
        };

        // BPF instruction opcodes
        const BPF_LD = 0x00;
        const BPF_H = 0x08; // Half-word (2 bytes)
        const BPF_ABS = 0x20;
        const BPF_JMP = 0x05;
        const BPF_JEQ = 0x10;
        const BPF_K = 0x00;
        const BPF_RET = 0x06;

        // Build the filter program
        const filter = [_]sock_filter{
            // Load 2 bytes at offset 46 (absolute)
            .{ .code = BPF_LD | BPF_H | BPF_ABS, .jt = 0, .jf = 0, .k = 46 },
            // Jump if equal to port (skip 0 if true, skip 1 if false)
            .{ .code = BPF_JMP | BPF_JEQ | BPF_K, .jt = 0, .jf = 1, .k = @as(u32, port) },
            // Return 0xffff (pass)
            .{ .code = BPF_RET | BPF_K, .jt = 0, .jf = 0, .k = 0xffff },
            // Return 0x0 (fail)
            .{ .code = BPF_RET | BPF_K, .jt = 0, .jf = 0, .k = 0x0 },
        };

        const fprog = sock_fprog{
            .len = filter.len,
            .filter = &filter,
        };

        // Attach filter to socket using setsockopt
        const SO_ATTACH_FILTER = 26;
        const rc = linux.setsockopt(
            self.fd,
            linux.SOL.SOCKET,
            SO_ATTACH_FILTER,
            @ptrCast(&fprog),
            @sizeOf(sock_fprog),
        );

        if (rc != 0) {
            return error.BpfAttachFailed;
        }
    }
};

const Writer = std.Io.Writer;