const std = @import("std");

const xev = @import("xev");

const RotInt = @import("./main.zig").RotInt;
const rotctl = @import("./rotctl.zig");

pub const Server = struct {
    // rotint initializes itself
    rotint: *RotInt = undefined,
    listen_socket: xev.TCP = undefined,
    client_socket: ?xev.TCP = null,
    read_completion: xev.Completion = undefined,
    write_completion: xev.Completion = undefined,
    // our messages are small. We may be able to get away with an even smaller buffer
    // size than this.
    read_buffer: [128]u8 = undefined,
    write_buffer: [128]u8 = undefined,
    should_disconnect: bool = false,

    pub fn listen(self: *Server, loop: *xev.Loop, address: std.net.Address) !void {
        self.listen_socket = try xev.TCP.init(address);
        try self.listen_socket.bind(address);
        try self.listen_socket.listen(1);
        self.listen_socket.accept(loop, &self.read_completion, Server, self, acceptCallback);
    }

    pub fn respond(self: *Server, loop: *xev.Loop, response: rotctl.RotReply) void {
        const rep = response.write(self.write_buffer[0..]) catch {
            self.rotint.warn("could not write response {s}", .{@tagName(response)});
            return;
        };
        if (self.client_socket) |socket|
            socket.write(
                loop,
                &self.write_completion,
                .{ .slice = rep },
                Server,
                self,
                writeCallback,
            );
    }

    fn acceptCallback(
        self_: ?*Server,
        loop: *xev.Loop,
        completion: *xev.Completion,
        result: xev.TCP.AcceptError!xev.TCP,
    ) xev.CallbackAction {
        _ = completion;

        const self = self_.?;

        self.client_socket = result catch {
            // different error sets for different backends makes handling individual
            // errors here odd. Either way, if this fails, I think we just want to
            // retry anyway (maybe this should be a fatal error instead?).
            return .rearm;
        };

        self.should_disconnect = false;
        // this is already using read_completion, but pass it explicitly here to be clear
        self.client_socket.?.read(
            loop,
            &self.read_completion,
            .{ .slice = self.read_buffer[0..] },
            Server,
            self,
            readCallback,
        );
        return .disarm;
    }

    fn readCallback(
        self_: ?*Server,
        loop: *xev.Loop,
        completion: *xev.Completion,
        socket: xev.TCP,
        buf: xev.ReadBuffer,
        result: xev.TCP.ReadError!usize,
    ) xev.CallbackAction {
        _ = socket;
        const self = self_.?;
        const n = result catch |err| switch (err) {
            // we can reuse the completion here in either of these error cases since we
            // disarm in both of them
            error.EOF => {
                self.disconnectClient(loop, completion);
                return .disarm;
            },
            else => {
                self.rotint.warn("server read unexpected err={}", .{err});
                // TODO: if read failed unexpectedly, will shutting down the socket like
                // this work?
                self.disconnectClient(loop, completion);
                return .disarm;
            },
        };

        // this is a request from gpredict to rotctl. We get rotint to take care of it
        // for us (it stores a copy of this buffer)
        self.rotint.handleControlRequest(buf.slice[0..n]);
        // we could disarm until we have generated and sent our response....
        return .rearm;
    }

    fn writeCallback(
        self_: ?*Server,
        loop: *xev.Loop,
        completion: *xev.Completion,
        socket: xev.TCP,
        buf: xev.WriteBuffer,
        result: xev.TCP.WriteError!usize,
    ) xev.CallbackAction {
        _ = socket;
        _ = buf;

        const self = self_.?;
        _ = result catch |err| {
            // we're in trouble. Probably this is due to an unexpected client disconnect?
            self.rotint.warn("server write unexpected err={}", .{err});
            self.disconnectClient(loop, completion);
            return .disarm;
        };

        if (self.should_disconnect) {
            self.disconnectClient(loop, completion);
            return .disarm;
        }

        return .disarm;
    }

    fn disconnectClient(self: *Server, loop: *xev.Loop, completion: *xev.Completion) void {
        if (self.client_socket) |sock| {
            sock.shutdown(loop, completion, Server, self, shutdownCallback);
            self.client_socket = null;
        }
    }

    fn shutdownCallback(
        self_: ?*Server,
        loop: *xev.Loop,
        completion: *xev.Completion,
        socket: xev.TCP,
        result: xev.TCP.ShutdownError!void,
    ) xev.CallbackAction {
        _ = result catch {};
        // shutdown only closes the write side of the socket. It's not entirely clear to
        // me if shutdown actually needs to be called before close.
        // This can be called with either the read or write completion.
        socket.close(loop, completion, Server, self_.?, closeCallback);
        return .disarm;
    }

    fn closeCallback(
        self_: ?*Server,
        loop: *xev.Loop,
        completion: *xev.Completion,
        socket: xev.TCP,
        result: xev.TCP.CloseError!void,
    ) xev.CallbackAction {
        _ = completion;
        _ = socket;
        _ = result catch unreachable;

        // listen for a new connection
        const self = self_.?;
        self.listen_socket.accept(loop, &self.read_completion, Server, self, acceptCallback);

        return .disarm;
    }
};

