package compose.utils

import androidx.compose.ui.Modifier
import androidx.compose.ui.node.ModifierNodeElement
import androidx.compose.ui.node.TraversableNode
import androidx.compose.ui.node.findNearestAncestor
import androidx.compose.ui.node.traverseChildren
import kotlin.random.Random


/**
 * A modifier element that can be used to tag, and then find nodes in the modifier node tree.
 */
class TagModifierElement<T>(


    /**
     * A traverse key that should be shared among instances of the same kind.
     * You can use [generateTraverseKey] to generate this.
     */
    val traverseKey: Any,


    /**
     * A piece of data that can be associated with the node.
     */
    val data: T


): ModifierNodeElement<TagModifierNode<T>>() {


    /**
     * The [TagModifierNode] node associated with this modifier element.
     */
    var modifierNode: TagModifierNode<T>? = null


    override fun create() = TagModifierNode(traverseKey, data)
        .also {
            modifierNode = it
        }


    override fun update(node: TagModifierNode<T>) {
        node.data = data
        modifierNode = node
    }


    override fun equals(other: Any?): Boolean {
        return other === this
    }


    override fun hashCode(): Int = System.identityHashCode(this)


}


/**
 * The modifier node actually in the tree.
 */
class TagModifierNode<T>(


    /**
     * The traverse key of this node.
     */
    override val traverseKey: Any,

    /**
     * The data associated with this node.
     */
    var data: T,


) : Modifier.Node(), TraversableNode


/**
 * Traverses all the given node's siblings that share the same traversal key (including itself).
 *
 * Note that this requires an ancestor [TagModifierNode] with [parentTraverseKey] to exist in the tree, as the
 * underlying API doesn't currently allow traversing siblings directly.
 */
fun <T> TagModifierNode<T>.traverseSiblings(
    parentTraverseKey: Any,
    block: (TagModifierNode<T>) -> Boolean
) {
    val parent = findNearestAncestor(key = parentTraverseKey) ?: error("No matching ancestor found")
    parent.traverseChildren(key = traverseKey) {
        @Suppress("UNCHECKED_CAST")
        block(it as TagModifierNode<T>)
    }
}


/**
 * Finds the sibling of the given node that precedes it in the tree by [count] nodes, or the first sibling if there
 * aren't [count] siblings preceding it.
 *
 * Note that this requires an ancestor [TagModifierNode] with [parentTraverseKey] to exist in the tree, as the
 * underlying API doesn't currently allow traversing siblings directly.
 */
fun <T> TagModifierNode<T>.findPrecedingSibling(
    parentTraverseKey: Any,
    count: Int
): TagModifierNode<T> {
    val deque = ArrayDeque<TagModifierNode<T>>()
    traverseSiblings(parentTraverseKey) {
        deque.addLast(it)
        if (deque.size > count + 1)
            deque.removeFirst()

        it != this
    }
    return deque.first()
}


/**
 * Finds the sibling of the given node that immediately precedes it in the tree, or the node itself if it is the first
 * child.
 *
 * Note that this requires an ancestor [TagModifierNode] with [parentTraverseKey] to exist in the tree, as the
 * underlying API doesn't currently allow traversing siblings directly.
 */
fun <T> TagModifierNode<T>.findPreviousSibling(parentTraverseKey: Any): TagModifierNode<T> =
    findPrecedingSibling(parentTraverseKey, 1)


/**
 * Finds the sibling of the given node that follows it in the tree by [count] nodes, or the last sibling if there
 * aren't [count] siblings following it.
 *
 * Note that this requires an ancestor [TagModifierNode] with [parentTraverseKey] to exist in the tree, as the
 * underlying API doesn't currently allow traversing siblings directly.
 */
fun <T> TagModifierNode<T>.findFollowingSibling(
    parentTraverseKey: Any,
    count: Int
): TagModifierNode<T> {
    var result: TagModifierNode<T>? = null
    var remainingCount = -1
    traverseSiblings(parentTraverseKey) {
        result = it
        if (it == this) {
            remainingCount = count - 1
            true
        }
        else if (remainingCount == 0) {
            false
        }
        else {
            remainingCount -= 1
            true
        }
    }
    return result!!
}


/**
 * Finds the sibling of the given node that immediately follows it in the tree, or the node itself if it is the last
 * child.
 *
 * Note that this requires an ancestor [TagModifierNode] with [parentTraverseKey] to exist in the tree, as the
 * underlying API doesn't currently allow traversing siblings directly.
 */
fun <T> TagModifierNode<T>.findNextSibling(parentTraverseKey: Any): TagModifierNode<T> =
    findFollowingSibling(parentTraverseKey, 1)


/**
 * Finds the first sibling of the given node, which may be the node itself.
 */
fun <T> TagModifierNode<T>.findFirstSibling(
    parentTraverseKey: Any
) : TagModifierNode<T> {
    var result: TagModifierNode<T>? = null
    traverseSiblings(parentTraverseKey) {
        result = it
        false
    }

    return result!!
}


/**
 * Finds the last sibling of the given node, which may be the node itself.
 */
fun <T> TagModifierNode<T>.findLastSibling(
    parentTraverseKey: Any
) : TagModifierNode<T> {
    var result: TagModifierNode<T>? = null
    traverseSiblings(parentTraverseKey) {
        result = it
        true
    }

    return result!!
}


/**
 * The characters
 */
private val CharPool : List<Char> = ('a'..'z') + ('A'..'Z') + ('0'..'9')


/**
 * Generates a random traverse key with the given prefix.
 */
fun generateTraverseKey(prefix: String): String =
    prefix + (1..10)
        .map {
            Random.nextInt(0, CharPool.size).let { CharPool[it] }
        }
        .joinToString("")