const std = @import("std");

pub const JSONError = error{
    IncorrectType,
    UnexpectedEOF,
    UnexpectedCharacter,
} || std.mem.Allocator.Error;

pub const JSONType = enum {
    Object,
    Array,
    String,
    Int,
    Float,
    Bool,
    Null,
};

pub const JSONObject = struct {
    allocator: std.mem.Allocator,
    children: MapType,

    const MapType = std.StringHashMap(*JSONValue);

    pub fn get(self: *JSONObject, name: []const u8) ?*JSONValue {
        return self.children.get(name);
    }

    fn init(allocator: std.mem.Allocator) JSONError!*JSONObject {
        var self = try allocator.create(JSONObject);
        self.allocator = allocator;
        self.children = MapType.init(allocator);

        return self;
    }

    pub fn deinit(self: *JSONObject) void {
        var iter = self.children.valueIterator();
        while (iter.next()) |child| {
            child.*.deinit();
        }

        var names = self.children.keyIterator();
        while (names.next()) |name| {
            self.allocator.free(name.*);
        }
        self.children.deinit();
        self.allocator.destroy(self);
    }
};

pub const JSONValue = struct {
    allocator: std.mem.Allocator,
    type: JSONType,
    vtable: VTable,

    const VTable = struct {
        getString: *const fn (self: *const JSONValue) JSONError![]const u8 = incorrectValue([]const u8),
        getArray: *const fn (self: *const JSONValue) JSONError![]*JSONValue = incorrectValue([]*JSONValue),
        getObject: *const fn (self: *const JSONValue) JSONError!*JSONObject = incorrectValue(*JSONObject),
        getInt: *const fn (self: *const JSONValue) JSONError!isize = incorrectValue(isize),
        getFloat: *const fn (self: *const JSONValue) JSONError!f64 = incorrectValue(f64),
        getBool: *const fn (self: *const JSONValue) JSONError!bool = incorrectValue(bool),

        deinit: *const fn (self: *JSONValue) void,
    };

    fn incorrectValue(comptime T: type) *const fn (*const JSONValue) JSONError!T {
        return struct {
            fn func(_: *const JSONValue) JSONError!T {
                return JSONError.IncorrectType;
            }
        }.func;
    }

    pub fn getBool(self: *const JSONValue) JSONError!bool {
        return self.vtable.getBool(self);
    }

    pub fn getInt(self: *const JSONValue) JSONError!isize {
        return self.vtable.getInt(self);
    }

    pub fn getFloat(self: *const JSONValue) JSONError!f64 {
        return self.vtable.getFloat(self);
    }

    pub fn getString(self: *const JSONValue) JSONError![]const u8 {
        return self.vtable.getString(self);
    }

    pub fn getArray(self: *const JSONValue) JSONError![]*JSONValue {
        return self.vtable.getArray(self);
    }

    pub fn getObject(self: *const JSONValue) JSONError!*JSONObject {
        return self.vtable.getObject(self);
    }

    pub fn deinit(self: *JSONValue) void {
        self.vtable.deinit(self);
    }

    fn createNull(allocator: std.mem.Allocator) JSONError!*JSONValue {
        var self = try allocator.create(JSONValue);
        self.type = .Null;
        self.vtable = .{ .deinit = basicDeinit };
        self.allocator = allocator;

        return self;
    }

    fn basicDeinit(self: *JSONValue) void {
        self.allocator.destroy(self);
    }
};

const JSONArrayValue = struct {
    value: JSONValue,
    array: []*JSONValue,

    fn init(allocator: std.mem.Allocator, array: []*JSONValue) !*JSONValue {
        var self = try allocator.create(JSONArrayValue);
        self.array = array;
        self.value = .{
            .type = .Array,
            .vtable = .{ .getArray = getValue, .deinit = deinit },
            .allocator = allocator,
        };

        return &self.value;
    }

    fn getValue(value: *const JSONValue) JSONError![]*JSONValue {
        const self: *const JSONArrayValue = @fieldParentPtr("value", value);
        return self.array;
    }

    fn deinit(value: *JSONValue) void {
        var self: *JSONArrayValue = @fieldParentPtr("value", value);
        for (self.array) |val| {
            val.deinit();
        }
        self.value.allocator.free(self.array);
        self.value.allocator.destroy(self);
    }
};

