package dev.austinzhu.algods.containers.tree

import dev.austinzhu.algods.containers.util.NAryTreePrinter
import dev.austinzhu.algods.containers.util.Operation
import kotlin.math.sqrt
import kotlin.random.Random

class Trie<T> : AbstractTrie<T>() {
    open inner class Node(var value: T? = null) : Tree.Node<String, T, Node> {
        val next = arrayOfNulls<Node>(Char.MAX_VALUE.code)

        override fun getChildren(): List<Node?> {
            class NodeWrapper(val char: Char, val node: Node?) : Node() {
                override fun toString(): String {
                    if (node == null) return ""
                    return if (node.value == null) "'$char'" else "'$char': ${node.value}"
                }

                override fun getChildren(): List<NodeWrapper?> {
                    return if (node == null) emptyList() else listOf(*node.next).indices.map {
                        if (node.next[it] != null) NodeWrapper(it.toChar(), node.next[it]) else null
                    }
                }
            }
            return NodeWrapper('\u0000', this).getChildren()
        }

        override fun hasAll() = next.all { it != null }

        override fun hasNone() = next.all { it == null }

        override fun toString() = "''"
    }

    companion object {
        fun init(size: Int, bound: Int): Trie<Int> {
            val capacity = Random.nextInt(size)
            val trie = Trie<Int>()
            val keys = Array(capacity) { it.toString() }
            val alphabets = ('a'..'z').toList()
            val prefixes = Array(sqrt(capacity.toFloat()).toInt()) {
                (1..capacity)
                    .map { Random.nextInt(0, alphabets.size) }
                    .map(alphabets::get)
                    .joinToString("")
            }
            for (k in keys) {
                val i = Random.nextInt(prefixes.size)
                trie.add(prefixes[i] + k, Random.nextInt(bound))
            }
            return trie
        }
    }

    @Operation
    override fun add(idx: String, value: T) {
        if (root == null) {
            root = Node()
        }
        var node = root!!
        for (c in idx) {
            val code = c.code
            if (node.next[code] == null) {
                node.next[code] = Node()
            }
            node = node.next[code]!!
        }
        node.value = value
    }

    @Operation
    override fun get(idx: String): T {
        var node = root ?: throw NoSuchElementException()
        for (c in idx) {
            val code = c.code
            node = node.next[code] ?: throw NoSuchElementException()
        }
        return node.value ?: throw NoSuchElementException()
    }

    @Operation
    override fun set(idx: String, value: T) {
        var last = root ?: throw NoSuchElementException()
        for (c in idx) {
            last = last.next[c.code] ?: throw NoSuchElementException()
        }
        last.value = value
    }

    @Operation
    override fun delete(idx: String): T {
        var del: T? = null
        fun helper(node: Node? = root, depth: Int = 0): Node? {
            if (node == null) return null
            if (depth == idx.length) {
                del = node.value
                node.value = null
            } else {
                val code = idx[depth].code
                node.next[code] = helper(node.next[code], depth + 1)
            }
            // remove subtrie if it is empty
            if (node.value != null) {
                return node
            }
            for (n in node.next) {
                if (n != null) {
                    return n
                }
            }
            return null
        }
        helper()
        return del ?: throw NoSuchElementException()
    }

    @Operation
    fun keysWithPrefix(prefix: String): ArrayList<String> {
        val res = ArrayList<String>()
        fun collect(node: Node?, prefixChars: MutableList<Char> = prefix.toMutableList()) {
            if (node == null) return
            if (node.value != null) {
                val str = prefixChars.fold("", String::plus)
                res.add(str)
            }
            for (c in 0 until 256) {
                prefixChars.add(c.toChar())
                collect(node.next[c], prefixChars)
                prefixChars.removeLast()
            }
        }

        var last: Node? = root
        for (c in prefix) {
            if (last == null) return res
            val code = c.code
            last = last.next[code]
        }
        collect(last)
        return res
    }

    override fun toString(): String {
        return NAryTreePrinter(root).toString()
    }
}
