defmodule Arcterex.Data.SplayTree do
  @moduledoc """
  A purely functional splay tree data structure.

  A splay tree is a self-adjusting binary search tree where recently accessed
  elements are moved to the root through a splaying operation. This provides
  amortized O(log n) performance for operations while optimizing access to
  frequently-used nodes.

  ## When to Use

  Splay trees are well-suited for:
  - Applications with non-uniform access patterns (locality of reference)
  - Caching and LRU-like behavior without explicit cache management
  - Ordered maps where recent keys are accessed more frequently
  - Spatial data structures with locality properties

  ## Key Characteristics

  - **Self-adjusting**: Automatically reorganizes based on access patterns
  - **Amortized O(log n)**: Individual operations may be O(n), but amortized over sequences
  - **Simple implementation**: No complex balancing rules like red-black trees
  - **Cache-friendly**: Recent items stay near the root

  ## Examples

      # Create an empty tree
      iex> tree = Arcterex.Data.SplayTree.new()
      iex> Arcterex.Data.SplayTree.empty?(tree)
      true

      # Insert key-value pairs
      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> Arcterex.Data.SplayTree.size(tree)
      3

      # Retrieve values
      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :key, "value")
      iex> Arcterex.Data.SplayTree.get(tree, :key, :default)
      "value"
      iex> Arcterex.Data.SplayTree.get(tree, :missing, :default)
      :default

      # Check for keys
      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> Arcterex.Data.SplayTree.has_key?(tree, :a)
      true
      iex> Arcterex.Data.SplayTree.has_key?(tree, :b)
      false

      # Get keys in sorted order
      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.keys(tree)
      [:a, :b, :c]

      # Custom comparator for reverse order
      iex> comparator = fn a, b ->
      ...>   cond do
      ...>     a > b -> :lt
      ...>     a < b -> :gt
      ...>     true -> :eq
      ...>   end
      ...> end
      iex> tree = Arcterex.Data.SplayTree.new(comparator)
      iex> tree = Arcterex.Data.SplayTree.put(tree, 1, "one")
      iex> tree = Arcterex.Data.SplayTree.put(tree, 2, "two")
      iex> tree = Arcterex.Data.SplayTree.put(tree, 3, "three")
      iex> Arcterex.Data.SplayTree.keys(tree)
      [3, 2, 1]
  """

  @opaque t() :: t(any(), any())
  @opaque t(key, value) :: %__MODULE__{
            root: node(key, value) | nil,
            size: non_neg_integer(),
            comparator: comparator()
          }

  @type comparator() :: (any(), any() -> comparison_result())
  @type comparison_result() :: :lt | :eq | :gt

  @typep node(key, value) :: {key, value, node(key, value) | nil, node(key, value) | nil}

  defstruct root: nil, size: 0, comparator: &__MODULE__.default_comparator/2

  @doc """
  Creates a new empty splay tree with the default comparator.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> Arcterex.Data.SplayTree.empty?(tree)
      true
      iex> Arcterex.Data.SplayTree.size(tree)
      0
  """
  @spec new() :: t()
  def new do
    %__MODULE__{}
  end

  @doc """
  Creates a new splay tree with a custom comparator or from a list.

  When given a comparator function, it should take two keys and return:
  - `:lt` if the first key is less than the second
  - `:eq` if the keys are equal
  - `:gt` if the first key is greater than the second

  When given a list of `{key, value}` tuples, creates a populated tree.

  ## Examples

      iex> reverse = fn a, b ->
      ...>   cond do
      ...>     a > b -> :lt
      ...>     a < b -> :gt
      ...>     true -> :eq
      ...>   end
      ...> end
      iex> tree = Arcterex.Data.SplayTree.new(reverse)
      iex> tree = Arcterex.Data.SplayTree.put(tree, 1, "a")
      iex> tree = Arcterex.Data.SplayTree.put(tree, 2, "b")
      iex> Arcterex.Data.SplayTree.keys(tree)
      [2, 1]

      iex> tree = Arcterex.Data.SplayTree.new([{:b, 2}, {:a, 1}, {:c, 3}])
      iex> Arcterex.Data.SplayTree.size(tree)
      3
  """
  @spec new(comparator()) :: t()
  def new(comparator) when is_function(comparator, 2) do
    %__MODULE__{comparator: comparator}
  end

  @spec new(list({k, v})) :: t(k, v) when k: any(), v: any()
  def new(list) when is_list(list) do
    Enum.reduce(list, %__MODULE__{}, fn {k, v}, tree -> put(tree, k, v) end)
  end

  @doc """
  Creates a new splay tree from a list with a custom comparator.

  ## Examples

      iex> desc = fn a, b ->
      ...>   cond do
      ...>     a > b -> :lt
      ...>     a < b -> :gt
      ...>     true -> :eq
      ...>   end
      ...> end
      iex> tree = Arcterex.Data.SplayTree.new([{1, :a}, {2, :b}], desc)
      iex> Arcterex.Data.SplayTree.keys(tree)
      [2, 1]
  """
  @spec new(list({k, v}), comparator()) :: t(k, v) when k: any(), v: any()
  def new(list, comparator) when is_list(list) and is_function(comparator, 2) do
    Enum.reduce(list, %__MODULE__{comparator: comparator}, fn {k, v}, tree -> put(tree, k, v) end)
  end

  @doc """
  Default comparison function for keys.

  Returns `:lt`, `:eq`, or `:gt` based on Elixir's term ordering.

  ## Examples

      iex> Arcterex.Data.SplayTree.default_comparator(1, 2)
      :lt
      iex> Arcterex.Data.SplayTree.default_comparator(2, 2)
      :eq
      iex> Arcterex.Data.SplayTree.default_comparator(3, 2)
      :gt
  """
  @spec default_comparator(any(), any()) :: comparison_result()
  def default_comparator(a, b) do
    cond do
      a < b -> :lt
      a > b -> :gt
      true -> :eq
    end
  end

  @doc """
  Inserts a key-value pair into the tree.

  Returns `{:ok, tree}` if the key was successfully inserted, or
  `{:error, :key_exists}` if the key already exists.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> {:ok, tree} = Arcterex.Data.SplayTree.insert(tree, :a, 1)
      iex> {:ok, tree} = Arcterex.Data.SplayTree.insert(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.size(tree)
      2

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> {:ok, tree} = Arcterex.Data.SplayTree.insert(tree, :a, 1)
      iex> Arcterex.Data.SplayTree.insert(tree, :a, 999)
      {:error, :key_exists}
  """
  @spec insert(t(k, v), k, v) :: {:ok, t(k, v)} | {:error, :key_exists} when k: any(), v: any()
  def insert(%__MODULE__{root: nil, size: 0, comparator: _comp} = tree, key, value) do
    {:ok, %{tree | root: {key, value, nil, nil}, size: 1}}
  end

  def insert(%__MODULE__{root: root, size: size, comparator: comp} = tree, key, value) do
    new_root = splay(root, key, comp)

    case new_root do
      {k, _, _, _} ->
        case comp.(key, k) do
          :eq ->
            # Key exists, don't insert
            {:error, :key_exists}

          :lt ->
            # Insert as new root with current root as right child
            {_, _, left, _} = new_root
            new_node = {key, value, left, new_root}

            {:ok,
             %{tree | root: put_elem(new_node, 3, put_elem(new_root, 2, nil)), size: size + 1}}

          :gt ->
            # Insert as new root with current root as left child
            {_, _, _, right} = new_root
            new_node = {key, value, new_root, right}

            {:ok,
             %{tree | root: put_elem(new_node, 2, put_elem(new_root, 3, nil)), size: size + 1}}
        end

      nil ->
        {:ok, %{tree | root: {key, value, nil, nil}, size: 1}}
    end
  end

  @doc """
  Inserts or updates a key-value pair in the tree.

  If the key exists, its value is updated. If not, a new entry is inserted.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      ...> |> Arcterex.Data.SplayTree.put(:a, 1)
      ...> |> Arcterex.Data.SplayTree.put(:a, 2)
      iex> Arcterex.Data.SplayTree.get(tree, :a, nil)
      2
      iex> Arcterex.Data.SplayTree.size(tree)
      1
  """
  @spec put(t(k, v), k, v) :: t(k, v) when k: any(), v: any()
  def put(%__MODULE__{root: nil, size: 0, comparator: _comp} = tree, key, value) do
    %{tree | root: {key, value, nil, nil}, size: 1}
  end

  def put(%__MODULE__{root: root, size: size, comparator: comp} = tree, key, value) do
    new_root = splay(root, key, comp)

    case new_root do
      {k, _, _, _} ->
        case comp.(key, k) do
          :eq ->
            # Update existing key
            %{tree | root: put_elem(new_root, 1, value)}

          :lt ->
            # Insert as new root with current root as right child
            {_, _, left, _} = new_root
            new_node = {key, value, left, new_root}
            %{tree | root: put_elem(new_node, 3, put_elem(new_root, 2, nil)), size: size + 1}

          :gt ->
            # Insert as new root with current root as left child
            {_, _, _, right} = new_root
            new_node = {key, value, new_root, right}
            %{tree | root: put_elem(new_node, 2, put_elem(new_root, 3, nil)), size: size + 1}
        end

      nil ->
        %{tree | root: {key, value, nil, nil}, size: 1}
    end
  end

  @doc """
  Retrieves the value associated with a key, or returns the default if not found.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new([{:a, 1}])
      iex> Arcterex.Data.SplayTree.get(tree, :a, :not_found)
      1
      iex> Arcterex.Data.SplayTree.get(tree, :a)
      1
      iex> Arcterex.Data.SplayTree.get(tree, :b, :not_found)
      :not_found
      iex> Arcterex.Data.SplayTree.get(tree, :b)
      nil
  """
  @spec get(t(k, v), k, default) :: v | default when k: any(), v: any(), default: any()
  def get(tree, key, default \\ nil)

  def get(%__MODULE__{root: nil}, _key, default), do: default

  def get(%__MODULE__{root: root, comparator: comp}, key, default) do
    case find_node(root, key, comp) do
      {^key, value, _, _} -> value
      _ -> default
    end
  end

  @doc """
  Accesses a key and splays the tree, returning both the value and rebalanced tree.

  This function performs splaying on access, which moves the accessed key to the root
  for faster subsequent access. Returns `{value, new_tree}` where value is `nil` if
  the key is not found.

  This is useful when you want to benefit from splay tree rebalancing for workloads
  with locality of reference.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new([{:a, 1}, {:b, 2}, {:c, 3}])
      iex> {value, tree2} = Arcterex.Data.SplayTree.access(tree, :b)
      iex> value
      2
      iex> {value2, _tree3} = Arcterex.Data.SplayTree.access(tree2, :b)
      iex> value2
      2
      iex> {value3, _tree4} = Arcterex.Data.SplayTree.access(tree, :z)
      iex> value3
      nil
  """
  @spec access(t(k, v), k) :: {v | nil, t(k, v)} when k: any(), v: any()
  def access(%__MODULE__{root: nil} = tree, _key), do: {nil, tree}

  def access(%__MODULE__{root: root, comparator: comp} = tree, key) do
    new_root = splay(root, key, comp)

    case new_root do
      {^key, value, _, _} ->
        {value, %{tree | root: new_root}}

      _ ->
        {nil, %{tree | root: new_root}}
    end
  end

  @doc """
  Fetches the value for a key, returning `{:ok, value}` or `:error`.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> Arcterex.Data.SplayTree.fetch(tree, :a)
      {:ok, 1}
      iex> Arcterex.Data.SplayTree.fetch(tree, :b)
      :error
  """
  @spec fetch(t(k, v), k) :: {:ok, v} | :error when k: any(), v: any()
  def fetch(%__MODULE__{root: nil}, _key), do: :error

  def fetch(%__MODULE__{root: root, comparator: comp}, key) do
    case find_node(root, key, comp) do
      {^key, value, _, _} -> {:ok, value}
      _ -> :error
    end
  end

  @doc """
  Deletes a key from the tree.

  Returns the tree unchanged if the key doesn't exist.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> tree = Arcterex.Data.SplayTree.delete(tree, :a)
      iex> Arcterex.Data.SplayTree.has_key?(tree, :a)
      false
      iex> Arcterex.Data.SplayTree.size(tree)
      1
  """
  @spec delete(t(k, v), k) :: t(k, v) when k: any(), v: any()
  def delete(%__MODULE__{root: nil} = tree, _key), do: tree

  def delete(%__MODULE__{root: root, size: size, comparator: comp} = tree, key) do
    new_root = splay(root, key, comp)

    case new_root do
      {k, _, left, right} ->
        case comp.(key, k) do
          :eq -> do_delete(tree, left, right, size)
          _ -> %{tree | root: new_root}
        end

      nil ->
        tree
    end
  end

  # Helper function to handle deletion when key is found
  defp do_delete(tree, nil, right, size) do
    %{tree | root: right, size: size - 1}
  end

  defp do_delete(tree, left, right, size) do
    # Splay the maximum element in left subtree to root
    max_node = splay_max(left, tree.comparator)
    # Attach right subtree to the right of max node
    new_root = put_elem(max_node, 3, right)
    %{tree | root: new_root, size: size - 1}
  end

  @doc """
  Checks if a key exists in the tree.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> Arcterex.Data.SplayTree.has_key?(tree, :a)
      true
      iex> Arcterex.Data.SplayTree.has_key?(tree, :b)
      false
  """
  @spec has_key?(t(k, v), k) :: boolean() when k: any(), v: any()
  def has_key?(%__MODULE__{root: nil}, _key), do: false

  def has_key?(%__MODULE__{root: root, comparator: comp}, key) do
    case find_node(root, key, comp) do
      {^key, _, _, _} -> true
      _ -> false
    end
  end

  @doc """
  Returns true if the tree is empty.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> Arcterex.Data.SplayTree.empty?(tree)
      true
      iex> {:ok, tree} = Arcterex.Data.SplayTree.insert(tree, :a, 1)
      iex> Arcterex.Data.SplayTree.empty?(tree)
      false
  """
  @spec empty?(t()) :: boolean()
  def empty?(%__MODULE__{size: 0}), do: true
  def empty?(%__MODULE__{}), do: false

  @doc """
  Returns the number of elements in the tree.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> Arcterex.Data.SplayTree.size(tree)
      0
      iex> {:ok, tree} = Arcterex.Data.SplayTree.insert(tree, :a, 1)
      iex> Arcterex.Data.SplayTree.size(tree)
      1
  """
  @spec size(t()) :: non_neg_integer()
  def size(%__MODULE__{size: size}), do: size

  @doc """
  Returns all keys in the tree in sorted order.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.keys(tree)
      [:a, :b, :c]
  """
  @spec keys(t(k, v)) :: [k] when k: any(), v: any()
  def keys(%__MODULE__{root: root}) do
    inorder_keys(root, [])
  end

  @doc """
  Returns all values in the tree in key-sorted order.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.values(tree)
      [1, 2, 3]
  """
  @spec values(t(k, v)) :: [v] when k: any(), v: any()
  def values(%__MODULE__{root: root}) do
    inorder_values(root, [])
  end

  @doc """
  Returns all key-value pairs in the tree in sorted order.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.to_list(tree)
      [a: 1, b: 2, c: 3]
  """
  @spec to_list(t(k, v)) :: [{k, v}] when k: any(), v: any()
  def to_list(%__MODULE__{root: root}) do
    inorder_pairs(root, [])
  end

  @doc """
  Inserts multiple key-value pairs into the tree.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      ...> |> Arcterex.Data.SplayTree.put_many([{:a, 1}, {:b, 2}, {:c, 3}])
      iex> Arcterex.Data.SplayTree.size(tree)
      3
      iex> Arcterex.Data.SplayTree.keys(tree)
      [:a, :b, :c]
  """
  @spec put_many(t(k, v), [{k, v}]) :: t(k, v) when k: any(), v: any()
  def put_many(%__MODULE__{} = tree, pairs) when is_list(pairs) do
    Enum.reduce(pairs, tree, fn {key, value}, acc ->
      put(acc, key, value)
    end)
  end

  @doc """
  Clears all elements from the tree.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> tree = Arcterex.Data.SplayTree.clear(tree)
      iex> Arcterex.Data.SplayTree.empty?(tree)
      true
      iex> Arcterex.Data.SplayTree.size(tree)
      0
  """
  @spec clear(t()) :: t()
  def clear(%__MODULE__{comparator: comp}) do
    %__MODULE__{comparator: comp}
  end

  @doc """
  Returns the height of the tree.

  The height is the length of the longest path from root to leaf.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> Arcterex.Data.SplayTree.height(tree)
      0
      iex> {:ok, tree} = Arcterex.Data.SplayTree.insert(tree, :a, 1)
      iex> Arcterex.Data.SplayTree.height(tree)
      1
      iex> {:ok, tree} = Arcterex.Data.SplayTree.insert(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.height(tree) >= 1
      true
  """
  @spec height(t()) :: non_neg_integer()
  def height(%__MODULE__{root: nil}), do: 0
  def height(%__MODULE__{root: root}), do: compute_height(root)

  @doc """
  Returns the key-value pair at the specified index in sorted order.

  Returns `{:ok, {key, value}}` if the index is valid, `:error` otherwise.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.at(tree, 0)
      {:ok, {:a, 1}}
      iex> Arcterex.Data.SplayTree.at(tree, 1)
      {:ok, {:b, 2}}
      iex> Arcterex.Data.SplayTree.at(tree, 10)
      :error
      iex> Arcterex.Data.SplayTree.at(tree, -1)
      :error
  """
  @spec at(t(k, v), integer()) :: {:ok, {k, v}} | :error when k: any(), v: any()
  def at(%__MODULE__{} = tree, index) when is_integer(index) and index >= 0 do
    case Enum.at(to_list(tree), index) do
      nil -> :error
      pair -> {:ok, pair}
    end
  end

  def at(%__MODULE__{}, _index), do: :error

  @doc """
  Returns the minimum key in the tree.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> Arcterex.Data.SplayTree.min_key(tree)
      nil
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.min_key(tree)
      :a
  """
  @spec min_key(t(k, v)) :: k | nil when k: any(), v: any()
  def min_key(%__MODULE__{root: nil}), do: nil

  def min_key(%__MODULE__{root: root, comparator: comp}) do
    {key, _, _, _} = splay_min(root, comp)
    key
  end

  @doc """
  Returns the maximum key in the tree.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> Arcterex.Data.SplayTree.max_key(tree)
      nil
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.max_key(tree)
      :c
  """
  @spec max_key(t(k, v)) :: k | nil when k: any(), v: any()
  def max_key(%__MODULE__{root: nil}), do: nil

  def max_key(%__MODULE__{root: root, comparator: comp}) do
    {key, _, _, _} = splay_max(root, comp)
    key
  end

  @doc """
  Returns the minimum entry (key-value pair) in the tree.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> Arcterex.Data.SplayTree.min_entry(tree)
      nil
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.min_entry(tree)
      {:a, 1}
  """
  @spec min_entry(t(k, v)) :: {k, v} | nil when k: any(), v: any()
  def min_entry(%__MODULE__{root: nil}), do: nil

  def min_entry(%__MODULE__{root: root, comparator: comp}) do
    {key, value, _, _} = splay_min(root, comp)
    {key, value}
  end

  @doc """
  Returns the maximum entry (key-value pair) in the tree.

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> Arcterex.Data.SplayTree.max_entry(tree)
      nil
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :b, 2)
      iex> Arcterex.Data.SplayTree.max_entry(tree)
      {:c, 3}
  """
  @spec max_entry(t(k, v)) :: {k, v} | nil when k: any(), v: any()
  def max_entry(%__MODULE__{root: nil}), do: nil

  def max_entry(%__MODULE__{root: root, comparator: comp}) do
    {key, value, _, _} = splay_max(root, comp)
    {key, value}
  end

  @doc """
  Returns the predecessor (largest key less than the given key).

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :e, 5)
      iex> Arcterex.Data.SplayTree.predecessor(tree, :e)
      :c
      iex> Arcterex.Data.SplayTree.predecessor(tree, :a)
      nil
  """
  @spec predecessor(t(k, v), k) :: k | nil when k: any(), v: any()
  def predecessor(%__MODULE__{root: nil}, _key), do: nil

  def predecessor(%__MODULE__{root: root, comparator: comp}, key) do
    splay_and_find_predecessor(root, key, comp, nil)
  end

  @doc """
  Returns the successor (smallest key greater than the given key).

  ## Examples

      iex> tree = Arcterex.Data.SplayTree.new()
      iex> tree = Arcterex.Data.SplayTree.put(tree, :c, 3)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :a, 1)
      iex> tree = Arcterex.Data.SplayTree.put(tree, :e, 5)
      iex> Arcterex.Data.SplayTree.successor(tree, :a)
      :c
      iex> Arcterex.Data.SplayTree.successor(tree, :e)
      nil
  """
  @spec successor(t(k, v), k) :: k | nil when k: any(), v: any()
  def successor(%__MODULE__{root: nil}, _key), do: nil

  def successor(%__MODULE__{root: root, comparator: comp}, key) do
    splay_and_find_successor(root, key, comp, nil)
  end

  # Private functions

  @spec splay(node(k, v) | nil, k, comparator()) :: node(k, v) | nil when k: any(), v: any()
  defp splay(nil, _key, _comp), do: nil

  defp splay(node, key, comp) do
    do_splay(node, key, comp)
  end

  @spec do_splay(node(k, v) | nil, k, comparator()) :: node(k, v) | nil when k: any(), v: any()
  defp do_splay(nil, _key, _comp), do: nil

  defp do_splay({k, v, left, right} = node, key, comp) do
    case comp.(key, k) do
      :eq -> node
      :lt -> splay_left(node, k, v, left, right, key, comp)
      :gt -> splay_right(node, k, v, left, right, key, comp)
    end
  end

  # Splay left: handles zig, zig-zig (left-left), and zig-zag (left-right)
  defp splay_left(node, _k, _v, nil, _right, _key, _comp), do: node

  defp splay_left(_node, k, v, {lk, lv, ll, lr}, right, key, comp) do
    case comp.(key, lk) do
      :eq -> rotate_right({k, v, right}, {lk, lv, ll, lr})
      :lt -> splay_left_left({k, v, right}, {lk, lv, ll, lr}, key, comp)
      :gt -> splay_left_right({k, v, right}, {lk, lv, ll, lr}, key, comp)
    end
  end

  # Zig-zig (left-left): double rotation right
  defp splay_left_left({k, v, right}, {lk, lv, nil, lr}, _key, _comp) do
    # Key not in tree, rotate right once
    {lk, lv, nil, {k, v, lr, right}}
  end

  defp splay_left_left({k, v, right}, {lk, lv, ll, lr}, key, comp) do
    # Recursively splay in left-left subtree, then double rotate
    {llk, llv, lll, llr} = do_splay(ll, key, comp)
    {llk, llv, lll, {lk, lv, llr, {k, v, lr, right}}}
  end

  # Zig-zag (left-right): rotate left then right
  defp splay_left_right({k, v, right}, {lk, lv, ll, nil}, _key, _comp) do
    # Key not in tree
    {lk, lv, ll, {k, v, nil, right}}
  end

  defp splay_left_right({k, v, right}, {lk, lv, ll, lr}, key, comp) do
    # Recursively splay in left-right subtree, then rotate
    {lrk, lrv, lrl, lrr} = do_splay(lr, key, comp)
    {lrk, lrv, {lk, lv, ll, lrl}, {k, v, lrr, right}}
  end

  # Splay right: handles zig, zig-zig (right-right), and zig-zag (right-left)
  defp splay_right(node, _k, _v, _left, nil, _key, _comp), do: node

  defp splay_right(_node, k, v, left, {rk, rv, rl, rr}, key, comp) do
    case comp.(key, rk) do
      :eq -> rotate_left({k, v, left}, {rk, rv, rl, rr})
      :gt -> splay_right_right({k, v, left}, {rk, rv, rl, rr}, key, comp)
      :lt -> splay_right_left({k, v, left}, {rk, rv, rl, rr}, key, comp)
    end
  end

  # Zig-zig (right-right): double rotation left
  defp splay_right_right({k, v, left}, {rk, rv, rl, nil}, _key, _comp) do
    # Key not in tree, rotate left once
    {rk, rv, {k, v, left, rl}, nil}
  end

  defp splay_right_right({k, v, left}, {rk, rv, rl, rr}, key, comp) do
    # Recursively splay in right-right subtree, then double rotate
    {rrk, rrv, rrl, rrr} = do_splay(rr, key, comp)
    {rrk, rrv, {rk, rv, {k, v, left, rl}, rrl}, rrr}
  end

  # Zig-zag (right-left): rotate right then left
  defp splay_right_left({k, v, left}, {rk, rv, nil, rr}, _key, _comp) do
    # Key not in tree
    {rk, rv, {k, v, left, nil}, rr}
  end

  defp splay_right_left({k, v, left}, {rk, rv, rl, rr}, key, comp) do
    # Recursively splay in right-left subtree, then rotate
    {rlk, rlv, rll, rlr} = do_splay(rl, key, comp)
    {rlk, rlv, {k, v, left, rll}, {rk, rv, rlr, rr}}
  end

  # Helper function for single right rotation
  defp rotate_right({k, v, right}, {lk, lv, ll, lr}) do
    {lk, lv, ll, {k, v, lr, right}}
  end

  # Helper function for single left rotation
  defp rotate_left({k, v, left}, {rk, rv, rl, rr}) do
    {rk, rv, {k, v, left, rl}, rr}
  end

  @spec splay_min(node(k, v), comparator()) :: node(k, v) when k: any(), v: any()
  defp splay_min({k, v, nil, right}, _comp), do: {k, v, nil, right}

  defp splay_min({k, v, left, right}, comp) do
    {lk, lv, ll, lr} = left

    case ll do
      nil ->
        # Left child is minimum, rotate right
        {lk, lv, nil, {k, v, lr, right}}

      _ ->
        # Continue to find minimum in left subtree
        min_node = splay_min(ll, comp)
        {mk, mv, ml, mr} = min_node
        # Rotate right
        {mk, mv, ml, {lk, lv, mr, {k, v, lr, right}}}
    end
  end

  @spec splay_max(node(k, v), comparator()) :: node(k, v) when k: any(), v: any()
  defp splay_max({k, v, left, nil}, _comp), do: {k, v, left, nil}

  defp splay_max({k, v, left, right}, comp) do
    {rk, rv, rl, rr} = right

    case rr do
      nil ->
        # Right child is maximum, rotate left
        {rk, rv, {k, v, left, rl}, nil}

      _ ->
        # Continue to find maximum in right subtree
        max_node = splay_max(rr, comp)
        {mk, mv, ml, mr} = max_node
        # Rotate left
        {mk, mv, {rk, rv, {k, v, left, rl}, ml}, mr}
    end
  end

  @spec find_node(node(k, v) | nil, k, comparator()) :: node(k, v) | nil
        when k: any(), v: any()
  defp find_node(nil, _key, _comp), do: nil

  defp find_node({k, _v, left, right} = node, key, comp) do
    case comp.(key, k) do
      :eq -> node
      :lt -> find_node(left, key, comp)
      :gt -> find_node(right, key, comp)
    end
  end

  defp inorder_keys(nil, acc), do: acc

  defp inorder_keys({k, _v, left, right}, acc) do
    acc = inorder_keys(right, acc)
    acc = [k | acc]
    inorder_keys(left, acc)
  end

  defp inorder_values(nil, acc), do: acc

  defp inorder_values({_k, v, left, right}, acc) do
    acc = inorder_values(right, acc)
    acc = [v | acc]
    inorder_values(left, acc)
  end

  defp inorder_pairs(nil, acc), do: acc

  defp inorder_pairs({k, v, left, right}, acc) do
    acc = inorder_pairs(right, acc)
    acc = [{k, v} | acc]
    inorder_pairs(left, acc)
  end

  defp compute_height({_k, _v, left, right}) do
    left_height =
      case left do
        nil -> 0
        _ -> compute_height(left)
      end

    right_height =
      case right do
        nil -> 0
        _ -> compute_height(right)
      end

    1 + max(left_height, right_height)
  end

  @spec splay_and_find_predecessor(node(k, v) | nil, k, comparator(), k | nil) :: k | nil
        when k: any(), v: any()
  defp splay_and_find_predecessor(nil, _key, _comp, result), do: result

  defp splay_and_find_predecessor({node_key, _value, left, right}, key, comp, result) do
    case comp.(key, node_key) do
      :lt ->
        # Key is less than current node, search left
        splay_and_find_predecessor(left, key, comp, result)

      :gt ->
        # Key is greater than current node, current node could be predecessor
        # Search right subtree for a closer one
        splay_and_find_predecessor(right, key, comp, node_key)

      :eq ->
        # Found exact key, predecessor is max of left subtree
        case left do
          nil ->
            result

          _ ->
            {max_key, _, _, _} = splay_max(left, comp)
            max_key
        end
    end
  end

  @spec splay_and_find_successor(node(k, v) | nil, k, comparator(), k | nil) :: k | nil
        when k: any(), v: any()
  defp splay_and_find_successor(nil, _key, _comp, result), do: result

  defp splay_and_find_successor({node_key, _value, left, right}, key, comp, result) do
    case comp.(key, node_key) do
      :lt ->
        # Key is less than current node, current node could be successor
        # Search left subtree for a closer one
        splay_and_find_successor(left, key, comp, node_key)

      :gt ->
        # Key is greater than current node, search right
        splay_and_find_successor(right, key, comp, result)

      :eq ->
        # Found exact key, successor is min of right subtree
        case right do
          nil ->
            result

          _ ->
            {min_key, _, _, _} = splay_min(right, comp)
            min_key
        end
    end
  end
end