const JSONObjectValue = struct {
    value: JSONValue,
    object: *JSONObject,

    fn init(allocator: std.mem.Allocator, object: *JSONObject) !*JSONValue {
        var self = try allocator.create(JSONObjectValue);
        self.object = object;
        self.value = .{
            .type = .Object,
            .vtable = .{ .getObject = getValue, .deinit = deinit },
            .allocator = allocator,
        };

        return &self.value;
    }

    fn getValue(value: *const JSONValue) JSONError!*JSONObject {
        const self: *const JSONObjectValue = @fieldParentPtr("value", value);
        return self.object;
    }

    fn deinit(value: *JSONValue) void {
        var self: *const JSONObjectValue = @fieldParentPtr("value", value);
        self.object.deinit();
        self.value.allocator.destroy(self);
    }
};

const JSONIntValue = struct {
    value: JSONValue,
    int: isize,

    fn init(allocator: std.mem.Allocator, value: isize) !*JSONValue {
        var self = try allocator.create(JSONIntValue);
        self.int = value;
        self.value = .{
            .type = .Int,
            .vtable = .{ .getInt = getValue, .deinit = deinit },
            .allocator = allocator,
        };

        return &self.value;
    }

    fn getValue(value: *const JSONValue) JSONError!isize {
        const self: *const JSONIntValue = @fieldParentPtr("value", value);
        return self.int;
    }

    fn deinit(value: *JSONValue) void {
        var self: *JSONIntValue = @fieldParentPtr("value", value);
        self.value.allocator.destroy(self);
    }
};

const JSONFloatValue = struct {
    value: JSONValue,
    float: f64,

    fn init(allocator: std.mem.Allocator, value: f64) !*JSONValue {
        var self = try allocator.create(JSONFloatValue);
        self.float = value;
        self.value = .{
            .type = .Float,
            .vtable = .{ .getFloat = getValue, .deinit = deinit },
            .allocator = allocator,
        };

        return &self.value;
    }

    fn getValue(value: *const JSONValue) JSONError!f64 {
        const self: *const JSONFloatValue = @fieldParentPtr("value", value);
        return self.float;
    }

    fn deinit(value: *JSONValue) void {
        var self: *JSONIntValue = @fieldParentPtr("value", value);
        self.value.allocator.destroy(self);
    }
};

const JSONBoolValue = struct {
    value: JSONValue,
    boolean: bool,

    fn init(allocator: std.mem.Allocator, value: bool) !*JSONValue {
        var self = try allocator.create(JSONBoolValue);
        self.boolean = value;
        self.value = .{
            .type = .Bool,
            .vtable = .{ .getBool = getValue, .deinit = deinit },
            .allocator = allocator,
        };

        return &self.value;
    }

    fn getValue(value: *const JSONValue) JSONError!bool {
        const self: *const JSONBoolValue = @fieldParentPtr("value", value);
        return self.boolean;
    }

    fn deinit(value: *JSONValue) void {
        var self: *JSONBoolValue = @fieldParentPtr("value", value);
        self.value.allocator.destroy(self);
    }
};

pub const JSONStringValue = struct {
    value: JSONValue,
    string: []const u8,

    pub fn init(allocator: std.mem.Allocator, string: []const u8) !*JSONValue {
        var self = try allocator.create(JSONStringValue);
        self.string = string;
        self.value = .{
            .type = .String,
            .vtable = .{ .getString = getValue, .deinit = deinit },
            .allocator = allocator,
        };

        return &self.value;
    }

    fn getValue(value: *const JSONValue) JSONError![]const u8 {
        const self: *const JSONStringValue = @fieldParentPtr("value", value);
        return self.string;
    }

    fn deinit(value: *JSONValue) void {
        var self: *JSONStringValue = @fieldParentPtr("value", value);
        self.value.allocator.free(self.string);
        self.value.allocator.destroy(self);
    }
};

pub fn parseFile(allocator: std.mem.Allocator, fileName: []const u8) !*JSONObject {
    const fileContents = blk: {
        var file = try std.fs.cwd().openFile(fileName, .{ .mode = .read_only });
        defer file.close();
        break :blk try file.readToEndAlloc(allocator, 64 * 1024 * 1024);
    };
    defer allocator.free(fileContents);

    var parser = try JSONParser.init(allocator, fileContents);
    return parser.parse();
}

