defmodule AOC.Day7 do
  use AOC.Day, day: 7

  def parse_input(lines) do
    Enum.map(lines, fn
      "$ cd " <> arg ->
        {:cd, arg}

      "$ ls" ->
        :ls

      "dir " <> arg ->
        {:dir, arg}

      line ->
        [size, file] = String.split(line)
        {:file, String.to_integer(size), file}
    end)
  end

  def part1(input) do
    input
    |> get_file_sizes()
    |> get_directory_sizes()
    |> Enum.map(fn {_, size} -> size end)
    |> Enum.filter(&(&1 <= 100_000))
    |> Enum.sum()
  end

  def part2(input) do
    file_sizes = get_file_sizes(input)

    total_size =
      file_sizes
      |> Enum.map(&elem(&1, 1))
      |> Enum.sum()

    size_needed = total_size - 40_000_000

    file_sizes
    |> get_directory_sizes()
    |> Enum.map(&elem(&1, 1))
    |> Enum.filter(fn size -> size >= size_needed end)
    |> Enum.min()
  end

  def get_file_sizes(input) do
    Enum.reduce(input, {[], []}, fn
      {:cd, "/"}, {_, sizes} ->
        {["/"], sizes}

      {:cd, ".."}, {[_ | path], sizes} ->
        {path, sizes}

      {:cd, dir}, {path, sizes} ->
        {[dir | path], sizes}

      {:file, size, file}, {path, sizes} ->
        {path, [{[file | path], size} | sizes]}

      _, acc ->
        acc
    end)
    |> elem(1)
  end

  def get_directory_sizes(file_sizes) do
    file_sizes
    |> Enum.reduce(%{}, fn {[_ | path], size}, acc ->
      path = Enum.reverse(path)
      update_directory_size(acc, size, path)
    end)
  end

  def update_directory_size(sizes, file_size, file_path) do
    Enum.reduce(file_path, {sizes, []}, fn dir, {sizes, cur_path} ->
      cur_path = [dir | cur_path]
      {Map.update(sizes, cur_path, file_size, &(&1 + file_size)), cur_path}
    end)
    |> elem(0)
  end
end