const std = @import("std");
pub const Block = @import("block.zig");
pub const Instruction = @import("instruction.zig").Instruction;
pub const Node = @import("node.zig");
pub const Grammar = @import("../grammar.zig");
const optimizer = @import("optimizer/root.zig");

const Self = @This();

grammar: *Grammar,
blocks: []Block,
node_pool: std.heap.MemoryPool(Node),

pub fn compile(grammar: *Grammar, allocator: std.mem.Allocator) !Self {
  const raw_blocks = try allocator.alloc(Block, grammar.non_terminals.len);
  defer {
    for (raw_blocks) |*block| {
      block.deinit(allocator);
    }
    allocator.free(raw_blocks);
  }

  var raw_node_pool: std.heap.MemoryPool(Node) = .init(allocator);
  defer raw_node_pool.deinit();

  for (raw_blocks, grammar.non_terminals) |*block, *non_terminal| {
    block.* = try .compile(non_terminal, &raw_node_pool, allocator);
  }

  var self: Self = .{
      .grammar = grammar,
      .blocks = try allocator.alloc(Block, grammar.non_terminals.len),
      .node_pool = .init(allocator),
  };

  for (self.blocks, raw_blocks) |*optimized, *raw| {
   try optimizer.optimize_block(raw_blocks, raw, optimized, &self.node_pool, allocator);
  }

  return self;
}

pub fn deinit(self: *Self, allocator: std.mem.Allocator) void {
  for (self.blocks) |*block| {
    block.deinit(allocator);
  }

  allocator.free(self.blocks);
  self.node_pool.deinit();
  self.* = undefined;
}

pub fn entry_point(self: *const Self) *Block {
    return &self.blocks[0];
}

test "terminals-single" {
  const allocator = std.testing.allocator;

  var grammar = try Grammar.parse("S", "S -> 'a'", allocator);
  defer grammar.deinit();

  var pex = try Self.compile(&grammar, allocator);
  defer pex.deinit(allocator);

  try std.testing.expectEqual(2, pex.blocks.len);
  try std.testing.expectEqual(1, pex.blocks[1].heads.items.len);
  try std.testing.expectEqual(
    Instruction{ .check = 'a' },
    pex.blocks[1].heads.items[0].instruction
  );
  try std.testing.expectEqual(1, pex.blocks[1].heads.items[0].next.items.len);
  try std.testing.expectEqual(
    Instruction{ .@"return" = void{} },
    pex.blocks[1].heads.items[0].next.items[0].instruction
  );
}

test "terminals-single-thread" {
  const allocator = std.testing.allocator;

  var grammar = try Grammar.parse("S", "S -> 'abc'", allocator);
  defer grammar.deinit();

  var pex = try Self.compile(&grammar, allocator);
  defer pex.deinit(allocator);

  try std.testing.expectEqual(2, pex.blocks.len);
  try std.testing.expectEqual(1, pex.blocks[1].heads.items.len);

  var current = pex.blocks[1].heads.items[0];
  for ("abc") |c| {
    try std.testing.expectEqual(
      Instruction{ .check = c },
      current.instruction,
    );
    try std.testing.expectEqual(1, current.next.items.len);
    current = current.next.items[0];
  }

  try std.testing.expectEqual(
    Instruction{ .@"return" = void{} },
    current.instruction,
  );
}

test "direct-left-recursion" {
  const allocator = std.testing.allocator;

  var grammar = try Grammar.parse("S", "S -> 'a' | S 'a'", allocator);
  defer grammar.deinit();

  var pex = try Self.compile(&grammar, allocator);
  defer pex.deinit(allocator);

  try std.testing.expectEqual(2, pex.blocks.len);
  try std.testing.expectEqual(1, pex.blocks[1].heads.items.len);

  try std.testing.expectEqual(
    Instruction{ .check = 'a' },
    pex.blocks[1].heads.items[0].instruction,
  );

  try std.testing.expectEqual(
    Instruction{ .@"return" = void{} },
    pex.blocks[1].heads.items[0].next.items[0].instruction,
  );

  try std.testing.expectEqual(
    Instruction{ .check = 'a' },
    pex.blocks[1].heads.items[0].next.items[1].next.items[0].instruction,
  );

  try std.testing.expectEqual(
    Instruction{ .@"return" = void{} },
    pex.blocks[1]
    .heads.items[0]
    .next.items[1]
    .next.items[0]
    .next.items[0]
    .instruction,
  );

  try std.testing.expect(
    pex.blocks[1]
    .heads.items[0]
    .next.items[1]
    .next.items[0]
    .next.items[1]
    .instruction == .jump
  );

  try std.testing.expectEqual(
    Instruction{ .@"return" = void{} },
    pex.blocks[1]
    .heads.items[0]
    .next.items[1]
    .next.items[0]
    .next.items[1]
    .next.items[0]
    .instruction,
  );
}