defmodule Day14 do
  use Bitwise

  @example """
  mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
  mem[8] = 11
  mem[7] = 101
  mem[8] = 0
  """

  @example2 """
  mask = 000000000000000000000000000000X1001X
  mem[42] = 100
  mask = 00000000000000000000000000000000X0XX
  mem[26] = 1
  """

  def part1(example \\ false) do
    parse_input(example)
    |> set_mem_values()
    |> Map.values()
    |> Enum.sum()
  end

  def part2(example \\ false) do
    parse_input(example)
    |> set_v2_mem_values()
    |> Map.values()
    |> Enum.sum()
  end

  def parse_input(example) do
    case example do
      true -> @example
      1 -> @example
      2 -> @example2
      _ -> File.read!("lib/day14/input.txt")
    end
    |> String.trim()
    |> String.split("\n")
    |> Enum.map(fn
      "mask = " <> mask ->
        String.to_charlist(mask)

      line ->
        [_, address, value] = Regex.run(~r/^mem\[(\d+)\] = (\d+)$/, line)
        {String.to_integer(address), String.to_integer(value)}
    end)
  end

  def set_mem_values(instructions, mask \\ nil, map \\ %{})

  def set_mem_values([], _, map) do
    map
  end

  def set_mem_values([new_mask | rest], _, map) when is_list(new_mask) do
    set_mem_values(rest, new_mask, map)
  end

  def set_mem_values([{address, value} | rest], mask, map) when not is_nil(mask) do
    masked_value = get_masked_value(mask, value)
    set_mem_values(rest, mask, Map.put(map, address, masked_value))
  end

  def get_masked_value(mask, value) do
    mask
    |> Enum.reverse()
    |> Enum.with_index()
    |> Enum.reduce(value, fn {mask_val, mask_index}, value ->
      case mask_val do
        ?X ->
          value

        ?0 ->
          value &&& bnot(1 <<< mask_index)

        ?1 ->
          value ||| 1 <<< mask_index
      end
    end)
  end

  def set_v2_mem_values(instructions, mask \\ nil, map \\ %{})

  def set_v2_mem_values([], _, map) do
    map
  end

  def set_v2_mem_values([new_mask | rest], _, map) when is_list(new_mask) do
    set_v2_mem_values(rest, new_mask, map)
  end

  def set_v2_mem_values([{address, value} | rest], mask, map) when not is_nil(mask) do
    new_map =
      calculate_all_addresses(address, mask)
      # |> IO.inspect()
      |> Enum.reduce(map, fn addr, map ->
        Map.put(map, addr, value)
      end)

    set_v2_mem_values(rest, mask, new_map)
  end

  def calculate_all_addresses(base_addr, mask) do
    base_addr
    |> Integer.to_string(2)
    |> String.pad_leading(length(mask), "0")
    |> String.to_charlist()
    |> Enum.zip(mask)
    |> Enum.map(fn
      {c, ?0} -> c
      {_, ?1} -> ?1
      {_, ?X} -> ?X
    end)
    |> expand_mask()
    |> Enum.map(fn addr ->
      addr |> List.to_string() |> String.to_integer(2)
    end)
  end

  def expand_mask([]) do
    [[]]
  end

  def expand_mask([?X | rest]) do
    rest_expanded = expand_mask(rest)

    zero = Enum.map(rest_expanded, fn mask -> [?0 | mask] end)
    one = Enum.map(rest_expanded, fn mask -> [?1 | mask] end)

    zero ++ one
  end

  def expand_mask([c | rest]) do
    rest
    |> expand_mask()
    |> Enum.map(fn mask -> [c | mask] end)
  end
end