import scala.collection.immutable.HashMap

class CritBitTrie[T](nodes: HashMap[String, CritBitTrie[T]], isExternalNode: Boolean /*, childrenSize: Int */, data: Option[T]) {
  def this(isExternalNode: Boolean, data: Option[T]) = this(HashMap[String, CritBitTrie[T]](), isExternalNode, 0, data)

  def add(word: String, value: T): CritBitTrie[T] = {
    if (word.isEmpty) {
      return new CritBitTrie[T](nodes, true, childrenSize, Option[T](value))
    }

    val maxPrefix = findMaxPrefix(word)
    if (maxPrefix.isEmpty) {
      val newLeaf = word -> new CritBitTrie[T](true, Option[T](value))
      return new CritBitTrie[T](nodes + newLeaf, isExternalNode, childrenSize + 1, data)
    }

    val keyVal = nodes.find(kv => kv._1.startsWith(maxPrefix)).get
    val edge = keyVal._1
    val trie = keyVal._2

    //добавление новой вершины
    if (edge == maxPrefix) {
      val newLeaf = edge -> trie.add(word.substring(edge.length), value)
      val newNodes = nodes + newLeaf
      val size = countWords(newNodes)
      new CritBitTrie[T](newNodes, isExternalNode, size, data)
    } else {
      val end = edge.substring(maxPrefix.length)
      val middleVertex = new CritBitTrie[T](HashMap(end -> trie), false, trie.size(), None)
      val newLeaf = maxPrefix -> middleVertex.add(word.substring(maxPrefix.length), value)

      val newNodes = (nodes - edge) + newLeaf
      val size = countWords(newNodes)
      new CritBitTrie[T](newNodes, isExternalNode, size, data)
    }
  }

  def remove(word: String): CritBitTrie[T] = {
    if (word.isEmpty) {
      return new CritBitTrie(nodes, false, childrenSize, None)
    }

    val maxPrefix = findMaxPrefix(word)
    if (maxPrefix.isEmpty) {
      return this
    }

    val keyVal = nodes.find(kv => kv._1.startsWith(maxPrefix)).get
    val edge = keyVal._1
    val trie = keyVal._2

    if (edge == maxPrefix) {
      val newLeaf = edge -> trie.remove(word.substring(edge.length))

      if (newLeaf._2.size() == 0) {
        val newNodes = nodes - edge
        val size = countWords(newNodes)
        new CritBitTrie[T](newNodes, isExternalNode, size, data)
      } else {
        newLeaf._2.nodesNumber() match {
          case 1 =>
            val nextLeaf = newLeaf._2.getSingleLeaf()
            if (newLeaf._2.size() == nextLeaf._2.size()) {
              val leaf = (edge + nextLeaf._1) -> nextLeaf._2
              val newNodes = (nodes - edge) + leaf
              val size = countWords(newNodes)
              new CritBitTrie[T](newNodes, isExternalNode, size, data)
            } else {
              val newNodes = nodes + newLeaf
              val size = countWords(newNodes)
              new CritBitTrie[T](newNodes, isExternalNode, size, data)
            }
          case _ =>
            val newNodes = nodes + newLeaf
            val size = countWords(newNodes)
            new CritBitTrie[T](newNodes, isExternalNode, size, data)
        }
      }
    } else {
      this
    }
  }

  def contains(word: String): Boolean = {
    get(word) != None
  }

  def get(word: String): Option[T] = {
    if (word.isEmpty) {
      if (isExternalNode) {
        return data
      } else {
        return None
      }
    }

    val maxPrefix = findMaxPrefix(word)
    if (maxPrefix.isEmpty) {
      return None
    }

    val keyVal = nodes.find(kv => kv._1.startsWith(maxPrefix)).get
    if (keyVal._1 != maxPrefix) {
      return None
    }
    keyVal._2.get(word.substring(maxPrefix.length))
  }

  def size(): Int = {
    if (isExternalNode) {
      childrenSize + 1
    } else {
      childrenSize
    }
  }

  def print(indent: Int = 0) {
    if (isExternalNode) println(" " * indent + "$(" + data.get + ")$")
    for (kv <- nodes) {
      println(" " * indent + kv._1)
      kv._2.print(indent + kv._1.length)
    }
  }

  private def nodesNumber() = nodes.size

  private def getSingleLeaf() = nodes.last


  //сколько слов (вэлью?) хранится в дереве
  private def countWords(nodes: HashMap[String, CritBitTrie[T]]) = nodes.foldLeft(0)(_ + _._2.size())

  private def findMaxPrefix(word: String): String = {
    for (edge <- nodes.keys) {
      val prefix = (word, edge).zipped.takeWhile(Function.tupled(_ == _)).unzip._1.mkString
      if (!prefix.isEmpty) {
        return prefix
      }
    }
    ""
  }
}

object CritBitTrie {
  def apply[T](): CritBitTrie[T] = {
    new CritBitTrie[T](false, None)
  }
}