pub const Client = struct {
    // rotint initializes itself. could use @fieldParentPtr("rotator", self) also
    rotint: *RotInt = undefined,
    socket: xev.TCP = undefined,
    read_completion: xev.Completion = undefined,
    write_completion: xev.Completion = undefined,
    read_buffer: [128]u8 = undefined,
    write_buffer: [128]u8 = undefined,

    pub fn connect(self: *Client, loop: *xev.Loop, address: std.net.Address) !void {
        self.socket = try xev.TCP.init(address);
        self.socket.connect(loop, &self.read_completion, address, Client, self, connectCallback);
    }

    pub fn sendCommand(self: *Client, loop: *xev.Loop, command: rotctl.RotCommand) void {
        const req = command.write(self.write_buffer[0..]) catch {
            self.rotint.warn("could not write command {s}", .{@tagName(command)});
            return;
        };
        self.socket.write(
            loop,
            &self.write_completion,
            .{ .slice = req },
            Client,
            self,
            writeCallback,
        );
    }

    fn disconnect(self: *Client, loop: *xev.Loop, completion: *xev.Completion) void {
        self.socket.shutdown(loop, completion, Client, self, shutdownCallback);
    }

    fn connectCallback(
        self_: ?*Client,
        loop: *xev.Loop,
        completion: *xev.Completion,
        socket: xev.TCP,
        result: xev.TCP.ConnectError!void,
    ) xev.CallbackAction {
        _ = completion;

        const self = self_.?;

        _ = result catch |err| {
            self.rotint.warn("connect failed {}", .{err});
            // Note: this loops instantaneously if the target address is on localhost
            // and is not being listened on. Actually just straight up crashes on macOS.
            // Probably need to use a timer to handle this better
            return .rearm;
        };

        self.rotint.stateEvent(.rotator_connected);
        socket.read(
            loop,
            &self.read_completion,
            .{ .slice = self.read_buffer[0..] },
            Client,
            self,
            readCallback,
        );
        return .disarm;
    }

    fn writeCallback(
        self_: ?*Client,
        loop: *xev.Loop,
        completion: *xev.Completion,
        socket: xev.TCP,
        buf: xev.WriteBuffer,
        result: xev.TCP.WriteError!usize,
    ) xev.CallbackAction {
        const self = self_.?;
        _ = result catch {
            self.disconnect(loop, completion);
            return .disarm;
        };
        _ = socket;
        _ = buf;

        return .disarm;
    }

    fn readCallback(
        self_: ?*Client,
        loop: *xev.Loop,
        completion: *xev.Completion,
        socket: xev.TCP,
        buf: xev.ReadBuffer,
        result: xev.TCP.ReadError!usize,
    ) xev.CallbackAction {
        _ = socket;

        const self = self_.?;
        const n = result catch {
            self.disconnect(loop, completion);
            return .disarm;
        };

        self.rotint.handleRotatorReply(buf.slice[0..n]);
        return .rearm;
    }

    fn shutdownCallback(
        self_: ?*Client,
        loop: *xev.Loop,
        completion: *xev.Completion,
        socket: xev.TCP,
        result: xev.TCP.ShutdownError!void,
    ) xev.CallbackAction {
        const self = self_.?;
        _ = result catch {};
        socket.close(loop, completion, Client, self, closeCallback);
        return .disarm;
    }

    fn closeCallback(
        self_: ?*Client,
        loop: *xev.Loop,
        completion: *xev.Completion,
        socket: xev.TCP,
        result: xev.TCP.CloseError!void,
    ) xev.CallbackAction {
        _ = self_;
        _ = loop;
        _ = completion;
        _ = socket;
        _ = result catch {};
        return .disarm;
    }
};
