defmodule Assembler do
  @asm """
    in res
    clt $res, 8 cmpRes
    jnz $cmpRes, lessThan
    ceq $res, 8, cmpRes
    jnz $cmpRes, equal
    out 1001
    hlt

  lessThan:
    out 999
    hlt

  equal:
    out 1000
    hlt

  res: 0
  cmpRes: 0
  """
  def test do
    @asm
    |> IO.inspect()
    |> assemble()
  end

  def assemble(asm) 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))
            }

          Regex.match?(~r/^\w+: \d+$/, line) ->
            [name, value] = String.split(line, ": ")
            value = String.to_integer(value)

            {
              [value | memory],
              Map.put(labels, name, length(memory))
            }

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

    # IO.inspect(labels)

    memory
    |> Enum.reverse()
    |> Enum.with_index()
    |> Enum.map(fn
      {{:label, name}, index} -> get_label(name, index, labels)
      {it, _} -> it
    end)
  end

  @doc """
  Get the address of a label.

  `_self` is a special label that resolves to the address of where it's being inserted into the program.

  ## Examples
    iex> Assembler.get_label("test", 1, %{"test" => 14})
    14
    iex> Assembler.get_label("_self", 1, %{})
    1
  """
  def get_label(name, index, labels)

  def get_label("_self", index, _) do
    index
  end

  def get_label(name, _, labels) do
    Map.fetch!(labels, name)
  end

  @doc """
  Raises the base to to the given power.

  ## Examples
    iex> Assembler.pow(10, 3)
    1000
    iex> Assembler.pow(5, 4)
    625
  """
  def pow(base, 1), do: base
  def pow(base, exp), do: base * pow(base, exp - 1)

  @doc """
  Parses assembly instruction parameter values and modes.

  ## Examples
    iex> Assembler.parse_params("1, 2, 3", [:read, :read, :read])
    {[1, 2, 3], 11100}
    iex> Assembler.parse_params("1, 2, 3", [:read, :read, :write])
    {[1, 2, 3], 1100}
    iex> Assembler.parse_params("1, 2, $3", [:read, :read, :write])
    {[1, 2, 3], 1100}
    iex> Assembler.parse_params("$1, 2, 3", [:read, :read, :write])
    {[1, 2, 3], 1000}
    iex> Assembler.parse_params("$label, 2, 3", [:read, :read, :write])
    {[{:label, "label"}, 2, 3], 1000}
    iex> Assembler.parse_params("1, label, 3", [:read, :read, :write])
    {[1, {:label, "label"}, 3], 1100}
  """
  def parse_params(params, param_types) do
    params
    |> String.split(",")
    |> Enum.map(&String.trim/1)
    |> 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 =
        case type do
          :read ->
            if String.starts_with?(param, "$") do
              modes
            else
              modes + pow(10, index + 2)
            end

          _ ->
            modes
        end

      {val, modes}
    end)
  end

  @doc """
  Assembles the given instruction and adds it to the memory (in reverse order).
  """
  def assemble_instruction(insn, memory \\ [])

  @doc """
  Assemble add instruction.

  ## Examples
    iex> Assembler.assemble_instruction("add 1, 2, 3")
    [3, 2, 1, 1101]
  """
  def assemble_instruction("add " <> params, memory) do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])

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

  @doc """
  Assemble multiply instruction.

  ## Examples
    iex> Assembler.assemble_instruction("mul 1, 2, 3")
    [3, 2, 1, 1102]
  """
  def assemble_instruction("mul " <> params, memory) do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])

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

  @doc """
  Assemble halt instruction.

  ## Examples
    iex> Assembler.assemble_instruction("hlt")
    [99]
  """
  def assemble_instruction("hlt", memory) do
    [99 | memory]
  end

  @doc """
  Assemble input instruction.

  ## Examples
    iex> Assembler.assemble_instruction("in 7")
    [7, 3]
  """
  def assemble_instruction("in " <> params, memory) do
    {[dest], modes} = parse_params(params, [:write])

    [dest, 3 + modes | memory]
  end

  @doc """
  Assemble input instruction.

  ## Examples
    iex> Assembler.assemble_instruction("out 7")
    [7, 104]
    iex> Assembler.assemble_instruction("out $7")
    [7, 4]
  """
  def assemble_instruction("out " <> params, memory) do
    {[val], modes} = parse_params(params, [:read])

    [val, 4 + modes | memory]
  end

  @doc """
  Assemble jump-if-non-zero instruction.

  ## Examples
    iex> Assembler.assemble_instruction("jnz 1, 2")
    [2, 1, 105]
    iex> Assembler.assemble_instruction("jnz $var, 2")
    [2, {:label, "var"}, 5]
  """
  def assemble_instruction("jnz " <> params, memory) do
    {[target, dest], modes} = parse_params(params, [:read, :read])

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

  @doc """
  Assemble jump-if-zero instruction.

  ## Examples
    iex> Assembler.assemble_instruction("jez 1, 2")
    [2, 1, 106]
    iex> Assembler.assemble_instruction("jez $var, 2")
    [2, {:label, "var"}, 6]
  """
  def assemble_instruction("jez " <> params, memory) do
    {[target, dest], modes} = parse_params(params, [:read, :read])

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

  @doc """
  Assemble compare-less-than instruction.

  ## Examples
    iex> Assembler.assemble_instruction("clt 1, 2, 3")
    [3, 2, 1, 1107]
    iex> Assembler.assemble_instruction("clt $1, $2, 3")
    [3, 2, 1, 7]
  """
  def assemble_instruction("clt " <> params, memory) do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])

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

  @doc """
  Assemble compare-equal instruction.

  ## Examples
    iex> Assembler.assemble_instruction("ceq 1, 2, 3")
    [3, 2, 1, 1108]
    iex> Assembler.assemble_instruction("ceq $1, $2, 3")
    [3, 2, 1, 8]
  """
  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