const std = @import("std");
const Grammar = @import("grammar.zig");
const NonTerminal = @import("non-terminal.zig");
const Rule = @import("rule.zig");
const Character = @import("character.zig").Character;
const Generator = @import("generator.zig").Generator;
const gss = @import("gss.zig");
const Pex = @import("pex/root.zig");

pub const Result = struct {
    const Info = struct {
        const Self = @This();

        max_heap_size: usize,
        number_of_nodes: usize,

        pub fn format(
            self: *const Self,
            comptime fmt: []const u8,
            options: std.fmt.FormatOptions,
            writer: anytype,
        ) !void {
            _ = fmt;
            _ = options;

            try writer.print("{:.2}, {} nodes", .{
                std.fmt.fmtIntSizeBin(self.max_heap_size),
                self.number_of_nodes,
            });
        }
    };

    is_accepted: bool,
    info: Info,
};

fn OpaqueContext(T: type) type {
    return struct {
        pub fn hash(_: @This(), h: T) u64 {
            return @intFromPtr(h);
        }

        pub fn eql(_: @This(), a: T, b: T) bool {
            return a == b;
        }
    };
}


pub fn reaches_end_of_entry(grammar: *Grammar, node: *gss.Node(*Pex.Node)) bool {
    if (node.parents.items.len == 0) {
        return true;
    }

    var pex_node = node.state;
    while (pex_node.next) |next| : (pex_node = next) {
        switch (pex_node.instruction) {
            .check => return false,
            .call => |n| {
                if (!grammar.non_terminal_by_id(n).first.is_set(Character.EPSILON)) {
                    return false;
                }
            },
            else => {},
        }
    }

    for (node.parents.items) |parent| {
        if (parent.state.next) |next| {
            parent.state = next;

            if (reaches_end_of_entry(grammar, parent)) {
                return true;
            }
        }
    }

    return false;
}

pub fn check(
    pex: *const Pex,
    input: []const u8,
    inner_allocator: std.mem.Allocator
) !Result {
    var arena = std.heap.ArenaAllocator.init(inner_allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    var queues = [2]std.ArrayList(*gss.Node(*Pex.Node)) {
        std.ArrayList(*gss.Node(*Pex.Node)).init(allocator),
        std.ArrayList(*gss.Node(*Pex.Node)).init(allocator),
    };
    defer for (queues) |queue| queue.deinit();
    var processing_queue = &queues[0];

    var graph = gss.Graph(*Pex.Node).init(allocator);

    for (pex.entry_point().heads.items) |head| {
        try processing_queue.append(try graph.add_toplevel(head));
    }

    var node_cache = std.HashMap(
        *Pex.Node,
        *gss.Node(*Pex.Node),
        OpaqueContext(*Pex.Node),
        80
    ).init(allocator);

    var forward_queue = std.ArrayList(struct { *Pex.Node, *gss.Node(*Pex.Node) }).init(allocator);
    var next_processing_queue = &queues[1];

    for (input) |character| {
        for (processing_queue.items) |base| {
            try forward_queue.append(.{ base.state, base });
            try node_cache.put(base.state, base);

            while (forward_queue.pop()) |checkpoint| {
                const pex_node, const last_node = checkpoint;

                switch (pex_node.instruction) {
                    .check => |t| if (t == character) {
                        if (pex_node.next) |next| {
                            const node = try graph.clone(last_node, next);
                            try next_processing_queue.append(node);
                        }
                    },
                    .@"return" => {
                        for (last_node.parents.items) |parent| {
                            if (parent.state.next) |next| {
                                try forward_queue.append(.{next, parent});
                            }
                        }
                    },
                    .call => |id| {
                        const non_terminal = pex.grammar.non_terminal_by_id(id);
                        for (non_terminal.rules(), pex.blocks[id].heads.items) |child_rule, head| {
                            if (
                                child_rule.first.is_set(character) or
                                (child_rule.first.is_set(Character.EPSILON) and non_terminal.follows.is_set(character))
                            ) {

                                if (node_cache.get(head)) |parent| {
                                    try parent.parents.append(last_node);
                                } else {
                                    const parent = try graph.clone(last_node, pex_node);
                                    const next = try graph.push(parent, head);
                                    try node_cache.put(head, next);
                                    try forward_queue.append(.{head, next});
                                }
                            }
                        }
                    },
                    .jump => |jump_next| {
                        if (pex_node.next) |next| {
                            try forward_queue.append(.{next, last_node});
                        }

                        for (jump_next.items) |next| {
                            try forward_queue.append(.{next, last_node});
                        }
                    },
                }
            }

            node_cache.clearRetainingCapacity();
            forward_queue.clearRetainingCapacity();
        }

        const swap = processing_queue;
        processing_queue = next_processing_queue;
        next_processing_queue = swap;
        next_processing_queue.clearRetainingCapacity();
    }

    for (processing_queue.items) |node| {
        if (reaches_end_of_entry(pex.grammar, node)) {
            return .{
                .is_accepted = true,
                .info = .{
                    .max_heap_size = arena.queryCapacity(),
                    .number_of_nodes = graph.number_of_nodes,
                },
            };
        }
    }

    return .{
        .is_accepted = false,
        .info = .{
            .max_heap_size = arena.queryCapacity(),
            .number_of_nodes = graph.number_of_nodes,
        },
    };
}

test "expr" {
    const text =
      \\S -> B A
      \\A -> '+' B A
      \\A -> ''
      \\B -> D C
      \\C -> '*' D C
      \\C -> ''
      \\D -> '(' S ')'
      \\D -> 'a'
      \\D -> 'b'
      ;

    const input = "b+a*b";

    const allocator = std.testing.allocator;

    var grammar = try Grammar.parse("S", text, allocator);
    defer grammar.deinit();

    try std.testing.expect(try check(&grammar, input, allocator));
}

test "simple 0 - success" {
    const text =
    \\S -> A S 'd'
    \\S -> B S
    \\S -> ''
    \\A -> 'a'
    \\A -> 'c'
    \\B -> 'a'
    \\B -> 'b'
    ;

    const input = "aad";

    const allocator = std.testing.allocator;

    var grammar = try Grammar.parse("S", text, allocator);
    defer grammar.deinit();

    try std.testing.expect(try check(&grammar, input, allocator));
}

test "simple 0 - fail" {
    const text =
    \\S -> A S 'd'
    \\S -> B S
    \\S -> ''
    \\A -> 'a'
    \\A -> 'c'
    \\B -> 'a'
    \\B -> 'b'
    ;

    const input = "accd";

    const allocator = std.testing.allocator;

    var grammar = try Grammar.parse("S", text, allocator);
    defer grammar.deinit();

    try std.testing.expect(!try check(&grammar, input, allocator));
}

test "simple 0 - fuzzy" {
    const text =
    \\S -> A S 'd'
    \\S -> B S
    \\S -> ''
    \\A -> 'a'
    \\A -> 'c'
    \\B -> 'a'
    \\B -> 'b'
    ;

    const allocator = std.testing.allocator;

    var grammar = try Grammar.parse("S", text, allocator);
    defer grammar.deinit();

    var generator = Generator(struct {
        const Self = @This();

        pub fn next(_: *Self, n: usize) usize {
            return std.crypto.random.uintLessThan(usize, n);
        }
    }){};

    for (0..100) |_| {
        const input = try generator.sentential_from_grammar(&grammar, 1000, allocator);
        defer allocator.free(input);

        try std.testing.expect(try check(&grammar, input, allocator));
    }

}