defmodule Day9 do
  def run do
    memory =
      File.read!("data.txt")
      |> String.trim()
      |> String.graphemes()
      |> Enum.map(&String.to_integer(&1))
      |> Enum.with_index()
      |> Enum.reduce([], fn {num, idx}, acc ->
        new_num = if rem(idx, 2) == 0, do: trunc(idx / 2), else: nil
        Enum.concat(acc, List.duplicate(new_num, num))
      end)

    dbg(memory)

    part1 = memory |> defrag1() |> checksum()
    dbg(part1)

    part2 = memory |> defrag2() |> checksum()
    dbg(part2)
  end

  defp checksum(memory) do
    Enum.with_index(memory)
    |> Enum.reduce(0, fn {el, idx}, acc ->
      if el, do: acc + idx * el, else: acc
    end)
  end

  defp defrag1(memory) do
    defrag1_step(%{memory: memory, free_index: 0, block_index: length(memory) - 1})
  end

  # 2 pointers, one tracking free space, one tracking blocks
  # increase the free one until we find a spot, decrease the block one until we find a block
  # if we have a free spot and a block, swap em
  defp defrag1_step(acc) do
    free = Enum.at(acc.memory, acc.free_index)
    block = Enum.at(acc.memory, acc.block_index)

    # if we've met in the middle, we're done recursing
    if acc.free_index >= acc.block_index do
      acc.memory
    else
      case {free, block} do
        {_, nil} ->
          %{acc | block_index: acc.block_index - 1}

        {f, _} when f != nil ->
          %{acc | free_index: acc.free_index + 1}

        {nil, block} ->
          memory =
            List.replace_at(
              List.replace_at(acc.memory, acc.free_index, block),
              acc.block_index,
              nil
            )

          %{
            acc
            | memory: memory,
              block_index: acc.block_index - 1,
              free_index: acc.free_index + 1
          }
      end
      |> defrag1_step()
    end
  end

  defp defrag2(memory) do
    # break the memory up into chunks. Each time the char changes, new chunk
    chunks =
      Enum.chunk_while(
        memory,
        [],
        fn el, acc ->
          if length(acc) >= 1 and el != hd(acc) do
            {:cont, Enum.reverse(acc), [el]}
          else
            {:cont, [el | acc]}
          end
        end,
        fn
          [] -> {:cont, []}
          acc -> {:cont, Enum.reverse(acc), []}
        end
      )

    # run it through the defragger and then flatten it
    defrag2_step(%{chunks: chunks, index: length(chunks) - 1})
    |> List.flatten()
  end

  defp defrag2_step(args) do
    chunks = args.chunks
    index = args.index

    # if we are at 0, stop recursing
    if index == 0 do
      chunks
    else
      file = Enum.at(chunks, index)

      # if this file is nils (empty space) skip it. recruse, picking next lowest chunk
      if is_nil(Enum.at(file, 0)) do
        defrag2_step(%{chunks: chunks, index: index - 1})
      else
        file_size = length(file)

        # find a chunk of empty space big enough starting from the beginning
        empty_chunk_index =
          Enum.find_index(chunks, fn chunk ->
            # if it's an empty space chunk and its big enough
            is_nil(Enum.at(chunk, 0)) and length(chunk) >= file_size
          end)

        if empty_chunk_index && empty_chunk_index < index do
          empty_chunk = Enum.at(chunks, empty_chunk_index)
          new_free_space_size = length(empty_chunk) - file_size

          # first replace the file with free space
          temp1 =
            List.replace_at(chunks, index, List.duplicate(nil, file_size))
            |> List.replace_at(empty_chunk_index, file)

          # if our file was smaller than the free space
          if new_free_space_size > 0 do
            # check the chunk _after_ the free space.
            next_chunk = Enum.at(chunks, empty_chunk_index + 1)

            new_list =
              if is_nil(Enum.at(next_chunk, 0)) do
                # if it is also free space, tack this on to it
                List.replace_at(
                  temp1,
                  empty_chunk_index + 1,
                  Enum.concat(next_chunk, List.duplicate(nil, new_free_space_size))
                )
              else
                # otherwise, create a new chunk after this one, filling up the space
                List.insert_at(
                  temp1,
                  empty_chunk_index + 1,
                  List.duplicate(nil, new_free_space_size)
                )
              end

            defrag2_step(%{chunks: new_list, index: index})
          else
            defrag2_step(%{chunks: temp1, index: index})
          end
        else
          # no free space for this file. recurse, check the next chunk
          defrag2_step(%{chunks: chunks, index: index - 1})
        end
      end
    end
  end
end

Day9.run()