const std = @import("std");

pub const MessageType = enum {
    info,
    connect,
    @"pub",
    hpub,
    sub,
    unsub,
    msg,
    hmsg,
    ping,
    pong,
    @"+ok",
    @"-err",

    fn parseMemEql(input: []const u8) ?MessageType {
        // if (std.mem.eql(u8, "INFO", input)) return .info;
        if (std.mem.eql(u8, "CONNECT", input)) return .connect;
        if (std.mem.eql(u8, "PUB", input)) return .@"pub";
        if (std.mem.eql(u8, "HPUB", input)) return .hpub;
        if (std.mem.eql(u8, "SUB", input)) return .sub;
        if (std.mem.eql(u8, "UNSUB", input)) return .unsub;
        // if (std.mem.eql(u8, "MSG", input)) return .msg;
        // if (std.mem.eql(u8, "HMSG", input)) return .hmsg;
        if (std.mem.eql(u8, "PING", input)) return .ping;
        if (std.mem.eql(u8, "PONG", input)) return .pong;
        // if (std.mem.eql(u8, "@"+OK"", input)) return .@"+ok";
        // if (std.mem.eql(u8, "@"-ERR"", input)) return .@"-err";
        return error.InvalidMessageType;
    }
};

pub const Message = union(MessageType) {
    info: ServerInfo,
    connect: Connect,
    @"pub": Pub,
    hpub: void,
    sub: Sub,
    unsub: Unsub,
    msg: Msg,
    hmsg: void,
    ping,
    pong,
    @"+ok": void,
    @"-err": []const u8,
    pub const ServerInfo = struct {
        /// The unique identifier of the NATS server.
        server_id: []const u8,
        /// The name of the NATS server.
        server_name: []const u8,
        /// The version of NATS.
        version: []const u8,
        /// The version of golang the NATS server was built with.
        go: []const u8 = "0.0.0",
        /// The IP address used to start the NATS server,
        /// by default this will be 0.0.0.0 and can be
        /// configured with -client_advertise host:port.
        host: []const u8 = "0.0.0.0",
        /// The port number the NATS server is configured
        /// to listen on.
        port: u16 = 4222,
        /// Whether the server supports headers.
        headers: bool = false,
        /// Maximum payload size, in bytes, that the server
        /// will accept from the client.
        max_payload: u64,
        /// An integer indicating the protocol version of
        /// the server. The server version 1.2.0 sets this
        /// to 1 to indicate that it supports the "Echo"
        /// feature.
        proto: u32 = 1,
    };
    pub const Connect = struct {
        verbose: bool = false,
        pedantic: bool = false,
        tls_required: bool = false,
        auth_token: ?[]const u8 = null,
        user: ?[]const u8 = null,
        pass: ?[]const u8 = null,
        name: ?[]const u8 = null,
        lang: []const u8,
        version: []const u8,
        protocol: u32,
        echo: ?bool = null,
        sig: ?[]const u8 = null,
        jwt: ?[]const u8 = null,
        no_responders: ?bool = null,
        headers: ?bool = null,
        nkey: ?[]const u8 = null,

        pub fn deinit(self: Connect, alloc: std.mem.Allocator) void {
            if (self.auth_token) |a| alloc.free(a);
            if (self.user) |u| alloc.free(u);
            if (self.pass) |p| alloc.free(p);
            if (self.name) |n| alloc.free(n);
            alloc.free(self.lang);
            alloc.free(self.version);
            if (self.sig) |s| alloc.free(s);
            if (self.jwt) |j| alloc.free(j);
            if (self.nkey) |n| alloc.free(n);
        }

        pub fn dupe(self: Connect, alloc: std.mem.Allocator) !Connect {
            var res = self;
            res.auth_token = if (self.auth_token) |a| try alloc.dupe(u8, a) else null;
            errdefer if (res.auth_token) |a| alloc.free(a);
            res.user = if (self.user) |u| try alloc.dupe(u8, u) else null;
            errdefer if (res.user) |u| alloc.free(u);
            res.pass = if (self.pass) |p| try alloc.dupe(u8, p) else null;
            errdefer if (res.pass) |p| alloc.free(p);
            res.name = if (self.name) |n| try alloc.dupe(u8, n) else null;
            errdefer if (res.name) |n| alloc.free(n);
            res.lang = try alloc.dupe(u8, self.lang);
            errdefer alloc.free(res.lang);
            res.version = try alloc.dupe(u8, self.version);
            errdefer alloc.free(res.version);
            res.sig = if (self.sig) |s| try alloc.dupe(u8, s) else null;
            errdefer if (res.sig) |s| alloc.free(s);
            res.jwt = if (self.jwt) |j| try alloc.dupe(u8, j) else null;
            errdefer if (res.jwt) |j| alloc.free(j);
            res.nkey = if (self.nkey) |n| try alloc.dupe(u8, n) else null;
            errdefer if (res.nkey) |n| alloc.free(n);
            return res;
        }
    };
    pub const Pub = struct {
        /// The destination subject to publish to.
        subject: []const u8,
        /// The reply subject that subscribers can use to send a response back to the publisher/requestor.
        reply_to: ?[]const u8 = null,
        /// The message payload data.
        payload: []const u8,

        pub fn deinit(self: Pub, alloc: std.mem.Allocator) void {
            alloc.free(self.subject);
            alloc.free(self.payload);
            if (self.reply_to) |r| alloc.free(r);
        }
    };
    pub const Sub = struct {
        /// The subject name to subscribe to.
        subject: []const u8,
        /// If specified, the subscriber will join this queue group.
        queue_group: ?[]const u8,
        /// A unique alphanumeric subscription ID, generated by the client.
        sid: []const u8,

        pub fn deinit(self: Sub, alloc: std.mem.Allocator) void {
            alloc.free(self.subject);
            alloc.free(self.sid);
            if (self.queue_group) |q| alloc.free(q);
        }
    };
    pub const Unsub = struct {
        /// The unique alphanumeric subscription ID of the subject to unsubscribe from.
        sid: []const u8,
        /// A number of messages to wait for before automatically unsubscribing.
        max_msgs: ?usize = null,
    };
    pub const Msg = struct {
        subject: []const u8,
        sid: []const u8,
        reply_to: ?[]const u8,
        payload: []const u8,

        pub fn deinit(self: Msg, alloc: std.mem.Allocator) void {
            alloc.free(self.subject);
            alloc.free(self.sid);
            if (self.reply_to) |r| alloc.free(r);
            alloc.free(self.payload);
        }

        pub fn dupe(self: Msg, alloc: std.mem.Allocator) !Msg {
            var res: Msg = undefined;
            res.subject = try alloc.dupe(u8, self.subject);
            errdefer alloc.free(res.subject);
            res.sid = try alloc.dupe(u8, self.sid);
            errdefer alloc.free(res.sid);
            res.reply_to = if (self.reply_to) |r| alloc.dupe(u8, r) else null;
            errdefer if (res.reply_to) |r| alloc.free(r);
            res.payload = try alloc.dupe(u8, self.payload);
            errdefer alloc.free(res.payload);
            return res;
        }
    };

    const client_types = std.StaticStringMap(MessageType).initComptime(
        .{
            // {"INFO", .info},
            .{ "CONNECT", .connect },
            .{ "PUB", .@"pub" },
            .{ "HPUB", .hpub },
            .{ "SUB", .sub },
            .{ "UNSUB", .unsub },
            // {"MSG", .msg},
            // {"HMSG", .hmsg},
            .{ "PING", .ping },
            .{ "PONG", .pong },
            // {"+OK", .@"+ok"},
            // {"-ERR", .@"-err"},
        },
    );
    fn parseStaticStringMap(input: []const u8) ?MessageType {
        return client_types.get(input);
    }

    pub const parse = parseStaticStringMap;

    /// An error should be handled by cleaning up this connection.
    pub fn next(alloc: std.mem.Allocator, in: *std.Io.Reader) !Message {
        var operation_string: std.ArrayList(u8) = blk: {
            var buf: ["CONTINUE".len + 1]u8 = undefined;
            break :blk .initBuffer(&buf);
        };

        while (in.peekByte()) |byte| {
            if (std.ascii.isUpper(byte)) {
                try operation_string.appendBounded(byte);
                in.toss(1);
            } else break;
        } else |err| return err;

        const operation = parse(operation_string.items) orelse {
            return error.InvalidOperation;
        };

        switch (operation) {
            .connect => {
                // for storing the json string
                var connect_string_writer_allocating: std.Io.Writer.Allocating = try .initCapacity(alloc, 1024);
                defer connect_string_writer_allocating.deinit();
                var connect_string_writer = connect_string_writer_allocating.writer;

                // for parsing the json string
                var connect_arena_allocator: std.heap.ArenaAllocator = .init(alloc);
                defer connect_arena_allocator.deinit();
                const connect_allocator = connect_arena_allocator.allocator();

                try in.discardAll(1); // throw away space

                // Should read the next JSON object to the fixed buffer writer.
                _ = try in.streamDelimiter(&connect_string_writer, '}');
                try connect_string_writer.writeByte('}');
                try expectStreamBytes(in, "}\r\n"); // discard '}\r\n'

                // TODO: should be CONNECTION allocator
                const res = try std.json.parseFromSliceLeaky(
                    Connect,
                    connect_allocator,
                    connect_string_writer.buffered(),
                    .{ .allocate = .alloc_always },
                );

                return .{ .connect = try res.dupe(alloc) };
            },
            .@"pub" => {
                try in.discardAll(1); // throw away space

                // Parse subject
                const subject: []const u8 = try readSubject(alloc, in);
                errdefer alloc.free(subject);

                // Parse byte count
                const byte_count = blk: {
                    var byte_count_list: std.ArrayList(u8) = try .initCapacity(alloc, 64);
                    defer byte_count_list.deinit(alloc);
                    while (in.peekByte()) |byte| {
                        if (std.ascii.isWhitespace(byte)) {
                            try expectStreamBytes(in, "\r\n");
                            break;
                        }
                        defer in.toss(1);

                        if (std.ascii.isDigit(byte)) {
                            try byte_count_list.append(alloc, byte);
                        } else {
                            return error.InvalidStream;
                        }
                    } else |err| return err;

                    break :blk try std.fmt.parseUnsigned(usize, byte_count_list.items, 10);
                };

                const payload = blk: {
                    const bytes = try alloc.alloc(u8, byte_count);
                    errdefer alloc.free(bytes);
                    try in.readSliceAll(bytes);
                    try expectStreamBytes(in, "\r\n");
                    break :blk bytes;
                };

                return .{
                    .@"pub" = .{
                        .subject = subject,
                        .payload = payload,
                    },
                };
            },
            .ping => {
                try expectStreamBytes(in, "\r\n");
                return .ping;
            },
            .pong => {
                try expectStreamBytes(in, "\r\n");
                return .pong;
            },
            .sub => {
                if (!std.ascii.isWhitespace(try in.takeByte())) {
                    @branchHint(.unlikely);
                    return error.InvalidStream;
                }
                const subject = try readSubject(alloc, in);
                errdefer alloc.free(subject);
                const second = blk: {
                    // Drop whitespace
                    while (in.peekByte()) |byte| {
                        if (std.ascii.isWhitespace(byte)) {
                            in.toss(1);
                        } else break;
                    } else |err| return err;

                    var acc: std.ArrayList(u8) = try .initCapacity(alloc, 32);
                    errdefer acc.deinit(alloc);
                    while (in.peekByte()) |byte| {
                        if (std.ascii.isWhitespace(byte)) break;
                        try acc.append(alloc, byte);
                        in.toss(1);
                    } else |err| return err;

                    break :blk try acc.toOwnedSlice(alloc);
                };
                errdefer alloc.free(second);
                const queue_group = if ((try in.peekByte()) != '\r') second else null;
                // We do not need an errdefer free for queue group, because it will only be second (already has errdefer free) or null.
                const sid = if (queue_group) |_| try alloc.dupe(u8, try in.takeDelimiterExclusive('\r')) else second;
                // if queue_group is null, that means sid is second, and already has an errdefer free.
                errdefer if (queue_group) |_| alloc.free(sid);
                try expectStreamBytes(in, "\r\n");
                return .{
                    .sub = .{
                        .subject = subject,
                        .queue_group = queue_group,
                        .sid = sid,
                    },
                };
            },
            .unsub => {
                if (!std.ascii.isWhitespace(try in.takeByte())) {
                    @branchHint(.unlikely);
                    return error.InvalidStream;
                }
                // Parse sid
                const sid = blk: {
                    var acc: std.ArrayList(u8) = try .initCapacity(alloc, 8);
                    errdefer acc.deinit(alloc);
                    while (in.peekByte()) |byte| {
                        if (std.ascii.isWhitespace(byte)) break;
                        try acc.append(alloc, byte);
                        in.toss(1);
                    } else |err| return err;
                    break :blk try acc.toOwnedSlice(alloc);
                };
                errdefer alloc.free(sid);

                if ((try in.peekByte()) == '\r') {
                    try expectStreamBytes(in, "\r\n");
                    return .{
                        .unsub = .{
                            .sid = sid,
                        },
                    };
                } else if (std.ascii.isWhitespace(try in.peekByte())) {
                    in.toss(1);
                    const max_msgs = blk: {
                        var max_msgs_list: std.ArrayList(u8) = try .initCapacity(alloc, 64);
                        errdefer max_msgs_list.deinit(alloc);
                        while (in.peekByte()) |byte| {
                            if (std.ascii.isWhitespace(byte)) {
                                try expectStreamBytes(in, "\r\n");
                                break;
                            }

                            if (std.ascii.isDigit(byte)) {
                                try max_msgs_list.append(alloc, byte);
                            } else {
                                return error.InvalidStream;
                            }
                        } else |err| return err;

                        break :blk try std.fmt.parseUnsigned(usize, max_msgs_list.items, 10);
                    };

                    return .{
                        .unsub = .{
                            .sid = sid,
                            .max_msgs = max_msgs,
                        },
                    };
                } else return error.InvalidStream;
            },
            else => |msg| std.debug.panic("Not implemented: {}\n", .{msg}),
        }
    }
};

