defmodule Idx do
  @moduledoc """
  Collection allowing access via dynamically created indices.

  The API is similar to `Map`. See `create_index/3` for more details
  about indices.

      iex> users = [%{name: "Bob", age: 20}, %{name: "Eve", age: 27}, %{name: "John", age: 45}]
      iex> idx = Idx.new(users, & &1.name)
      iex> Idx.get(idx, "Bob")
      %{name: "Bob", age: 20}
      iex> idx |> Enum.to_list() |> Enum.sort()
      users

  """
  alias __MODULE__.{Key, Primary}

  @behaviour Access

  @enforce_keys [Primary, :indices, :lazy_indices]
  defstruct @enforce_keys

  @type t :: %{:__struct__ => __MODULE__, any => any}

  @typedoc """
  Index function. Returns a key based on the value to be indexed.
  Must be pure (always return the same key for the same value).
  """
  @type index :: (value -> key)
  @type index_name :: atom

  @type key :: any
  @type value :: any

  @type full_key :: primary_key | non_primary_key
  @type primary_key :: key
  @opaque non_primary_key :: {Key, index_name, key}

  @doc """
  Creates a new Idx instance.
  """
  @spec new(Enumerable.t(), index) :: t
  def new(enum \\ [], primary_index) do
    map = Map.new(enum, &{{Primary, primary_index.(&1)}, &1})
    %__MODULE__{Primary => primary_index, indices: %{}, lazy_indices: %{}} |> Map.merge(map)
  end

  @doc """
  Allows accessing a value by a non-primary key in the idx.

  See `create_index/3` for details.
  """
  @spec key(index_name, non_primary_key) :: full_key
  def key(index_name, key) do
    {Key, index_name, key}
  end

  @doc """
  Creates a new non-primary index on the `Idx` instance.

      iex> users = [%{name: "Bob", age: 20}, %{name: "Eve", age: 27}, %{name: "John", age: 45}]
      iex> idx = Idx.new(users, & &1.name)
      iex> idx = Idx.create_index(idx, :initial, &String.first(&1.name))
      iex> Idx.get(idx, Idx.key(:initial, "J"))
      %{name: "John", age: 45}

  If `lazy` option is set to true, the index keys won't be precomputed and stored,
  instead it will always be calculated on demand. This will slow down access, but
  speed up insertion and operations relying on other indices.
  """
  @spec create_index(t, index_name, index, lazy?: boolean()) :: t
  def create_index(%__MODULE__{} = idx, name, fun, options \\ []) do
    %__MODULE__{indices: indices, lazy_indices: lazy_indices} = idx

    if Map.has_key?(indices, name) or Map.has_key?(lazy_indices, name) do
      raise ArgumentError, "Index #{inspect(name)} already present"
    end

    if Keyword.get(options, :lazy?, false) do
      %{idx | lazy_indices: Map.put(lazy_indices, name, fun)}
    else
      data = idx |> to_map() |> Map.new(fn {key, value} -> {fun.(value), key} end)

      %{idx | indices: Map.put(indices, name, fun)}
      |> Map.put({__MODULE__, name}, data)
    end
  end

  @spec drop_index(t, index_name) :: t
  def drop_index(%__MODULE__{} = idx, name) do
    %__MODULE__{indices: indices, lazy_indices: lazy_indices} = idx
    {eager_fun, indices} = Map.pop(indices, name)
    {lazy_fun, lazy_indices} = Map.pop(lazy_indices, name)

    unless eager_fun || lazy_fun do
      raise ArgumentError, "Unknown index #{inspect(name)}"
    end

    idx = Map.delete(idx, {__MODULE__, name})
    %{idx | indices: indices, lazy_indices: lazy_indices}
  end

  @spec put(t, value) :: t
  def put(%__MODULE__{Primary => primary_index, indices: indices} = idx, value)
      when indices == %{} do
    Map.put(idx, {Primary, primary_index.(value)}, value)
  end

  def put(%__MODULE__{Primary => primary_index, indices: indices} = idx, value) do
    primary_key = primary_index.(value)
    idx = Map.put(idx, {Primary, primary_index.(value)}, value)

    Enum.reduce(indices, idx, fn {name, fun}, idx ->
      name = {__MODULE__, name}
      %{^name => data} = idx
      %{idx | name => Map.put(data, fun.(value), primary_key)}
    end)
  end

  @spec fetch(t, full_key) :: {:ok, value} | :error
  def fetch(%__MODULE__{} = idx, key) do
    key = resolve_key(idx, key)

    case idx do
      %{^key => value} -> {:ok, value}
      %{} -> :error
    end
  end

  @spec fetch!(t, full_key) :: value
  def fetch!(idx, key) do
    key = resolve_key!(idx, key)
    %{^key => value} = idx
    value
  end

  @spec get(t, full_key, value | nil) :: value | nil
  def get(idx, key, default \\ nil) do
    case fetch(idx, key) do
      {:ok, value} -> value
      :error -> default
    end
  end

  @spec pop!(t, full_key) :: value
  def pop!(idx, key) do
    key = resolve_key!(idx, key)
    {value, idx} = Map.pop!(idx, key)
    {value, remove_value_from_indices(idx, value)}
  end

  def pop(idx, key, default \\ nil) do
    key = resolve_key(idx, key)

    case :maps.take(key, idx) do
      {value, idx} -> {value, remove_value_from_indices(idx, value)}
      :error -> {default, idx}
    end
  end

  defp remove_value_from_indices(%Idx{indices: indices} = idx, value) do
    Enum.reduce(indices, idx, fn {name, fun}, idx ->
      name = {__MODULE__, name}
      %{^name => data} = idx
      key = fun.(value)
      data = Map.delete(data, key)
      %{idx | name => data}
    end)
  end

  @spec update!(t, full_key, (value -> value)) :: t
  def update!(idx, key, fun) do
    key = resolve_key!(idx, key)
    {value, idx} = pop!(idx, key)
    put(idx, fun.(value))
  end

  @doc """
  Updates the value under the key by calling `fun`. No indices can change after the update.

  Thanks to the guarantee that the indices (including the primary index)
  won't change, this function is faster than `update!/3`. However, if the
  indices change, it will leave the `idx` in an invalid state and lead
  to undefined behaviour.
  """
  @spec fast_update!(t, full_key, (value -> value)) :: t
  def fast_update!(idx, key, fun) do
    key = resolve_key!(idx, key)
    Map.update!(idx, key, fun)
  end

  @spec get_and_update!(t, full_key, (value -> {get, update} | :pop)) :: {get, t}
        when get: any, update: value
  def get_and_update!(idx, key, fun) do
    key = resolve_key!(idx, key)
    {value, idx} = pop!(idx, key)

    case fun.(value) do
      {get, update} -> {get, put(idx, update)}
      :pop -> {value, idx}
    end
  end

  @spec get_and_update(t, full_key, (value | nil -> {get, update} | :pop)) :: {get, t}
        when get: any, update: value
  def get_and_update(idx, key, fun) do
    key = resolve_key(idx, key)
    {value, idx} = pop(idx, key)

    case fun.(value) do
      {get, update} -> {get, put(idx, update)}
      :pop -> {value, idx}
    end
  end

  @spec size(t) :: non_neg_integer
  def size(%Idx{indices: indices} = idx) do
    map_size(idx) - map_size(Idx.__struct__()) - map_size(indices)
  end

  @spec to_list(t) :: [value()]
  def to_list(%Idx{} = idx) do
    :maps.fold(
      fn
        {Primary, _key}, value, acc -> [value | acc]
        _key, _value, acc -> acc
      end,
      [],
      idx
    )
    |> :lists.reverse()
  end

  @doc """
  Converts `idx` to a map, where keys are the primary keys of the `idx`.
  """
  @spec to_map(t) :: %{key => value}
  def to_map(%Idx{} = idx) do
    Enum.reduce(Map.from_struct(idx), %{}, fn
      {{Primary, key}, value}, acc -> Map.put(acc, key, value)
      _kv, acc -> acc
    end)
  end

  @spec member?(t, value) :: boolean
  def member?(%Idx{Idx.Primary => primary} = idx, value) do
    {:ok, value} == Idx.fetch(idx, primary.(value))
  end

  @spec primary_key!(t, index_name, non_primary_key) :: primary_key
  def primary_key!(%{__struct__: Idx} = idx, name, key) do
    name = {__MODULE__, name}
    %{^name => %{^key => primary}} = idx
    primary
  end

  @doc """
  Returns the primary key of a value under the given non-primary key.
  """
  @spec primary_key(t, index_name, non_primary_key) :: {:ok, primary_key} | :error
  def primary_key(%Idx{} = idx, name, key) do
    name = {__MODULE__, name}

    case idx do
      %{^name => %{^key => primary}} -> {:ok, primary}
      %{} -> :error
    end
  end

  defp resolve_key(_idx, {Primary, primary}) do
    {Primary, primary}
  end

  defp resolve_key(idx, {Key, name, key}) do
    data_ref = {__MODULE__, name}

    case idx do
      %{^data_ref => %{^key => primary}} -> {Primary, primary}
      %{lazy_indices: %{^name => fun}} -> resolve_key_lazy(idx, fun, key)
      %{} -> Key.Imaginary
    end
  end

  defp resolve_key(_idx, primary) do
    {Primary, primary}
  end

  defp resolve_key!(_idx, {Primary, primary}) do
    {Primary, primary}
  end

  defp resolve_key!(idx, {Key, name, key}) do
    case resolve_key(idx, {Key, name, key}) do
      Key.Imaginary -> raise "Unknown key #{inspect(key)} of index #{inspect(name)}"
      primary_key -> primary_key
    end
  end

  defp resolve_key!(_idx, primary) do
    {Primary, primary}
  end

  defp resolve_key_lazy(idx, fun, key) do
    Enum.find_value(Map.from_struct(idx), Key.Imaginary, fn
      {{Primary, primary_key}, value} -> if fun.(value) === key, do: {Primary, primary_key}
      _entry -> nil
    end)
  end
end

defimpl Enumerable, for: Idx do
  def count(idx) do
    {:ok, Idx.size(idx)}
  end

  def member?(idx, value) do
    {:ok, Idx.member?(idx, value)}
  end

  def slice(idx) do
    size = Idx.size(idx)
    {:ok, size, &Idx.to_list/1}
  end

  def reduce(idx, acc, fun) do
    Enumerable.List.reduce(Idx.to_list(idx), acc, fun)
  end
end

defimpl Collectable, for: Idx do
  def into(idx) do
    fun = fn
      idx, {:cont, value} -> Idx.put(idx, value)
      idx, :done -> idx
      _idx, :halt -> :ok
    end

    {idx, fun}
  end
end

defimpl Inspect, for: Idx do
  import Inspect.Algebra

  def inspect(%Idx{Idx.Primary => primary} = idx, opts) do
    opts = %Inspect.Opts{opts | charlists: :as_lists}

    concat([
      "#Idx<",
      to_doc(Idx.to_list(idx), opts),
      ", ",
      "indices: ",
      to_doc(
        [primary: primary] ++ Enum.to_list(idx.indices) ++ Enum.to_list(idx.lazy_indices),
        opts
      ),
      ">"
    ])
  end
end