package me.eater.threedom.dom

import me.eater.threedom.dom.event.*
import me.eater.threedom.dom.render.IRenderNode
import me.eater.threedom.dom.render.IRenderTarget
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) ->
            if (event.child.shouldIndexLocation) {
                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 (tagName in event.tagNames) {
                byTag.getOrPut(tagName, ::mutableSetOf).add(event.node.nodeId)
            }
        }

        on<NodeClassListUpdate.Removed> { (event) ->
            for (tagName in event.tagNames) {
                val set = byTag[tagName] ?: continue

                set.remove(event.node.nodeId)

                if (set.size == 0) {
                    byTag.remove(tagName)
                }
            }
        }

        on<NodeModelUpdate> { (event) ->
            updateKDTreeForNode(event.node)
        }

        on<NodeLocationIndexStateChange> { (event) ->
            if (event.shouldIndex) {
                kdTree.add(event.node)
            } else {
                kdTree.remove(event.node)
            }
        }
    }
    private val kdTree = KDTree(this)

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

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

        for (child in node.recursiveIterator()) {
            child.updateAbsolute()
            if (node.shouldIndexLocation) {
                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 getNodesByTag(tagName: String): Sequence<INode<*>> =
        byTag[tagName]?.asSequence()?.mapNotNull(allNodes::get) ?: emptySequence()

    @Suppress("UNCHECKED_CAST")
    override fun <T : INode<T>> getNodesByClass(className: String): Sequence<T> =
        byType[className]?.asSequence()?.mapNotNull { nodeId ->
            allNodes[nodeId]?.let {
                (it.className == className) as? T
            }
        } ?: emptySequence()

    @Suppress("UNCHECKED_CAST")
    override fun <T : IRenderTarget<C>, C> getNodesByRenderTarget(targetType: String): Sequence<IRenderNode<*, T, C>> =
        byRenderTarget[targetType]?.asSequence()?.mapNotNull(allNodes::get) as? Sequence<IRenderNode<*, T, C>>
            ?: emptySequence()


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

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

        for (className in node.tagList) {
            byTag.getOrPut(className, ::mutableSetOf).add(node.nodeId)
        }

        byType.getOrPut(node.javaClass.name, ::mutableSetOf).add(node.nodeId)

        if (node is IRenderNode<*, *, *>) {
            byRenderTarget.getOrPut(node.target.type, ::mutableSetOf).add(node.nodeId)
        }

        if (node.shouldIndexLocation) {
            kdTree.add(node)
        }
    }

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

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

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


        byType[node.javaClass.name]?.let {
            it.remove(node.nodeId)

            if (it.size == 0) {
                byType.remove(node.javaClass.name)
            }
        }

        if (node is IRenderNode<*, *, *>) {
            byRenderTarget[node.target.type]?.let {
                it.remove(node.nodeId)

                if (it.size == 0) {
                    byRenderTarget.remove(node.target.type)
                }
            }
        }

        if (node.shouldIndexLocation) {
            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)
}