defmodule Assembler do
  import NimbleParsec

  label =
    ascii_string([?a..?z, ?A..?Z, ?_], min: 1)
    |> ignore(ascii_char([?:]))
    |> ignore(repeat(ascii_char([?\s])))

  data_label =
    label
    |> choice([
      integer(min: 1),
      ignore(ascii_char([?-]))
      |> integer(min: 1)
      |> tag(:neg)
    ])
    |> tag(:data_label)

  param =
    optional(ignore(ascii_string([?\s], min: 1)))
    |> ascii_string([{:not, ?,}], min: 1)

  instruction =
    ignore(ascii_string([?\s], 2))
    |> ascii_string([?a..?z], min: 1)
    |> repeat(param |> ignore(ascii_char([?,])))
    |> optional(param)
    |> tag(:insn)

  defparsec(
    :parse_line,
    choice([
      data_label,
      label |> tag(:label),
      instruction,
      eos()
    ])
  )

  @asm """
    in res
    clt $res, _self + 1, 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
    assemble(@asm)
  end

  def assemble(asm) do
    {memory, labels} =
      asm
      |> String.trim_trailing("\n")
      |> String.split("\n")
      |> Enum.map(fn line ->
        case parse_line(line) do
          {:ok, res, _, _, _, _} ->
            res

          err ->
            raise "Unable to parse line: '#{line}':\n#{inspect(err)}"
        end
      end)
      |> Enum.reduce({[], %{}}, &assemble_line/2)

    memory
    |> Enum.reverse()
    |> Enum.with_index()
    |> Enum.map(fn
      {{:expr, expr}, index} ->
        ExpressionEvaluator.eval(expr, fn
          "_self" -> index
          name -> Map.fetch!(labels, name)
        end)

      {it, _} ->
        it
    end)
  end

  def assemble_line([], acc), do: acc

  def assemble_line([label: [name]], {memory, labels}) do
    {
      memory,
      Map.put(labels, name, length(memory))
    }
  end

  def assemble_line([data_label: [name, value]], {memory, labels}) do
    {
      [value | memory],
      Map.put(labels, name, length(memory))
    }
  end

  def assemble_line([insn: insn], {memory, labels}) do
    {
      assemble_insn(insn, memory),
      labels
    }
  end

  def assemble_insn(["add" | params], memory) when length(params) == 3 do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])
    [dest, b, a, 1 + modes | memory]
  end

  def assemble_insn(["mul" | params], memory) when length(params) == 3 do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])
    [dest, b, a, 2 + modes | memory]
  end

  def assemble_insn(["in", dest], memory) do
    {[dest], modes} = parse_params([dest], [:write])
    [dest, 3 + modes | memory]
  end

  def assemble_insn(["out", src], memory) do
    {[src], modes} = parse_params([src], [:read])
    [src, 4 + modes | memory]
  end

  def assemble_insn(["jnz" | params], memory) when length(params) == 2 do
    {[target, dest], modes} = parse_params(params, [:read, :read])
    [dest, target, 5 + modes | memory]
  end

  def assemble_insn(["jez" | params], memory) when length(params) == 2 do
    {[target, dest], modes} = parse_params(params, [:read, :read])
    [dest, target, 6 + modes | memory]
  end

  def assemble_insn(["clt" | params], memory) when length(params) == 3 do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])
    [dest, b, a, 7 + modes | memory]
  end

  def assemble_insn(["ceq" | params], memory) when length(params) == 3 do
    {[a, b, dest], modes} = parse_params(params, [:read, :read, :write])
    [dest, b, a, 8 + modes | memory]
  end

  def assemble_insn(["hlt"], memory) do
    [99 | memory]
  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])
    {[{:expr, "label"}, 2, 3], 1000}
    iex> Assembler.parse_params(["1", "label", "3"], [:read, :read, :write])
    {[1, {:expr, "label"}, 3], 1100}
  """
  def parse_params(params, param_types) do
    params
    |> 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)

          _ ->
            case param do
              "$" <> param ->
                {:expr, param}

              _ ->
                {:expr, param}
            end
        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
end