defmodule Assembler do
  @asm """
    in #30
    clt 30 #8 #31
    jnz 31 #lessThan
    ceq 30 #8 #31
    jnz 31 #equal
    out #1001
    hlt

  lessThan:
    out #999
    hlt

  equal:
    out #1000
    hlt
  """
  def test do
    @asm
    |> IO.inspect()
    |> assemble(32)
  end

  def assemble(asm, pad_length \\ 0) do
    {memory, labels} =
      asm
      |> String.split("\n")
      |> Enum.reject(&(String.length(&1) == 0))
      |> Enum.reduce({[], %{}}, fn line, {memory, labels} ->
        cond do
          Regex.match?(~r/^\s+$/, line) ->
            {memory, labels}

          String.starts_with?(line, "  ") ->
            {
              assemble_instruction(String.slice(line, 2..-1), memory),
              labels
            }

          String.ends_with?(line, ":") ->
            {
              memory,
              Map.put(labels, String.slice(line, 0..-2), length(memory))
            }

          true ->
            IO.inspect("Ignoring line: #{line}")
        end
      end)

    memory =
      Enum.map(memory, fn
        {:label, name} -> Map.fetch!(labels, name)
        it -> it
      end)

    memory =
      if pad_length > length(memory) do
        repeat(0, pad_length - length(memory)) ++ memory
      else
        memory
      end

    IO.inspect(labels)
    Enum.reverse(memory)
  end

  def repeat(_, 0), do: []
  def repeat(n, count), do: [n | repeat(n, count - 1)]

  def pow(base, 1), do: base
  def pow(base, exp), do: base * pow(base, exp - 1)

  def parse_params(params, param_types) do
    params
    |> String.split(" ")
    |> Enum.zip(param_types)
    |> Enum.with_index()
    |> Enum.map_reduce(0, fn {{param, type}, index}, modes ->
      val =
        case Regex.run(~r/^#?(\d+)$/, param) do
          [_, digits] ->
            String.to_integer(digits)

          _ ->
            [_, name] = Regex.run(~r/^#(\w+)$/, param)
            {:label, name}
        end

      modes =
        if type == :read && String.starts_with?(param, "#") do
          modes + pow(10, index + 2)
        else
          modes
        end

      {val, modes}
    end)
  end

  def assemble_instruction(insn, memory \\ [])

  def assemble_instruction("add " <> params, memory) do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])

    [dest, b, a, 1 + modes | memory]
  end

  def assemble_instruction("mul " <> params, memory) do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])

    [dest, b, a, 2 + modes | memory]
  end

  def assemble_instruction("hlt", memory) do
    [99 | memory]
  end

  def assemble_instruction("in " <> params, memory) do
    {[dest], modes} = parse_params(params, [:write])

    [dest, 3 + modes | memory]
  end

  def assemble_instruction("out " <> params, memory) do
    {[val], modes} = parse_params(params, [:read])

    [val, 4 + modes | memory]
  end

  def assemble_instruction("jnz " <> params, memory) do
    {[target, dest], modes} = parse_params(params, [:read, :read])

    [dest, target, 5 + modes | memory]
  end

  def assemble_instruction("jez " <> params, memory) do
    {[target, dest], modes} = parse_params(params, [:read, :read])

    [dest, target, 6 + modes | memory]
  end

  def assemble_instruction("clt " <> params, memory) do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])

    [dest, b, a, 7 + modes | memory]
  end

  def assemble_instruction("ceq " <> params, memory) do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])

    [dest, b, a, 8 + modes | memory]
  end
end