threedom/threedom/src/main/kotlin/me/eater/threedom/dom/Document.kt

package me.eater.threedom.dom

import me.eater.threedom.dom.event.DOMTreeUpdate
import me.eater.threedom.dom.event.NodeClassListUpdate
import me.eater.threedom.dom.event.NodeIDUpdate
import me.eater.threedom.dom.event.NodeModelUpdate
import me.eater.threedom.event.Event
import me.eater.threedom.generated.EventNames
import me.eater.threedom.utils.KDTree
import org.joml.Vector3dc
import kotlin.reflect.KClass

class Document : IDocument {
    val root: PlainNode = PlainNode(this)
    private val eventTree = EventTree {
        on<DOMTreeUpdate.Remove> { (event) ->
            removeNodeFromSearch(event.child)
        }

        on<DOMTreeUpdate.Insert> { (event) ->
            addNodeToSearch(event.child)
        }

        on<DOMTreeUpdate.Move> { (event) ->
            updateKDTreeForNode(event.child)
        }

        on<NodeIDUpdate> { (event) ->
            event.old?.let {
                byId.remove(it, event.node.nodeId)
            }

            event.new?.let {
                byId.putIfAbsent(it, event.node.nodeId)
            }
        }

        on<NodeClassListUpdate.Added> { (event) ->
            for (className in event.classNames) {
                byClass.getOrPut(className, ::mutableSetOf).add(event.node.nodeId)
            }
        }

        on<NodeClassListUpdate.Removed> { (event) ->
            for (className in event.classNames) {
                val set = byClass[className] ?: continue

                set.remove(event.node.nodeId)

                if (set.size == 0) {
                    byClass.remove(className)
                }
            }
        }

        on<NodeModelUpdate> { (event) ->
            updateKDTreeForNode(event.node)
        }
    }
    private val kdTree = KDTree(this)

    private val allNodes: MutableMap<Long, INode<*>> = mutableMapOf()
    private val byId: MutableMap<String, Long> = mutableMapOf()
    private val byClass: MutableMap<String, MutableSet<Long>> = mutableMapOf()

    private fun updateKDTreeForNode(node: INode<*>) {
        node.updateAbsolute()
        kdTree.update(node)

        for (child in node.recursiveIterator()) {
            child.updateAbsolute()
            kdTree.update(child)
        }
    }

    override fun <T> addEventListener(eventName: String, refNode: INode<*>, block: (Event<T>) -> Unit) {
        @Suppress("UNCHECKED_CAST")
        eventTree.addEventListener(eventName, refNode, block as (Event<*>) -> Unit)
    }

    fun <T> addEventListener(eventName: String, block: (Event<T>) -> Unit) = addEventListener(eventName, root, block)

    fun <T> addTopLevelEventListener(eventName: String, block: (Event<T>) -> Unit) =
        @Suppress("UNCHECKED_CAST")
        eventTree.addTopLevelEventListener(eventName, block as (Event<*>) -> Unit)

    override fun trigger(eventName: String, targetNode: INode<*>, event: Event<*>) =
        eventTree.trigger(eventName, targetNode, event)


    override fun addNode(newNode: INode<*>) = root.addNode(newNode)
    override fun removeNode(refNode: INode<*>) = root.removeNode(refNode)

    override fun deleteNode(refNode: INode<*>) {
        refNode.detachFromDocument()
        eventTree.removeNode(refNode)
    }

    override fun removeAll() = root.removeAll()

    override fun replaceNode(newNode: INode<*>, refNode: INode<*>): Boolean = root.replaceNode(newNode, refNode)

    override fun hasChild(refNode: INode<*>): Boolean = root.hasChild(refNode)

    override fun sequence(): Sequence<INode<*>> = root.sequence()
    override fun iterator(): Iterator<INode<*>> = root.iterator()

    override fun getNodeById(id: String): INode<*>? = byId[id]?.let(allNodes::get)

    override fun getNodesByClassName(className: String): Sequence<INode<*>> =
        byClass[className]?.asSequence()?.mapNotNull(allNodes::get) ?: emptySequence()


    private fun addNodeToSearch(node: INode<*>) {
        allNodes[node.nodeId] = node

        node.id?.let { byId.putIfAbsent(it, node.nodeId) }

        for (className in node.classList) {
            byClass.getOrPut(className, ::mutableSetOf).add(node.nodeId)
        }

        kdTree.add(node)
    }

    private fun removeNodeFromSearch(node: INode<*>) {
        allNodes.remove(node.nodeId)

        node.id?.let { byId.remove(it, node.nodeId) }

        for (className in node.classList) {
            val set = byClass.get(className) ?: continue
            set.remove(node.nodeId)
            if (set.size == 0) {
                byClass.remove(className)
            }
        }

        kdTree.remove(node)
    }

    inline fun <reified T> on(topLevel: Boolean = false, noinline block: (Event<T>) -> Unit) = if (topLevel)
        addTopLevelEventListener(EventNames.getEventName<T>(), block)
    else
        addEventListener(EventNames.getEventName<T>(), block)

    override fun <T : INode<T>> createNode(nodeType: KClass<T>): T =
        nodeType.java.getConstructor(IDocument::class.java).newInstance(this)

    override fun getNodeByNodeId(nodeId: Long): INode<*>? = this.allNodes[nodeId]

    override fun findAt(vec: Vector3dc) = kdTree.find(vec)
    override fun rebalance() = kdTree.rebalance()
    override fun findInRegion(pointA: Vector3dc, pointB: Vector3dc) = kdTree.findInRegion(pointA, pointB)
    override fun findInRange(origin: Vector3dc, range: Number) = kdTree.findInRange(origin, range)
}