pub fn parseString(allocator: std.mem.Allocator, string: []const u8) !*JSONObject {
    var parser = try JSONParser.init(allocator, string);
    return parser.parse();
}

const JSONParser = struct {
    allocator: std.mem.Allocator,
    iter: std.unicode.Utf8Iterator,
    current: u21 = 0,
    end: bool = false,

    fn init(allocator: std.mem.Allocator, input: []const u8) !JSONParser {
        return JSONParser{
            .allocator = allocator,
            .iter = (try std.unicode.Utf8View.init(input)).iterator(),
        };
    }

    fn next(self: *JSONParser) void {
        if (self.iter.nextCodepoint()) |code_point| {
            self.current = code_point;
        } else {
            self.end = true;
        }
    }

    fn consume(self: *JSONParser) u21 {
        const code_point = self.current;
        self.next();
        return code_point;
    }

    fn skipWhitespace(self: *JSONParser) JSONError!void {
        while (!self.end) {
            switch (self.current) {
                ' ', '\t', '\n', '\r' => {},
                else => return,
            }
            self.next();
        }

        return JSONError.UnexpectedEOF;
    }

    fn expectNext(self: *JSONParser, code_point: u21) JSONError!void {
        //std.debug.print("Expecting {u} ({X}) to be {u} ({X})\n", .{ self.current, self.current, code_point, code_point });
        if (self.current != code_point or self.end) {
            return JSONError.UnexpectedCharacter;
        }
        self.next();
    }

    fn parse(self: *JSONParser) JSONError!*JSONObject {
        self.next();
        try self.skipWhitespace();
        return self.parseObject();
    }

    fn parseObject(self: *JSONParser) JSONError!*JSONObject {
        const obj = try JSONObject.init(self.allocator);
        errdefer obj.deinit();

        try self.expectNext('{');

        try self.skipWhitespace();

        if (self.current != '}') {
            while (true) {
                try self.skipWhitespace();
                const valueName = try self.parseString();
                errdefer self.allocator.free(valueName);
                try self.skipWhitespace();

                //std.debug.print("Value name: {s}\n", .{valueName});

                try self.expectNext(':');

                const value = try self.parseValue();
                errdefer value.deinit();

                try obj.children.put(valueName, value);

                // Do-while would have been better here
                if (self.current != ',') {
                    break;
                }
                self.next();
            }
        }

        try self.expectNext('}');

        return obj;
    }

    fn parseArray(self: *JSONParser) JSONError![]*JSONValue {
        var array = std.ArrayList(*JSONValue).init(self.allocator);
        errdefer {
            for (array.items) |value| {
                value.deinit();
            }
            array.deinit();
        }

        try self.expectNext('[');

        try self.skipWhitespace();

        if (self.current != ']') {
            while (true) {
                try array.append(try self.parseValue());
                if (self.current != ',') {
                    break;
                }
                self.next();
            }
        }

        try self.expectNext(']');

        return array.toOwnedSlice();
    }

    fn parseValue(self: *JSONParser) JSONError!*JSONValue {
        try self.skipWhitespace();

        var value: *JSONValue = undefined;

        switch (self.current) {
            '"' => {
                //std.debug.print("Parsing string\n", .{});
                const str = try self.parseString();
                value = try JSONStringValue.init(self.allocator, str);
            },
            '-', '0'...'9' => {
                //std.debug.print("Parsing number\n", .{});
                value = try self.parseNumberValue();
            },
            '{' => {
                //std.debug.print("Parsing object\n", .{});
                const obj = try self.parseObject();
                value = try JSONObjectValue.init(self.allocator, obj);
            },
            '[' => {
                //std.debug.print("Parsing array\n", .{});
                const arr = try self.parseArray();
                value = try JSONArrayValue.init(self.allocator, arr);
            },
            't', 'f' => {
                //std.debug.print("Parsing bool\n", .{});
                const b = try self.parseBool();
                value = try JSONBoolValue.init(self.allocator, b);
            },
            'n' => {
                //std.debug.print("Parsing null\n", .{});
                self.next();
                try self.expectNext('u');
                try self.expectNext('l');
                try self.expectNext('l');

                value = try JSONValue.createNull(self.allocator);
            },
            else => {
                std.debug.print("Unexpected character: {u}\n", .{self.current});
                return JSONError.UnexpectedCharacter;
            },
        }

        self.skipWhitespace() catch |err| {
            value.deinit();
            return err;
        };

        return value;
    }

    fn parseBool(self: *JSONParser) JSONError!bool {
        var b = false;

        switch (self.current) {
            't' => {
                self.next();
                try self.expectNext('r');
                try self.expectNext('u');
                try self.expectNext('e');

                b = true;
            },
            'f' => {
                self.next();
                try self.expectNext('a');
                try self.expectNext('l');
                try self.expectNext('s');
                try self.expectNext('e');

                b = false;
            },
            else => return JSONError.UnexpectedCharacter,
        }

        return b;
    }

    fn parseString(self: *JSONParser) JSONError![]const u8 {
        var string = std.ArrayList(u8).init(self.allocator);
        defer string.deinit();

        try self.expectNext('"');

        while (self.current != '"' and !self.end) {
            switch (self.consume()) {
                '\\' => {
                    if (self.end) return JSONError.UnexpectedEOF;
                    const code_point = try self.escapedCharacter();
                    var utf8: [4]u8 = undefined;
                    const len = std.unicode.utf8Encode(code_point, &utf8) catch return JSONError.UnexpectedCharacter;
                    try string.appendSlice(utf8[0..len]);
                },
                0...std.ascii.control_code.us, std.ascii.control_code.del => return JSONError.UnexpectedCharacter,
                else => |c| {
                    var buf: [4]u8 = undefined;
                    const len = std.unicode.utf8Encode(c, &buf) catch return JSONError.UnexpectedCharacter;
                    try string.appendSlice(buf[0..len]);
                },
            }
        }

        try self.expectNext('"');

        return string.toOwnedSlice();
    }

    fn escapedCharacter(self: *JSONParser) JSONError!u21 {
        return switch (self.consume()) {
            '"' => '"',
            '\\' => '\\',
            '/' => '/',
            'b' => '\x08',
            'f' => '\x0C',
            'n' => '\n',
            'r' => '\r',
            't' => '\t',
            'u' => {
                // I hate this
                if (self.end) return JSONError.UnexpectedEOF;
                const c1 = self.consume();
                var b1: [1]u8 = undefined;
                _ = std.unicode.utf8Encode(c1, &b1) catch return JSONError.UnexpectedCharacter;
                if (self.end) return JSONError.UnexpectedEOF;
                const c2 = self.consume();
                var b2: [1]u8 = undefined;
                _ = std.unicode.utf8Encode(c2, &b2) catch return JSONError.UnexpectedCharacter;
                if (self.end) return JSONError.UnexpectedEOF;
                const c3 = self.consume();
                var b3: [1]u8 = undefined;
                _ = std.unicode.utf8Encode(c3, &b3) catch return JSONError.UnexpectedCharacter;
                if (self.end) return JSONError.UnexpectedEOF;
                const c4 = self.consume();
                var b4: [1]u8 = undefined;
                _ = std.unicode.utf8Encode(c4, &b4) catch return JSONError.UnexpectedCharacter;

                return std.fmt.parseInt(u21, &[4]u8{ b1[0], b2[0], b3[0], b4[0] }, 16) catch JSONError.UnexpectedCharacter;
            },
            else => JSONError.UnexpectedCharacter,
        };
    }

    fn parseNumberValue(self: *JSONParser) JSONError!*JSONValue {
        var buf = std.ArrayList(u8).init(self.allocator);
        defer buf.deinit();
        var float = false;
        var exp: isize = 0;

        if (self.current == '-') {
            self.next();
            try buf.append('-');
        }

        if (self.end) {
            return JSONError.UnexpectedEOF;
        }

        if (self.current != '0') {
            while (!self.end) {
                switch (self.current) {
                    '0'...'9' => |c| {
                        try buf.append(@intCast(c));
                    },
                    else => break,
                }
                self.next();
            }

            if (self.current == '.') {
                try buf.append('.');
                float = true;
                self.next();
            }
        } else {
            try self.expectNext('.');
            try buf.append('.');
            float = true;
        }

        if (float) {
            while (!self.end) {
                switch (self.current) {
                    '0'...'9' => |c| {
                        try buf.append(@intCast(c));
                    },
                    else => break,
                }
                self.next();
            }
        }

        if (self.current == 'e' or self.current == 'E') {
            self.next();

            var expBuf = std.ArrayList(u8).init(self.allocator);
            defer expBuf.deinit();

            while (!self.end) {
                switch (self.current) {
                    '-' => try expBuf.append('-'),
                    '0'...'9' => |c| try expBuf.append(@intCast(c)),
                    else => break,
                }
                self.next();
            }

            if (expBuf.items.len == 0) {
                return JSONError.UnexpectedCharacter;
            }

            exp = std.fmt.parseInt(isize, expBuf.items, 10) catch return JSONError.UnexpectedCharacter;
        }

        if (float or exp != 0) {
            var value: f64 = std.fmt.parseFloat(f64, buf.items) catch unreachable;

            if (exp != 0) {
                value *= std.math.pow(f64, 10.0, @floatFromInt(exp));
            }

            return try JSONFloatValue.init(self.allocator, value);
        } else {
            const value: isize = std.fmt.parseInt(isize, buf.items, 10) catch unreachable;
            return try JSONIntValue.init(self.allocator, value);
        }
    }
};

