package moe.kageru.spektacles

import kotlin.collections.LinkedHashMap


// maybe try this instead?
// https://arrow-kt.io/docs/recursion/intro/
sealed class Node<K, T : Any>(open val value: Any) {
  class Leaf<K, T : Any>(override val value: List<T>) : Node<K, T>(value)
  class Split<K, T : Any>(val children: Map<K, Node<K, T>>) : Node<K, T>(children)
  class Value<K, T : Any>(override val value: T) : Node<K, T>(value)

  fun <R : Any> map(op: (T) -> R): Node<K, R> {
    return when (this) {
      is Leaf -> Leaf(value.map(op))
      is Value -> Value(op(value))
      is Split -> Split(children.mapValues { it.value.map(op) })
    }
  }

  fun groupBy(keyGenerator: (T) -> K): Node<K, T> {
    return when (this) {
      is Leaf -> Split(value.groupBy(keyGenerator).mapValues { Leaf<K, T>(it.value) })
      is Value -> this
      is Split -> Split(children.mapValues { it.value.groupBy(keyGenerator) })
    }
  }

  override fun toString() = asString()

  private fun asString(depth: Int = 0): String {
    return when (this) {
      is Leaf -> value.joinToString(indent(depth), prefix = "\n${indent(depth)}")
      is Value -> value.toString()
      is Split -> children.entries.joinToString("") { "\n${indent(depth)}${it.key}: ${it.value.asString(depth + 1)}" }
    }
  }

  fun countLeafs(): Node<K, Int> {
    return when (this) {
      is Leaf -> Value(value.size)
      is Value -> Value(1)
      is Split -> Split(children.mapValues { it.value.countLeafs() })
    }
  }

  private fun indent(depth: Int) = (0..(depth * 2)).joinToString("") { " " }
}

fun <K> Node<K, Int>.sort(): Node<K, Int> {
  return when (this) {
    is Node.Leaf -> this
    is Node.Value -> this
    is Node.Split -> Node.Split(this.children.sortByValues().mapValues { it.value.sort() })
  }
}

private fun <K> Map<K, Node<K, Int>>.sortByValues(): Map<K, Node<K, Int>> {
  val sortedEntries = entries.sortedBy {
    when (val v = it.value) {
      is Node.Value -> v.value
      is Node.Leaf -> v.value.size
      is Node.Split -> Int.MAX_VALUE
    }
  }
  return LinkedHashMap<K, Node<K, Int>>().apply {
    for ((key, value) in sortedEntries) {
      this[key] = value
    }
  }
}