defmodule MST.Height do
  @moduledoc """
  Key-depth computation for the AT Protocol Merkle Search Tree.

  Each key's depth (also called "layer" or "height") is derived by SHA-256
  hashing the key and counting the number of leading zero bits, divided by two
  (rounding down). This gives a theoretical fanout of 4: each additional level
  of depth is four times rarer than the previous.

  Spec: https://atproto.com/specs/repository#mst-structure
  """

  @doc """
  Computes the MST depth for a given key.

  SHA-256 hashes `key` and counts the number of leading zero bits in the
  binary output, then divides by 2 (floor). Returns a non-negative integer;
  depth 0 is the most common (probability ~75% per key), each higher depth
  is four times rarer.

  ## Examples

      iex> MST.Height.for_key("2653ae71")
      0

      iex> MST.Height.for_key("blue")
      1

      iex> MST.Height.for_key("app.bsky.feed.post/454397e440ec")
      4

      iex> MST.Height.for_key("app.bsky.feed.post/9adeb165882c")
      8

  """
  @spec for_key(binary()) :: non_neg_integer()
  def for_key(key) when is_binary(key) do
    :crypto.hash(:sha256, key)
    |> leading_zero_bits()
    |> div(2)
  end

  # ---------------------------------------------------------------------------
  # Private helpers
  # ---------------------------------------------------------------------------

  @spec leading_zero_bits(binary()) :: non_neg_integer()
  defp leading_zero_bits(<<>>), do: 0

  defp leading_zero_bits(<<byte, rest::binary>>) do
    lz = leading_zeros_in_byte(byte)

    if lz == 8 do
      8 + leading_zero_bits(rest)
    else
      lz
    end
  end

  # Returns the count of leading zero bits in a single byte (0–8).
  @spec leading_zeros_in_byte(byte()) :: 0..8
  defp leading_zeros_in_byte(0), do: 8
  defp leading_zeros_in_byte(b) when b >= 128, do: 0
  defp leading_zeros_in_byte(b) when b >= 64, do: 1
  defp leading_zeros_in_byte(b) when b >= 32, do: 2
  defp leading_zeros_in_byte(b) when b >= 16, do: 3
  defp leading_zeros_in_byte(b) when b >= 8, do: 4
  defp leading_zeros_in_byte(b) when b >= 4, do: 5
  defp leading_zeros_in_byte(b) when b >= 2, do: 6
  defp leading_zeros_in_byte(_), do: 7
end