test "types" {
    var root = try parseString(std.testing.allocator, "{\"string\": \"abc\", \"int\": 12, \"float\": 3.4, \"array\": [], \"object\": {}, \"null\": null, \"bool\": true}");
    defer root.deinit();

    if (root.get("string")) |value| {
        try std.testing.expect(value.type == .String);
    } else {
        return error.FailedToParse;
    }

    if (root.get("int")) |value| {
        try std.testing.expect(value.type == .Int);
    } else {
        return error.FailedToParse;
    }

    if (root.get("float")) |value| {
        try std.testing.expect(value.type == .Float);
    } else {
        return error.FailedToParse;
    }

    if (root.get("array")) |value| {
        try std.testing.expect(value.type == .Array);
    } else {
        return error.FailedToParse;
    }

    if (root.get("object")) |value| {
        try std.testing.expect(value.type == .Object);
    } else {
        return error.FailedToParse;
    }

    if (root.get("null")) |value| {
        try std.testing.expect(value.type == .Null);
    } else {
        return error.FailedToParse;
    }

    if (root.get("bool")) |value| {
        try std.testing.expect(value.type == .Bool);
    } else {
        return error.FailedToParse;
    }
}

test "ints" {
    const root = try parseString(std.testing.allocator, "{\"one\": 1, \"two\": 2, \"three\": 3, \"fifteen\": 15, \"negative\": -15}");
    defer root.deinit();

    if (root.get("one")) |value| {
        try std.testing.expect(value.type == .Int);
        try std.testing.expect(try value.getInt() == 1);
    } else {
        return error.FailedToParse;
    }

    if (root.get("two")) |value| {
        try std.testing.expect(value.type == .Int);
        try std.testing.expect(try value.getInt() == 2);
    } else {
        return error.FailedToParse;
    }

    if (root.get("three")) |value| {
        try std.testing.expect(value.type == .Int);
        try std.testing.expect(try value.getInt() == 3);
    } else {
        return error.FailedToParse;
    }

    if (root.get("fifteen")) |value| {
        try std.testing.expect(value.type == .Int);
        try std.testing.expect(try value.getInt() == 15);
    } else {
        return error.FailedToParse;
    }

    if (root.get("negative")) |value| {
        try std.testing.expect(value.type == .Int);
        try std.testing.expect(try value.getInt() == -15);
    } else {
        return error.FailedToParse;
    }
}

