defmodule Rephex.State do
  @moduledoc """
  Define Rephex state by `use`.
  Defined state must be initialized by `init/1` in `Phoenix.LiveView.mount/3`.

  ## Example

      defmodule ExampleWeb.State do
        @type t :: %{
                count: integer(),
                add_twice_async: %AsyncResult{}
              }

        @initial_state %{
          count: 0,
        }

        use Rephex.State, initial_state: @initial_state

        def add_count(socket, %{amount: amount} = _payload) when is_integer(amount) do
          # You can use `update_state`, `update_state_in` and `put_state_in` to update state
          update_state_in(socket, [:count], &(&1 + amount))
        end
      end

      defmodule ExampleWeb.AccountLive.Index do
        use ExampleWeb, :live_view
        use Rephex.LiveView

        @impl true
        def mount(_params, _session, socket) do
          {:ok, socket |> ExampleWeb.State.init()}
        end
      end
  """
  alias Phoenix.LiveView.Socket

  @root Rephex.root()
  @selector_keys_key :__selector_keys__

  def init(%Socket{} = socket, %{} = initial_state) do
    selector_keys = get_selector_keys_as_ordered(initial_state)

    initial_state = Map.put(initial_state, @selector_keys_key, selector_keys)

    socket
    |> Phoenix.Component.assign(@root, initial_state)
    |> update_selectors()
  end

  defp get_selector_keys_as_ordered(%{} = state) do
    selector_module_to_key =
      state
      |> Enum.filter(fn {_k, v} -> is_struct(v, Rephex.Selector) end)
      |> Enum.map(fn {k, v} -> {v.module, k} end)
      |> Enum.into(%{})

    selector_module_to_key
    |> Map.keys()
    |> Rephex.Selector.sort_selectors_by_update_order()
    |> Enum.map(&Map.fetch!(selector_module_to_key, &1))
  end

  def update_selectors(%Socket{} = socket) do
    selector_keys = Rephex.State.Assigns.get_state(socket)[@selector_keys_key]

    if selector_keys == nil do
      raise """
      Rephex.State.Assigns.update_selectors/1:
      Selector modules are not found in the state.
      """
    end

    if not is_list(selector_keys) do
      raise """
      Rephex.State.Assigns.update_selectors/1:
      Selector modules must be a list.
      """
    end

    Enum.reduce(selector_keys, socket, fn key, socket ->
      selector = Rephex.State.Assigns.get_state_in(socket, [key])
      state = socket.assigns[@root]

      Rephex.Selector.update(selector, state)
      |> case do
        nil -> socket
        new_selector -> Rephex.State.Assigns.put_state_in(socket, [key], new_selector)
      end
    end)
  end

  defmacro __using__(opt) do
    initial_state = Keyword.fetch!(opt, :initial_state)

    quote do
      import Rephex.State.Assigns

      def init(%Socket{} = socket) do
        socket |> Rephex.State.init(unquote(initial_state))
      end
    end
  end
end

defmodule Rephex.State.Assigns do
  alias Phoenix.LiveView.Socket

  @root Rephex.root()
  @selector_keys_key :__selector_keys__

  @doc """
  Update Rephex state.

  ## Example

      def add_count(socket, %{amount: amount} = _payload) do
        update_state(socket, fn state -> %{state | count: state.count + amount} end)
      end

  """
  def update_state(%Socket{parent_pid: parent_pid}, _fun) when parent_pid != nil,
    do: raise("Use this function only in LiveView (root).")

  def update_state(%Socket{} = socket, fun) when is_function(fun, 1) do
    # I should keep meta data in the state
    selector_keys = Rephex.State.Assigns.get_state(socket)[@selector_keys_key]

    socket
    |> Phoenix.Component.update(@root, fun)
    |> Phoenix.Component.update(@root, &Map.put(&1, @selector_keys_key, selector_keys))
    |> Rephex.State.update_selectors()
  end

  @doc """
  Update Rephex state by `put_in/3`.

  ## Example

      def put_value(socket, %{key: k, value: v} = _payload) do
        put_state_in(socket, [:items, k], v)
      end

  """
  def put_state_in(%Socket{parent_pid: parent_pid}, _keys, _value) when parent_pid != nil,
    do: raise("Use this function only in LiveView (root).")

  def put_state_in(%Socket{} = socket, keys, value) when is_list(keys) do
    update_state(socket, &put_in(&1, keys, value))
  end

  @doc """
  Update Rephex state by `update_in/3`.

  ## Example

      import Rephex.State.Assigns

      def mlt_count(socket, %{mlt: mlt} = _payload) do
        update_state_in(socket, [:count], &(&1 * mlt))
      end
  """
  def update_state_in(%Socket{parent_pid: parent_pid}, _keys, _fun) when parent_pid != nil,
    do: raise("Use this function only in LiveView (root).")

  def update_state_in(%Socket{} = socket, keys, fun) when is_list(keys) and is_function(fun, 1) do
    update_state(socket, &update_in(&1, keys, fun))
  end

  def get_state(%Socket{parent_pid: parent_pid}) when parent_pid != nil,
    do: raise("Use this function only in LiveView (root).")

  def get_state(%Socket{} = socket) do
    socket.assigns[@root]
  end

  def get_state_in(%Socket{parent_pid: parent_pid}, _keys) when parent_pid != nil,
    do: raise("Use this function only in LiveView (root).")

  def get_state_in(%Socket{} = socket, keys) when is_list(keys) do
    get_state(socket) |> get_in(keys)
  end
end