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

  lessThan:
    out #999
    hlt

  equal:
    out #1000
    hlt

  res: 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)

    memory =
      Enum.map(memory, fn
        {:label, name} -> Map.fetch!(labels, name)
        it -> it
      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