test "floats" {
    const root = try parseString(std.testing.allocator, "{\"one\": 1.2, \"ten\": 1e1, \"three\": -3.4, \"four\": 3.4e12, \"five\": 1e-2}");
    defer root.deinit();

    if (root.get("one")) |value| {
        try std.testing.expect(value.type == .Float);
        try std.testing.expect(try value.getFloat() - 1.2 < 0.1);
    } else {
        return error.FailedToParse;
    }

    if (root.get("ten")) |value| {
        try std.testing.expect(value.type == .Float);
        try std.testing.expect(try value.getFloat() - 10 < 0.1);
    } else {
        return error.FailedToParse;
    }

    if (root.get("three")) |value| {
        try std.testing.expect(value.type == .Float);
        try std.testing.expect(try value.getFloat() - -3.4 < 0.1);
    } else {
        return error.FailedToParse;
    }

    if (root.get("four")) |value| {
        try std.testing.expect(value.type == .Float);
        try std.testing.expect(try value.getFloat() - 3400000000000 < 0.1);
    } else {
        return error.FailedToParse;
    }

    if (root.get("five")) |value| {
        try std.testing.expect(value.type == .Float);
        try std.testing.expect(try value.getFloat() - 0.01 < 0.001);
    } else {
        return error.FailedToParse;
    }
}

