const std = @import("std");
const Allocator = std.mem.Allocator;

pub fn main(init: std.process.Init.Minimal) !void {
    var debug_allocator: std.heap.DebugAllocator(.{ .stack_trace_frames = 0 }) = .{};
    defer _ = debug_allocator.deinit();
    const gpa = debug_allocator.allocator();

    var arena_allocator = std.heap.ArenaAllocator.init(gpa);
    defer arena_allocator.deinit();
    const arena = arena_allocator.allocator();

    var threaded_io = std.Io.Threaded.init(gpa, .{
        .environ = init.environ,
    });
    defer threaded_io.deinit();
    const io = threaded_io.io();
    const cwd = std.Io.Dir.cwd();

    const args = try init.args.toSlice(arena);
    if (args.len != 3) {
        const stderr = std.debug.lockStderr(&.{});
        stderr.file_writer.interface.print("Usage: {s} <input-file> <output-path>", .{args[0]}) catch {};
        std.process.exit(1);
    }

    const input_path = args[1];
    const output_path = args[2];

    const input = try cwd.readFileAlloc(io, input_path, gpa, .unlimited);
    defer gpa.free(input);

    var values: std.StringArrayHashMapUnmanaged([]const []const u8) = .empty;
    var properties: std.ArrayList([]const u8) = .empty;
    var headers: std.ArrayList([]const u8) = .empty;

    var value_name: ?[]const u8 = null;
    var it = std.mem.tokenizeAny(u8, input, "\r\n");
    while (it.next()) |line_untrimmed| {
        const line = std.mem.trim(u8, line_untrimmed, " \t");
        if (line.len == 0 or line[0] == '#') continue;
        if (std.mem.startsWith(u8, line, "const ") or std.mem.startsWith(u8, line, "pub const ")) {
            try headers.append(arena, line);
            continue;
        }
        if (line[0] == '.') {
            if (value_name == null) {
                std.debug.panic("property not attached to a value:\n\"{s}\"", .{line});
            }
            try properties.append(arena, line);
            continue;
        }

        if (value_name) |name| {
            const old = try values.fetchPut(arena, name, try properties.toOwnedSlice(arena));
            if (old != null) std.debug.panic("duplicate value \"{s}\"", .{name});
        }
        value_name = line;
    }

    if (value_name) |name| {
        const old = try values.fetchPut(arena, name, try properties.toOwnedSlice(arena));
        if (old != null) std.debug.panic("duplicate value \"{s}\"", .{name});
    }

    const sorted_list = try arena.dupe([]const u8, values.keys());
    defer arena.free(sorted_list);
    std.mem.sort([]const u8, sorted_list, {}, struct {
        pub fn lessThan(_: void, a_str: []const u8, b_str: []const u8) bool {
            return std.mem.lessThan(u8, a_str, b_str);
        }
    }.lessThan);

    var longest_name: usize = 0;
    var shortest_name: usize = std.math.maxInt(usize);

    var builder = try DafsaBuilder.init(arena);
    defer builder.deinit();
    for (sorted_list) |name| {
        try builder.insert(name);
        longest_name = @max(name.len, longest_name);
        shortest_name = @min(name.len, shortest_name);
    }
    try builder.finish();
    builder.calcNumbers();

    // As a sanity check, confirm that the minimal perfect hashing doesn't
    // have any collisions
    {
        var index_set = std.AutoHashMap(usize, void).init(arena);
        defer index_set.deinit();

        for (values.keys()) |name| {
            const index = builder.getUniqueIndex(name).?;
            const result = try index_set.getOrPut(index);
            if (result.found_existing) {
                std.debug.panic("clobbered {}, name={s}\n", .{ index, name });
            }
        }
    }

    const output_file = try cwd.createFile(io, output_path, .{});
    defer output_file.close(io);

    var output_buf: [4196]u8 = undefined;
    var output_writer = output_file.writer(io, &output_buf);
    const writer = &output_writer.interface;

    try writer.print(
        \\//! Autogenerated by GenerateDef from {s}, do not edit
        \\
        \\const std = @import("std");
        \\
        \\pub fn with(comptime Properties: type) type {{
        \\return struct {{
        \\
    , .{input_path});
    for (headers.items) |line| {
        try writer.print("{s}\n", .{line});
    }

    var values_array = try arena.alloc(Value, values.count());
    defer arena.free(values_array);

    for (values.keys(), values.values()) |name, props| {
        const unique_index = builder.getUniqueIndex(name).?;
        const data_index = unique_index - 1;
        values_array[data_index] = .{ .name = name, .properties = props };
    }

    try writer.writeAll(
        \\
        \\/// Integer starting at 0 derived from the unique index,
        \\/// corresponds with the data array index.
        \\pub const Tag = enum(u16) {
    );
    for (values_array) |value| {
        try writer.print("    {f},\n", .{std.zig.fmtId(value.name)});
    }
    try writer.writeAll(
        \\};
        \\
        \\pub fn fromName(name: []const u8) ?Properties {
        \\    const data_index = tagFromName(name) orelse return null;
        \\    return data[@intFromEnum(data_index)];
        \\}
        \\
        \\pub fn tagFromName(name: []const u8) ?Tag {
        \\    const unique_index = uniqueIndex(name) orelse return null;
        \\    return @enumFromInt(unique_index - 1);
        \\}
        \\
        \\pub fn fromTag(tag: Tag) Properties {
        \\    return data[@intFromEnum(tag)];
        \\}
        \\
        \\pub fn nameFromTagIntoBuf(tag: Tag, name_buf: []u8) []u8 {
        \\    std.debug.assert(name_buf.len >= longest_name);
        \\    const unique_index = @intFromEnum(tag) + 1;
        \\    return nameFromUniqueIndex(unique_index, name_buf);
        \\}
        \\
        \\pub fn nameFromTag(tag: Tag) NameBuf {
        \\    var name_buf: NameBuf = undefined;
        \\    const unique_index = @intFromEnum(tag) + 1;
        \\    const name = nameFromUniqueIndex(unique_index, &name_buf.buf);
        \\    name_buf.len = @intCast(name.len);
        \\    return name_buf;
        \\}
        \\
        \\pub const NameBuf = struct {
        \\    buf: [longest_name]u8 = undefined,
        \\    len: std.math.IntFittingRange(0, longest_name),
        \\
        \\    pub fn span(self: *const NameBuf) []const u8 {
        \\        return self.buf[0..self.len];
        \\    }
        \\};
        \\
        \\pub fn exists(name: []const u8) bool {
        \\    if (name.len < shortest_name or name.len > longest_name) return false;
        \\
        \\    var index: u16 = 0;
        \\    for (name) |c| {
        \\        index = findInList(dafsa[index].child_index, c) orelse return false;
        \\    }
        \\    return dafsa[index].end_of_word;
        \\}
        \\
        \\
    );
    try writer.print("pub const shortest_name = {};\n", .{shortest_name});
    try writer.print("pub const longest_name = {};\n\n", .{longest_name});
    try writer.writeAll(
        \\/// Search siblings of `first_child_index` for the `char`
        \\/// If found, returns the index of the node within the `dafsa` array.
        \\/// Otherwise, returns `null`.
        \\pub fn findInList(first_child_index: u16, char: u8) ?u16 {
        \\
    );
    try writer.print("    @setEvalBranchQuota({d});\n", .{values.count() * 2});
    try writer.writeAll(
        \\    var index = first_child_index;
        \\    while (true) {
        \\        if (dafsa[index].char == char) return index;
        \\        if (dafsa[index].end_of_list) return null;
        \\        index += 1;
        \\    }
        \\    unreachable;
        \\}
        \\
        \\/// Returns a unique (minimal perfect hash) index (starting at 1) for the `name`,
        \\/// or null if the name was not found.
        \\pub fn uniqueIndex(name: []const u8) ?u16 {
        \\    if (name.len < shortest_name or name.len > longest_name) return null;
        \\
        \\    var index: u16 = 0;
        \\    var node_index: u16 = 0;
        \\
        \\    for (name) |c| {
        \\        const child_index = findInList(dafsa[node_index].child_index, c) orelse return null;
        \\        var sibling_index = dafsa[node_index].child_index;
        \\        while (true) {
        \\            const sibling_c = dafsa[sibling_index].char;
        \\            std.debug.assert(sibling_c != 0);
        \\            if (sibling_c < c) {
        \\                index += dafsa[sibling_index].number;
        \\            }
        \\            if (dafsa[sibling_index].end_of_list) break;
        \\            sibling_index += 1;
        \\        }
        \\        node_index = child_index;
        \\        if (dafsa[node_index].end_of_word) index += 1;
        \\    }
        \\
        \\    if (!dafsa[node_index].end_of_word) return null;
        \\
        \\    return index;
        \\}
        \\
        \\/// Returns a slice of `buf` with the name associated with the given `index`.
        \\/// This function should only be called with an `index` that
        \\/// is already known to exist within the `dafsa`, e.g. an index
        \\/// returned from `uniqueIndex`.
        \\pub fn nameFromUniqueIndex(index: u16, buf: []u8) []u8 {
        \\    std.debug.assert(index >= 1 and index <= data.len);
        \\
        \\    var node_index: u16 = 0;
        \\    var count: u16 = index;
        \\    var w = std.Io.Writer.fixed(buf);
        \\
        \\    while (true) {
        \\        var sibling_index = dafsa[node_index].child_index;
        \\        while (true) {
        \\            if (dafsa[sibling_index].number > 0 and dafsa[sibling_index].number < count) {
        \\                count -= dafsa[sibling_index].number;
        \\            } else {
        \\                w.writeByte(dafsa[sibling_index].char) catch unreachable;
        \\                node_index = sibling_index;
        \\                if (dafsa[node_index].end_of_word) {
        \\                    count -= 1;
        \\                }
        \\                break;
        \\            }
        \\
        \\            if (dafsa[sibling_index].end_of_list) break;
        \\            sibling_index += 1;
        \\        }
        \\        if (count == 0) break;
        \\    }
        \\
        \\    return w.buffered();
        \\}
        \\
        \\
    );

    try writer.writeAll(
        \\const Node = packed struct {
        \\    char: u8,
        \\    /// Nodes are numbered with "an integer which gives the number of words that
        \\    /// would be accepted by the automaton starting from that state." This numbering
        \\    /// allows calculating "a one-to-one correspondence between the integers 1 to L
        \\    /// (L is the number of words accepted by the automaton) and the words themselves."
        \\    ///
        \\    /// Essentially, this allows us to have a minimal perfect hashing scheme such that
        \\    /// it's possible to store & lookup the properties of each builtin using a separate array.
        \\    number: std.math.IntFittingRange(0, data.len),
        \\    /// If true, this node is the end of a valid builtin.
        \\    /// Note: This does not necessarily mean that this node does not have child nodes.
        \\    end_of_word: bool,
        \\    /// If true, this node is the end of a sibling list.
        \\    /// If false, then (index + 1) will contain the next sibling.
        \\    end_of_list: bool,
        \\    /// Index of the first child of this node.
        \\    child_index: u16,
        \\};
        \\
        \\
    );
    try builder.writeDafsa(writer);
    try writeData(writer, values_array);
    try writer.writeAll(
        \\};
        \\}
        \\
    );
    try writer.flush();
}

const Value = struct {
    name: []const u8,
    properties: []const []const u8,
};

fn writeData(writer: *std.Io.Writer, values: []const Value) !void {
    try writer.writeAll("pub const data = blk: {\n");
    try writer.print("    @setEvalBranchQuota({d});\n", .{values.len * 9});
    try writer.writeAll("    break :blk [_]Properties{\n");
    for (values) |value| {
        try writer.writeAll("        .{");
        for (value.properties, 0..) |property, j| {
            if (j != 0) try writer.writeByte(',');
            try writer.writeByte(' ');
            try writer.writeAll(property);
        }
        if (value.properties.len != 0) try writer.writeByte(' ');
        try writer.writeAll("},\n");
    }
    try writer.writeAll("    };\n");
    try writer.writeAll("};\n");
}

const DafsaBuilder = struct {
    root: *Node,
    arena: std.heap.ArenaAllocator.State,
    allocator: Allocator,
    unchecked_nodes: std.ArrayListUnmanaged(UncheckedNode),
    minimized_nodes: std.HashMapUnmanaged(*Node, *Node, Node.DuplicateContext, std.hash_map.default_max_load_percentage),
    previous_word_buf: [128]u8 = undefined,
    previous_word: []u8 = &[_]u8{},

    const UncheckedNode = struct {
        parent: *Node,
        char: u8,
        child: *Node,
    };

    pub fn init(allocator: Allocator) !DafsaBuilder {
        var arena = std.heap.ArenaAllocator.init(allocator);
        errdefer arena.deinit();

        const root = try arena.allocator().create(Node);
        root.* = .{};
        return .{
            .root = root,
            .allocator = allocator,
            .arena = arena.state,
            .unchecked_nodes = .empty,
            .minimized_nodes = .{},
        };
    }

    pub fn deinit(db: *DafsaBuilder) void {
        db.arena.promote(db.allocator).deinit();
        db.unchecked_nodes.deinit(db.allocator);
        db.minimized_nodes.deinit(db.allocator);
        db.* = undefined;
    }

    const Node = struct {
        children: [256]?*Node = @splat(null),
        is_terminal: bool = false,
        number: usize = 0,

        const DuplicateContext = struct {
            pub fn hash(ctx: @This(), key: *Node) u64 {
                _ = ctx;
                var hasher = std.hash.Wyhash.init(0);
                std.hash.autoHash(&hasher, key.children);
                std.hash.autoHash(&hasher, key.is_terminal);
                return hasher.final();
            }

            pub fn eql(ctx: @This(), a: *Node, b: *Node) bool {
                _ = ctx;
                return a.is_terminal == b.is_terminal and std.mem.eql(?*Node, &a.children, &b.children);
            }
        };

        pub fn calcNumbers(db: *Node) void {
            db.number = @intFromBool(db.is_terminal);
            for (db.children) |maybe_child| {
                const child = maybe_child orelse continue;
                // A node's number is the sum of the
                // numbers of its immediate child nodes.
                child.calcNumbers();
                db.number += child.number;
            }
        }

        pub fn numDirectChildren(db: *const Node) u8 {
            var num: u8 = 0;
            for (db.children) |child| {
                if (child != null) num += 1;
            }
            return num;
        }
    };

    pub fn insert(db: *DafsaBuilder, str: []const u8) !void {
        if (std.mem.order(u8, str, db.previous_word) == .lt) {
            @panic("insertion order must be sorted");
        }

        var common_prefix_len: usize = 0;
        for (0..@min(str.len, db.previous_word.len)) |i| {
            if (str[i] != db.previous_word[i]) break;
            common_prefix_len += 1;
        }

        try db.minimize(common_prefix_len);

        var node = if (db.unchecked_nodes.items.len == 0)
            db.root
        else
            db.unchecked_nodes.getLast().child;

        for (str[common_prefix_len..]) |c| {
            std.debug.assert(node.children[c] == null);

            var arena = db.arena.promote(db.allocator);
            const child = try arena.allocator().create(Node);
            db.arena = arena.state;

            child.* = .{};
            node.children[c] = child;
            try db.unchecked_nodes.append(db.allocator, .{
                .parent = node,
                .char = c,
                .child = child,
            });
            node = node.children[c].?;
        }
        node.is_terminal = true;

        db.previous_word = db.previous_word_buf[0..str.len];
        @memcpy(db.previous_word, str);
    }

    pub fn minimize(db: *DafsaBuilder, down_to: usize) !void {
        if (db.unchecked_nodes.items.len == 0) return;
        while (db.unchecked_nodes.items.len > down_to) {
            const unchecked_node = db.unchecked_nodes.pop().?;
            if (db.minimized_nodes.getPtr(unchecked_node.child)) |child| {
                unchecked_node.parent.children[unchecked_node.char] = child.*;
            } else {
                try db.minimized_nodes.put(db.allocator, unchecked_node.child, unchecked_node.child);
            }
        }
    }

    pub fn finish(db: *DafsaBuilder) !void {
        try db.minimize(0);
    }

    fn nodeCount(db: *const DafsaBuilder) usize {
        return db.minimized_nodes.count();
    }

    fn edgeCount(db: *const DafsaBuilder) usize {
        var count: usize = 0;
        var it = db.minimized_nodes.iterator();
        while (it.next()) |entry| {
            for (entry.key_ptr.*.children) |child| {
                if (child != null) count += 1;
            }
        }
        return count;
    }

    fn contains(db: *const DafsaBuilder, str: []const u8) bool {
        var node = db.root;
        for (str) |c| {
            node = node.children[c] orelse return false;
        }
        return node.is_terminal;
    }

    fn calcNumbers(db: *const DafsaBuilder) void {
        db.root.calcNumbers();
    }

    fn getUniqueIndex(db: *const DafsaBuilder, str: []const u8) ?usize {
        var index: usize = 0;
        var node = db.root;

        for (str) |c| {
            const child = node.children[c] orelse return null;
            for (node.children, 0..) |sibling, sibling_c| {
                if (sibling == null) continue;
                if (sibling_c < c) {
                    index += sibling.?.number;
                }
            }
            node = child;
            if (node.is_terminal) index += 1;
        }

        return index;
    }

    fn writeDafsa(db: *const DafsaBuilder, writer: *std.Io.Writer) !void {
        try writer.writeAll("const dafsa = [_]Node{\n");

        // write root
        try writer.writeAll("    .{ .char = 0, .end_of_word = false, .end_of_list = true, .number = 0, .child_index = 1 },\n");

        var queue: std.ArrayList(*Node) = .empty;
        defer queue.deinit(db.allocator);

        var child_indexes: std.AutoHashMapUnmanaged(*Node, usize) = .empty;
        defer child_indexes.deinit(db.allocator);

        try child_indexes.ensureTotalCapacity(db.allocator, @intCast(db.edgeCount()));

        var first_available_index: usize = db.root.numDirectChildren() + 1;
        first_available_index = try db.writeDafsaChildren(db.root, writer, &queue, &child_indexes, first_available_index);

        while (queue.items.len > 0) {
            // TODO: something with better time complexity
            const node = queue.orderedRemove(0);

            first_available_index = try db.writeDafsaChildren(node, writer, &queue, &child_indexes, first_available_index);
        }

        try writer.writeAll("};\n");
    }

    fn writeDafsaChildren(
        db: *const DafsaBuilder,
        node: *Node,
        writer: *std.Io.Writer,
        queue: *std.ArrayList(*Node),
        child_indexes: *std.AutoHashMapUnmanaged(*Node, usize),
        first_available_index: usize,
    ) !usize {
        var cur_available_index = first_available_index;
        const num_children = node.numDirectChildren();
        var child_i: usize = 0;
        for (node.children, 0..) |maybe_child, c_usize| {
            const child = maybe_child orelse continue;
            const c: u8 = @intCast(c_usize);
            const is_last_child = child_i == num_children - 1;

            if (!child_indexes.contains(child)) {
                const child_num_children = child.numDirectChildren();
                if (child_num_children > 0) {
                    child_indexes.putAssumeCapacityNoClobber(child, cur_available_index);
                    cur_available_index += child_num_children;
                }
                try queue.append(db.allocator, child);
            }

            try writer.print(
                "    .{{ .char = '{c}', .end_of_word = {}, .end_of_list = {}, .number = {}, .child_index = {} }},\n",
                .{ c, child.is_terminal, is_last_child, child.number, child_indexes.get(child) orelse 0 },
            );

            child_i += 1;
        }
        return cur_available_index;
    }
};
