defmodule Clickhousex.Codec.Binary.Extractor do
  @moduledoc """
  Allows modules that `use` this module to create efficient extractor functions that speak clickhouse's binary protocol.

  To define extractors, annotate a function with the `extract` attribute like this:


      @extract length: :varint
      def extract_length(<<data::binary>>, length, other_param) do
        do_something_with_length(data, length, other_param)
      end

      def do_something_with_length(_data, length, other_param) do
        {other_param, length}
      end

  In the above example, a function named `extract_length/2` will be created, which, when passed a binary, will
  extract the length varint from it, and call the function above, passing the unparsed part of the binary and the extracted
  length varint to it.

  Usage looks like this

      {:ok, binary_from_network} = :gen_tcp.recv(conn, 0)
      {:this_is_passed_along, length} = extract_length(binary_from_network, :this_is_passed_along)


  If there isn't enough data to parse, a resume tuple is returned. The second element of the tuple is a function that when
  called with more data, picks up the parse operation where it left off.


      {:resume, resume_fn} = extract_length(<<>>, :this_is_passed_along)
      {:ok, data} = :gen_tcp.recv(conn, 0)
      {:this_is_passed_along, length} = resume_fn.(data)


  # Performance
  All functions generated by this module take advantage of binary optimizations, resuse match contexts and won't create sub-binaries.

  # Completeness
  The following extractors are implemented:

    1. Variable length integers `:varint`
    1. Signed integers: `:i8`, `:i16`, `:i32`, `i64`
    1. Unsigned integers: `:u8`, `:u16`, `:u32`, `:u64`
    1. Floats: `:f32`, `:f64`
    1. Strings: `:string`
    1. Booleans: `:boolean`
    1. Dates: `:date`, `:datetime`
    1. Lists of the above scalar types `{:list, scalar}`
    1. Nullable instances of all the above `{:nullable, scalar}` or `{:list, {:nullable, scalar}}`
  """

  defmacro __using__(_) do
    quote do
      use Bitwise
      Module.register_attribute(__MODULE__, :extract, accumulate: true)
      Module.register_attribute(__MODULE__, :extractors, accumulate: true)
      @on_definition {unquote(__MODULE__), :on_definition}
      @before_compile unquote(__MODULE__)
    end
  end

  @doc false
  defmacro __before_compile__(env) do
    for {name, visibility, args, [extractors]} <- Module.get_attribute(env.module, :extractors),
        {arg_name, arg_type} <- extractors do
      [_ | non_binary_args] = args
      extractor_args = reject_argument(non_binary_args, arg_name)

      landing_call =
        quote do
          unquote(name)(rest, unquote_splicing(non_binary_args))
        end

      extractor_fn_name = unique_name(name)

      jump_functions =
        build_jump_fn(name, extractor_fn_name, extractor_args)
        |> rewrite_visibility(visibility)
        |> collapse_blocks()

      extractors =
        arg_type
        |> build_extractor(arg_name, extractor_fn_name, landing_call, args)
        |> rewrite_visibility(visibility)

      quote do
        unquote_splicing(jump_functions)
        unquote(extractors)
      end
    end
    |> collapse_blocks()
  end

  @doc false
  def on_definition(env, visibility, name, args, _guards, _body) do
    extractors = Module.get_attribute(env.module, :extract)
    Module.delete_attribute(env.module, :extract)
    Module.put_attribute(env.module, :extractors, {name, visibility, args, extractors})
  end

  defp build_jump_fn(base_fn_name, extractor_fn_name, extractor_args) do
    quote do
      def unquote(base_fn_name)(<<>>, unquote_splicing(extractor_args)) do
        {:resume, &unquote(extractor_fn_name)(&1, unquote_splicing(extractor_args))}
      end

      def unquote(base_fn_name)(<<rest::binary>>, unquote_splicing(extractor_args)) do
        unquote(extractor_fn_name)(rest, unquote_splicing(extractor_args))
      end
    end
  end

  defp build_extractor(:varint, arg_name, extractor_name, landing_call, [_ | non_binary_args]) do
    extractor_args = reject_argument(non_binary_args, arg_name)
    int_variable = Macro.var(arg_name, nil)

    vars = quote do: [a, b, c, d, e, f, g, h, i, j]

    # ZigZag encoding is defined for arbitrary sized integers, but for
    # our purposes up to 10 parts are enough. Let's unroll the decoding loop.
    extractor_clauses =
      for parts_count <- 1..10 do
        vars_for_clause = Enum.take(vars, parts_count)
        pattern = varint_pattern(vars_for_clause)
        decoding = varint_decoding(vars_for_clause)

        quote do
          def unquote(extractor_name)(unquote(pattern), unquote_splicing(extractor_args)) do
            unquote(int_variable) = unquote(decoding)
            unquote(landing_call)
          end
        end
      end

    quote do
      def unquote(extractor_name)(<<>>, unquote_splicing(extractor_args)) do
        {:resume, &unquote(extractor_name)(&1, unquote_splicing(extractor_args))}
      end

      unquote_splicing(extractor_clauses)

      def unquote(extractor_name)(<<rest::binary>>, unquote_splicing(extractor_args)) do
        {:resume, fn more_data -> unquote(extractor_name)(rest <> more_data, unquote_splicing(extractor_args)) end}
      end
    end
  end

  # `vars` are variables for binding varint parts, from high to low
  defp varint_pattern([_ | _] = vars) do
    [last | rest] = Enum.reverse(vars)
    tag = quote do: 1 :: size(1)
    init = quote do: [0 :: size(1), unquote(last) :: size(7), rest :: binary]
    patterns = Enum.reduce(rest, init, &[tag, quote(do: unquote(&1) :: size(7)) | &2])
    {:<<>>, [], patterns}
  end

  # `vars` are varint parts, from high to low
  defp varint_decoding([_ | _] = vars) do
    vars
    |> Enum.reverse()
    |> Enum.with_index()
    |> Enum.map(fn
      {var, 0} -> var
      {var, index} -> {:<<<, [], [var, index * 7]}
    end)
    |> Enum.reduce(&{:|||, [], [&2, &1]})
  end

  @int_extractors [
    {:i64, :signed, 64},
    {:u64, :unsigned, 64},
    {:i32, :signed, 32},
    {:u32, :unsigned, 32},
    {:i16, :signed, 16},
    {:u16, :unsigned, 16},
    {:i8, :signed, 8},
    {:u8, :unsigned, 8}
  ]

  for {type_name, signed, width} <- @int_extractors do
    defp build_extractor(unquote(type_name), arg_name, extractor_name, landing_call, [_ | args]) do
      extractor_args = reject_argument(args, arg_name)
      value_variable = Macro.var(arg_name, nil)
      width = unquote(width)
      signedness = Macro.var(unquote(signed), nil)

      match =
        quote do
          <<unquote(value_variable)::little-unquote(signedness)-size(unquote(width)), rest::binary>>
        end

      quote do
        def unquote(extractor_name)(unquote(match), unquote_splicing(extractor_args)) do
          unquote(landing_call)
        end

        def unquote(extractor_name)(<<>>, unquote_splicing(extractor_args)) do
          {:resume, &unquote(extractor_name)(&1, unquote_splicing(extractor_args))}
        end

        def unquote(extractor_name)(<<data::binary>>, unquote_splicing(extractor_args)) do
          {:resume, &unquote(extractor_name)(data <> &1, unquote_splicing(extractor_args))}
        end
      end
    end
  end

  # Float extractors
  for width <- [32, 64],
      type_name = :"f#{width}" do
    defp build_extractor(unquote(type_name), arg_name, extractor_name, landing_call, [_ | args]) do
      extractor_args = reject_argument(args, arg_name)
      value_variable = Macro.var(arg_name, nil)
      width = unquote(width)

      quote do
        def unquote(extractor_name)(
              <<unquote(value_variable)::little-signed-float-size(unquote(width)), rest::binary>>,
              unquote_splicing(extractor_args)
            ) do
          unquote(landing_call)
        end

        def unquote(extractor_name)(<<>>, unquote_splicing(extractor_args)) do
          {:resume, &unquote(extractor_name)(&1, unquote_splicing(extractor_args))}
        end

        def unquote(extractor_name)(<<rest::binary>>, unquote_splicing(extractor_args)) do
          {:resume, &unquote(extractor_name)(rest <> &1, unquote_splicing(extractor_args))}
        end
      end
    end
  end

  defp build_extractor(:boolean, arg_name, extractor_name, landing_call, [_ | args]) do
    extractor_args = reject_argument(args, arg_name)
    value_variable = Macro.var(arg_name, nil)

    quote do
      def unquote(extractor_name)(<<1, rest::binary>>, unquote_splicing(extractor_args)) do
        unquote(value_variable) = true
        unquote(landing_call)
      end

      def unquote(extractor_name)(<<0, rest::binary>>, unquote_splicing(extractor_args)) do
        unquote(value_variable) = false
        unquote(landing_call)
      end

      def unquote(extractor_name)(<<>>, unquote_splicing(extractor_args)) do
        {:resume, &unquote(extractor_name)(&1, unquote_splicing(extractor_args))}
      end
    end
  end

  defp build_extractor(:date, arg_name, extractor_name, landing_call, [_ | args]) do
    extractor_args = reject_argument(args, arg_name)
    value_variable = Macro.var(arg_name, nil)

    quote do
      def unquote(extractor_name)(
            <<days_since_epoch::little-unsigned-size(16), rest::binary>>,
            unquote_splicing(extractor_args)
          ) do
        {:ok, date} = Date.new(1970, 01, 01)

        unquote(value_variable) = Date.add(date, days_since_epoch)
        unquote(landing_call)
      end

      def unquote(extractor_name)(<<>>, unquote_splicing(extractor_args)) do
        {:resume, &unquote(extractor_name)(&1, unquote_splicing(extractor_args))}
      end

      def unquote(extractor_name)(<<rest::binary>>, unquote_splicing(extractor_args)) do
        {:resume, &unquote(extractor_name)(rest <> &1, unquote_splicing(extractor_args))}
      end
    end
  end

  defp build_extractor(:datetime, arg_name, extractor_name, landing_call, [_ | args]) do
    extractor_args = reject_argument(args, arg_name)
    value_variable = Macro.var(arg_name, nil)

    quote do
      def unquote(extractor_name)(
            <<seconds_since_epoch::little-unsigned-size(32), rest::binary>>,
            unquote_splicing(extractor_args)
          ) do
        {:ok, date_time} = NaiveDateTime.new(1970, 1, 1, 0, 0, 0)

        unquote(value_variable) = NaiveDateTime.add(date_time, seconds_since_epoch)
        unquote(landing_call)
      end

      def unquote(extractor_name)(<<>>, unquote_splicing(extractor_args)) do
        {:resume, &unquote(extractor_name)(&1, unquote_splicing(extractor_args))}
      end

      def unquote(extractor_name)(<<rest::binary>>, unquote_splicing(extractor_args)) do
        {:resume, &unquote(extractor_name)(rest <> &1, unquote_splicing(extractor_args))}
      end
    end
  end

  defp build_extractor({:nullable, type}, arg_name, extractor_name, landing_call, [_ | non_binary_args] = args) do
    extractor_args = reject_argument(non_binary_args, arg_name)
    value_variable = Macro.var(arg_name, nil)
    value_extractor_name = :"#{extractor_name}_value"

    value_extractors =
      type
      |> build_extractor(arg_name, value_extractor_name, landing_call, args)
      |> collapse_blocks()

    quote do
      unquote_splicing(value_extractors)

      def unquote(extractor_name)(<<>>, unquote_splicing(extractor_args)) do
        {:resume, &unquote(extractor_name)(&1, unquote_splicing(extractor_args))}
      end

      def unquote(extractor_name)(<<0, rest::binary>>, unquote_splicing(extractor_args)) do
        unquote(value_extractor_name)(rest, unquote_splicing(extractor_args))
      end

      def unquote(extractor_name)(<<1, rest::binary>>, unquote_splicing(extractor_args)) do
        unquote(value_variable) = nil
        unquote(landing_call)
      end
    end
  end

  defp build_extractor(:string, arg_name, extractor_name, landing_call, [binary_arg | non_binary_args]) do
    extractor_args = reject_argument(non_binary_args, arg_name)

    length_variable_name = unique_name("string_length")
    length_variable = Macro.var(length_variable_name, nil)
    length_extractor_name = :"#{extractor_name}_length"
    length_extractor_args = extractor_args

    length_landing_call =
      quote do
        unquote(extractor_name)(rest, unquote_splicing(extractor_args), unquote(length_variable))
      end

    length_extractors =
      build_extractor(
        :varint,
        length_variable_name,
        length_extractor_name,
        length_landing_call,
        [binary_arg | length_extractor_args] ++ [length_variable]
      )
      |> collapse_blocks()

    value_arg = Macro.var(arg_name, nil)

    # The string extractor call chain looks like this:
    # top_level function -> length_extractor -> value_extractor
    quote do
      # Size exctractors
      unquote_splicing(length_extractors)

      # Value extractors

      # Empty string optimization, prevents concatenating large data to an empty string and
      # reallocating the large data
      def unquote(extractor_name)(<<>>, unquote_splicing(extractor_args), unquote(length_variable)) do
        {:resume, &unquote(extractor_name)(&1, unquote_splicing(extractor_args), unquote(length_variable))}
      end

      def unquote(extractor_name)(<<rest::binary>>, unquote_splicing(extractor_args), unquote(length_variable)) do
        case rest do
          <<unquote(value_arg)::binary-size(unquote(length_variable)), rest::binary>> ->
            unquote(landing_call)

          _ ->
            {:resume, &unquote(extractor_name)(rest <> &1, unquote_splicing(extractor_args), unquote(length_variable))}
        end
      end

      # Starts the size extractor chain
      def unquote(extractor_name)(<<b::binary>>, unquote_splicing(extractor_args)) do
        unquote(length_extractor_name)(b, unquote_splicing(extractor_args))
      end
    end
  end

  defp build_extractor({:array, item_type}, arg_name, extractor_name, landing_call, args) do
    build_extractor({:list, item_type}, arg_name, extractor_name, landing_call, args)
  end

  defp build_extractor({:list, item_type}, arg_name, extractor_name, landing_call, [binary_arg | non_binary_args]) do
    extractor_args = reject_argument(non_binary_args, arg_name)

    length_extractor_name = :"#{extractor_name}_list_length"
    length_name = :length |> unique_name()
    length_variable = length_name |> Macro.var(nil)
    length_extractor_args = [binary_arg | extractor_args] ++ [length_variable]

    list_extractor_name = unique_name("#{extractor_name}_list")
    item_name = :item |> unique_name()
    item_variable = Macro.var(item_name, nil)

    item_accumulator_variable = Macro.var(arg_name, nil)
    count_variable = Macro.var(:"#{extractor_name}_count", nil)
    item_extractor_name = unique_name("#{extractor_name}_item")
    item_extractor_call_args = extractor_args ++ [count_variable, item_accumulator_variable]
    item_extractor_args = [binary_arg] ++ item_extractor_call_args
    list_extractor_args = extractor_args

    length_landing_call =
      quote do
        unquote(item_extractor_name)(rest, unquote_splicing(extractor_args), unquote(length_variable), [])
      end

    list_landing_call =
      quote do
        unquote(list_extractor_name)(
          rest,
          unquote_splicing(list_extractor_args),
          unquote(count_variable) - 1,
          unquote(item_variable),
          unquote(item_accumulator_variable)
        )
      end

    item_extractors =
      item_type
      |> build_extractor(item_name, item_extractor_name, list_landing_call, item_extractor_args)
      |> collapse_blocks

    length_extractors =
      :varint
      |> build_extractor(length_name, length_extractor_name, length_landing_call, length_extractor_args)
      |> collapse_blocks()

    quote do
      def unquote(extractor_name)(<<>>, unquote_splicing(extractor_args)) do
        {:resume, &unquote(extractor_name)(&1, unquote_splicing(extractor_args))}
      end

      # Starts the chain by calling the length extractor
      def unquote(extractor_name)(<<rest::binary>>, unquote_splicing(extractor_args)) do
        unquote(length_extractor_name)(rest, unquote_splicing(extractor_args))
      end

      unquote_splicing(length_extractors)
      unquote_splicing(item_extractors)

      # This clause matches when we've extracted all items (remaining count is 0)
      def unquote(list_extractor_name)(
            <<rest::binary>>,
            unquote_splicing(list_extractor_args),
            0,
            unquote(item_variable),
            unquote(item_accumulator_variable)
          ) do
        unquote(item_accumulator_variable) = Enum.reverse([unquote(item_variable) | unquote(item_accumulator_variable)])
        unquote(landing_call)
      end

      # This matches when there's more work to do. It accumulates the extracted item
      # and calls the item extractor again
      def unquote(list_extractor_name)(
            <<rest::binary>>,
            unquote_splicing(list_extractor_args),
            unquote(count_variable),
            unquote(item_variable),
            unquote(item_accumulator_variable)
          ) do
        unquote(item_accumulator_variable) = [unquote(item_variable) | unquote(item_accumulator_variable)]
        unquote(item_extractor_name)(rest, unquote_splicing(item_extractor_call_args))
      end
    end
  end

  # Helper functions

  defp rewrite_visibility(ast, :def) do
    ast
  end

  defp rewrite_visibility(ast, :defp) do
    Macro.prewalk(ast, fn
      {:def, context, rest} -> {:defp, context, rest}
      other -> other
    end)
  end

  defp collapse_blocks({:__block__, _, defs}) do
    defs
  end

  defp collapse_blocks(ast) when is_list(ast) do
    Enum.reduce(ast, [], fn
      {:__block__, _context, clauses}, acc ->
        acc ++ clauses

      _, acc ->
        acc
    end)
    |> Enum.reverse()
  end

  defp collapse_blocks(ast) do
    [ast]
  end

  defp reject_argument(args, arg_name) do
    Enum.reject(args, fn
      {^arg_name, _, _} -> true
      _ -> false
    end)
  end

  defp unique_name(base_name) do
    unique = System.unique_integer([:positive, :monotonic])
    :"#{base_name}_#{unique}"
  end
end