test "bool" {
    const root = try parseString(std.testing.allocator, "{\"true\": true, \"false\": false}");
    defer root.deinit();

    if (root.get("true")) |value| {
        try std.testing.expect(value.type == .Bool);
        try std.testing.expect(try value.getBool());
    } else {
        return error.FailedToParse;
    }

    if (root.get("false")) |value| {
        try std.testing.expect(value.type == .Bool);
        try std.testing.expect(!try value.getBool());
    } else {
        return error.FailedToParse;
    }
}

test "string" {
    const root = try parseString(std.testing.allocator, "{\"one\": \"abc 123\\r\\t\\n\\u2611 \\\" \\\\ \\/\"}");
    defer root.deinit();

    if (root.get("one")) |value| {
        try std.testing.expect(value.type == .String);
        try std.testing.expectEqualStrings("abc 123\r\t\n☑ \" \\ /", try value.getString());
    } else {
        return error.FailedToParse;
    }
}

test "null" {
    const root = try parseString(std.testing.allocator, "{\"null\": null}");
    defer root.deinit();

    if (root.get("null")) |value| {
        try std.testing.expect(value.type == .Null);
    } else {
        return error.FailedToParse;
    }
}

test "array" {
    const root = try parseString(std.testing.allocator, "{\"empty\": [], \"two\": [1, 2], \"three\": [ 3.1, 3.2, 3.3 ]}");
    defer root.deinit();

    if (root.get("empty")) |value| {
        try std.testing.expect(value.type == .Array);
        try std.testing.expect((try value.getArray()).len == 0);
    } else {
        return error.FailedToParse;
    }

    if (root.get("two")) |value| {
        try std.testing.expect(value.type == .Array);
        const arr = try value.getArray();
        try std.testing.expect(arr.len == 2);

        try std.testing.expect(arr[0].type == .Int);
        try std.testing.expect(try arr[0].getInt() == 1);

        try std.testing.expect(arr[1].type == .Int);
        try std.testing.expect(try arr[1].getInt() == 2);
    } else {
        return error.FailedToParse;
    }

    if (root.get("three")) |value| {
        try std.testing.expect(value.type == .Array);
        const arr = try value.getArray();
        try std.testing.expect(arr.len == 3);

        try std.testing.expect(arr[0].type == .Float);
        try std.testing.expect(try arr[0].getFloat() - 3.1 < 0.1);

        try std.testing.expect(arr[1].type == .Float);
        try std.testing.expect(try arr[1].getFloat() - 3.2 < 0.1);

        try std.testing.expect(arr[2].type == .Float);
        try std.testing.expect(try arr[2].getFloat() - 3.3 < 0.1);
    } else {
        return error.FailedToParse;
    }
}

test "objects" {
    const root = try parseString(std.testing.allocator, "{\"empty\": {}, \"two\": {\"abc\": 1}}");
    defer root.deinit();

    if (root.get("empty")) |value| {
        try std.testing.expect(value.type == .Object);
        const obj = try value.getObject();
        try std.testing.expect(obj.children.count() == 0);
    } else {
        return error.FailedToParse;
    }

    if (root.get("two")) |value| {
        try std.testing.expect(value.type == .Object);
        const obj = try value.getObject();
        try std.testing.expect(obj.children.count() == 1);

        if (obj.get("abc")) |value2| {
            try std.testing.expect(value2.type == .Int);
            try std.testing.expect(try value2.getInt() == 1);
        } else {
            return error.FailedToParse;
        }
    } else {
        return error.FailedToParse;
    }
}

test "file parsing" {
    const root = try parseFile(std.testing.allocator, "test/test.json");
    root.deinit();
}