defmodule SData.MerkleSearchTree do
  @moduledoc"""
  A Merkle search tree.

  A node of the tree is
    {
      level,
      hash_of_node | nil,
      [ 
        { k_low_bound, v, hash_of_node | nil },
        { k_low_bound, v, hash_of_node | nil },
        ...
      }
    }
  """

  alias SData.PageStore, as: Store

  defstruct [:root, :store]

  defmodule Page do
    defstruct [:level, :low, :list]
  end

  defimpl SData.Page, for: Page do
    def refs(page) do
      refs = for {_, _, h} <- page.list, h != nil, do: h
      if page.low != nil do
        [ page.low | refs ]
      else
        refs
      end
    end
  end


  def new(store \\ nil) do
    store = store || SData.LocalStore.new
    %SData.MerkleSearchTree{
      root:  nil,
      store: store
    }
  end

  def insert(state, key, value) do
    level = calc_level(key)
    {hash, store} = insert_at(state.store, state.root, key, level, value)
    %{ state | root: hash, store: store }
  end

  def remove(_state, _key) do
    # TODO
  end

  defp insert_at(store, root, key, level, value) do
    {new_page, store} = if root == nil do
      { %Page{ level: level, low: nil, list: [ { key, value, nil } ] }, store }
    else
      %Page{ level: plevel, low: low, list: lst } = Store.get(store, root)
      [ { k0, _, _} | _ ] = lst
      cond do
        plevel < level ->
          {low, high, store} = split(store, root, key)
          { %Page{ level: level, low: low, list: [ { key, value, high } ] }, store }
        plevel == level ->
          store = Store.free(store, root)
          if key < k0 do
            {low2a, low2b, store} = split(store, low, key)
            { %Page{ level: level, low: low2a, list: [ { key, value, low2b } | lst] }, store }
          else
            {new_lst, store} = aux_insert_after_first(store, lst, key, value)
            { %Page{ level: plevel, low: low, list: new_lst }, store }
          end
        plevel > level ->
          store = Store.free(store, root)
          if key < k0 do
            {new_low, store} = insert_at(store, low, key, level, value)
            { %Page{ level: plevel, low: new_low, list: lst }, store }
          else 
            {new_lst, store} = aux_insert_sub_after_first(store, lst, key, level, value)
            { %Page{ level: plevel, low: low, list: new_lst }, store }
          end
      end
    end
    Store.put(store, new_page)    # returns {hash, store}
  end

  defp split(store, hash, key) do
    if hash == nil do
      {nil, nil, store}
    else
      %Page{ level: level, low: low, list: lst } = Store.get(store, hash)
      store = Store.free(store, hash)
      [ { k0, _, _} | _ ] = lst
      cond do
        key < k0 ->
          {lowlow, lowhi, store} = split(store, low, key)
          newp2 = %Page{ level: level, low: lowhi, list: lst }
          {newp2h, store} = Store.put(store, newp2)
          {lowlow, newp2h, store}
        key > k0 ->
          {lst1, p2, store} = split_aux(store, lst, key, level)
          newp1 = %Page{ level: level, low: low, list: lst1}
          {newp1h, store} = Store.put(store, newp1)
          {newp1h, p2, store}
      end 
    end
  end

  defp split_aux(store, lst, key, level) do
    case lst do
      [ {k1, _, _} | _ ] when k1 == key ->
        raise "Bad logic"
      [ {k1, v1, r1} ] ->
        {r1l, r1h, store} = split(store, r1, key)
        { [{k1, v1, r1l}], r1h, store }
      [ {k1, v1, r1}, {k2, v2, r2} | rst ] when key < k2->
        {r1l, r1h, store} = split(store, r1, key)
        p2 = %Page{ level: level, low: r1h, list: [ {k2, v2, r2} | rst ] }
        {p2h, store} = Store.put(store, p2)
        { [{k1, v1, r1l}], p2h, store }
      [ fst | rst ] ->
        {rst2, hi, store} = split_aux(store, rst, key, level)
        { [ fst | rst2 ], hi, store }
    end
  end

  defp aux_insert_after_first(store, lst, key, value) do
    case lst do
      [ {k1, _old_value, r1} | rst ] when k1 == key ->
        { [ {k1, value, r1} | rst ], store }
      [ {k1, v1, r1} ] ->
        {r1a, r1b, new_store} = split(store, r1, key)
        { [ {k1, v1, r1a}, {key, value, r1b} ], new_store }
      [ {k1, v1, r1}, {k2, v2, r2} | rst ] when key < k2->
        {r1a, r1b, new_store} = split(store, r1, key)
        { [ {k1, v1, r1a}, {key, value, r1b}, {k2, v2, r2} | rst ], new_store }
      [ fst | rst ] ->
        {rst2, new_store} = aux_insert_after_first(store, rst, key, value)
        { [ fst | rst2 ], new_store }
    end
  end

  defp aux_insert_sub_after_first(store, lst, key, level, value) do
    case lst do
      [ {k1, _, _} | _ ] when k1 == key ->
        raise "Bad logic"
      [ {k1, v1, r1} ] ->
        {r1new, store_new} = insert_at(store, r1, key, level, value)
        { [{k1, v1, r1new}], store_new }
      [ {k1, v1, r1}, {k2, v2, r2} | rst ] when key < k2->
        {r1new, store_new} = insert_at(store, r1, key, level, value)
        { [{k1, v1, r1new}, {k2, v2, r2} | rst], store_new }
      [ fst | rst ] ->
        {rst2, new_store} = aux_insert_sub_after_first(store, rst, key, level, value)
        { [ fst | rst2 ], new_store }
    end
  end

  def get(state, key) do
    get(state.store, state.root, key)
  end

  defp get(store, root, key) do
    case root do
      nil -> nil
      _ ->
        %Page{ level: _, low: low, list: lst } = Store.get(store, root)
        get_aux(store, low, lst, key)
    end
  end

  defp get_aux(store, low, lst, key) do
    case lst do
      [] ->
        get(store, low, key)
      [ {k, v, _} | _ ] when key == k ->
        v
      [ {k, _, _} | _ ] when key < k ->
        get(store, low, key)
      [ {k, _, low2} | rst ] when key > k ->
        get_aux(store, low2, rst, key)
    end
  end

  def dump(state) do
    dump(state.store, state.root, "")
  end

  defp dump(store, root, lvl) do
    case root do
      nil ->
        IO.puts(lvl <> "nil")
      _ ->
        %Page{ level: level, low: low, list: lst} = Store.get(store, root)
        IO.puts(lvl <> "#{root|>Base.encode16} (#{level})")
        dump(store, low, lvl <> "  ")
        for {k, v, r} <- lst do
          IO.puts(lvl<>"- #{inspect k} => #{inspect v}")
          dump(store, r, lvl <> "  ")
        end
    end
  end

  defp calc_level(key) do
    key
      |> SData.term_hash
      |> Base.encode16
      |> String.to_charlist
      |> count_leading_zeroes
  end

  defp count_leading_zeroes('0' ++ rest) do
    1 + count_leading_zeroes(rest)
  end
  defp count_leading_zeroes(_) do
    0
  end

end