fn readSubject(alloc: std.mem.Allocator, in: *std.Io.Reader) ![]const u8 {
    var subject_list: std.ArrayList(u8) = try .initCapacity(alloc, 1024);
    errdefer subject_list.deinit(alloc);

    // Handle the first character
    {
        const byte = try in.takeByte();
        if (std.ascii.isWhitespace(byte) or byte == '.')
            return error.InvalidStream;

        try subject_list.append(alloc, byte);
    }

    while (in.takeByte()) |byte| {
        if (std.ascii.isWhitespace(byte)) break;
        if (std.ascii.isAscii(byte)) {
            if (byte == '.') {
                const next_byte = try in.peekByte();
                if (next_byte == '.' or std.ascii.isWhitespace(next_byte))
                    return error.InvalidSubject;
            }
            try subject_list.append(alloc, byte);
        }
    } else |err| return err;
    return subject_list.toOwnedSlice(alloc);
}

inline fn expectStreamBytes(reader: *std.Io.Reader, expected: []const u8) !void {
    if (!std.mem.eql(u8, try reader.take(expected.len), expected)) {
        @branchHint(.unlikely);
        return error.InvalidStream;
    }
}

test "parsing a stream" {
    const input = "CONNECT {\"verbose\":false,\"pedantic\":false,\"tls_required\":fa" ++
        "lse,\"name\":\"NATS CLI Version v0.2.4\",\"lang\":\"go\",\"version\":\"1.43" ++
        ".0\",\"protocol\":1,\"echo\":true,\"headers\":true,\"no_responders\":true}\r" ++
        "\nPUB hi 3\r\nfoo\r\n";
    var reader: std.Io.Reader = .fixed(input);
    var arena: std.heap.ArenaAllocator = .init(std.testing.allocator);
    defer arena.deinit();
    const gpa = arena.allocator();

    {
        const msg: Message = try Message.next(gpa, &reader);
        const expected: Message = .{ .connect = .{
            .verbose = false,
            .pedantic = false,
            .tls_required = false,
            .name = "NATS CLI Version v0.2.4",
            .lang = "go",
            .version = "1.43.0",
            .protocol = 1,
            .echo = true,
            .headers = true,
            .no_responders = true,
        } };

        try std.testing.expectEqualDeep(expected, msg);
    }
    {
        const msg: Message = try Message.next(gpa, &reader);
        const expected: Message = .{ .@"pub" = .{
            .subject = "hi",
            .payload = "foo",
        } };
        try std.testing.expectEqualDeep(expected, msg);
    }
}