defmodule AOC.Day5 do
  use AOC.Day, day: 5, trim: false

  def parse_input(lines) do
    [init, ops] =
      lines
      |> Enum.chunk_by(&(&1 == "\n"))
      |> Enum.reject(&(&1 == ["\n"]))

    {parse_init(init), parse_ops(ops)}
  end

  def parse_init(init) do
    init
    |> Enum.split(-1)
    |> elem(0)
    |> Enum.map(fn line ->
      line
      |> String.to_charlist()
      |> tl()
      |> Enum.take_every(4)
    end)
    |> Enum.zip_with(&Function.identity/1)
    |> Enum.map(fn stack ->
      Enum.reject(stack, &(&1 == ?\s))
    end)
  end

  def parse_ops(ops) do
    Enum.map(ops, fn op ->
      %{"count" => count, "from" => from, "to" => to} =
        Regex.named_captures(~r/move (?<count>\d+) from (?<from>\d+) to (?<to>\d+)\n/, op)

      %{
        count: String.to_integer(count),
        from: String.to_integer(from) - 1,
        to: String.to_integer(to) - 1
      }
    end)
  end

  def part1(input) do
    move(input, true)
  end

  def part2(input) do
    move(input, false)
  end

  def move({init, ops}, reverse) do
    ops
    |> Enum.reduce(init, fn %{count: count, from: from, to: to}, acc ->
      {acc, removed} = remove(acc, from, count)
      removed = if reverse, do: Enum.reverse(removed), else: removed
      add(acc, to, removed)
    end)
    |> Enum.map(&hd/1)
  end

  def remove(stacks, from, count) do
    {removed, new} = Enum.split(Enum.at(stacks, from), count)
    stacks = List.update_at(stacks, from, fn _ -> new end)
    {stacks, removed}
  end

  def add(stacks, to, removed) do
    added = Enum.concat(removed, Enum.at(stacks, to))
    List.update_at(stacks, to, fn _ -> added end)
  end
end