package ru.dbotthepony.mc.otm.graph

import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap
import it.unimi.dsi.fastutil.objects.ReferenceOpenHashSet
import net.minecraft.core.BlockPos
import net.minecraft.core.Direction
import net.minecraft.core.SectionPos
import net.minecraft.server.level.ServerLevel
import net.minecraft.world.level.block.entity.BlockEntity
import net.minecraftforge.common.capabilities.Capability
import ru.dbotthepony.mc.otm.addTicker
import ru.dbotthepony.mc.otm.core.math.plus
import ru.dbotthepony.mc.otm.core.orNull
import ru.dbotthepony.mc.otm.core.util.IConditionalTickable
import ru.dbotthepony.mc.otm.core.util.ITickable
import java.util.*
import kotlin.collections.ArrayList

open class GraphNode<N : GraphNode<N, G>, G : GraphNodeList<N, G>>(val graphFactory: () -> G) {
	interface Link {
		val opposite: Link
		operator fun unaryMinus() = opposite
	}

	data class DirectionLink(val direction: Direction) : Link {
		override val opposite: Link
			get() = wrapped[direction.opposite.ordinal]

		override fun toString(): String {
			return "Link[$direction]"
		}
	}

	private val neighbours = Object2ObjectOpenHashMap<Link, N>()
	val neighboursView: Map<Link, N> = Collections.unmodifiableMap(neighbours)

	var graph: G = graphFactory.invoke()
		internal set

	init {
		check(graph.addNodeQueued(this as N))
	}

	private var seen: Int = 0

	fun beginTicking() {
		require(this is IConditionalTickable || this is ITickable) { "Node must implement either ITickable or IConditionalTickable to tick" }
		graph.beginTicking(this as N)
	}

	operator fun get(key: Link): N? = neighbours[key]

	operator fun set(key: Link, node: N?) {
		set(key, node, false)
	}

	fun set(key: Link, target: N?, allowReplacement: Boolean) {
		check(isValid) { "Can not neighbour any node while this node is invalid" }

		val old = neighbours[key]
		if (old === target) return

		if (old != null)
			breakConnection(this as N, old, key)

		if (target != null) {
			require(target.isValid) { "Can not neighbour invalid node" }

			val opposite = key.opposite

			if (allowReplacement) {
				target.neighbours[opposite]?.let {
					breakConnection(target, it, opposite)
				}

				check(target.neighbours[opposite] == null) { "$target didn't break connection at direction $opposite" }
			} else {
				check(target.neighbours[opposite] == null) { "Trying to form connection from $this to $target at $key, but $target already has neighbour at $opposite (${target.neighbours[opposite]})!" }
			}

			target.neighbours[opposite] = this as N
			neighbours[key] = target
			target.graph.merge(graph)
			target.onNeighbour(opposite)
			onNeighbour(key)
		} else {
			neighbours.remove(key)
		}
	}

	var top: N?
		get() = neighbours[UP]
		set(value) {
			set(UP, value)
		}

	var bottom: N?
		get() = neighbours[DOWN]
		set(value) {
			set(DOWN, value)
		}

	var south: N?
		get() = neighbours[SOUTH]
		set(value) {
			set(SOUTH, value)
		}

	var north: N?
		get() = neighbours[NORTH]
		set(value) {
			set(NORTH, value)
		}

	var west: N?
		get() = neighbours[WEST]
		set(value) {
			set(WEST, value)
		}

	var east: N?
		get() = neighbours[EAST]
		set(value) {
			set(EAST, value)
		}

	protected open fun onNeighbour(link: Link) {}
	protected open fun onUnNeighbour(link: Link) {}

	protected open fun invalidate() {}
	protected open fun revive() {}

	var isValid: Boolean = true
		set(value) {
			if (value == field) return
			field = value

			if (!value) {
				val neighbours = neighbours.entries.map { it.key to it.value }

				for ((link, node) in neighbours) {
					removeNeighbours(this as N, node, link)
				}

				graph.removeNode(this as N)
				rebuildGraphs(neighbours.map { it.second })
				invalidate()
			} else {
				revive()
			}
		}

	fun discover(
		level: ServerLevel,
		blockPos: BlockPos,
		nodeGetter: (BlockEntity) -> N?
	) {
		if (!isValid) return

		level.addTicker {
			isValid && !discoverStep(level, blockPos, nodeGetter)
		}
	}

	fun discover(
		level: ServerLevel,
		blockPos: BlockPos,
		capability: Capability<out N>
	) {
		if (!isValid) return

		level.addTicker {
			isValid && !discoverStep(level, blockPos) { it.getCapability(capability).orNull() }
		}
	}

	fun discover(blockEntity: BlockEntity, nodeGetter: (BlockEntity) -> N?) {
		discover(blockEntity.level as? ServerLevel ?: return, blockEntity.blockPos, nodeGetter)
	}

	fun discover(blockEntity: BlockEntity, capability: Capability<out N>) {
		discover(blockEntity.level as? ServerLevel ?: return, blockEntity.blockPos, capability)
	}

	fun discoverStep(
		level: ServerLevel,
		blockPos: BlockPos,
		nodeGetter: (BlockEntity) -> N?,
	): Boolean {
		if (!isValid) return false
		var fullDiscovery = true

		for (dir in wrapped) {
			val offset = blockPos + dir.direction
			val chunk = level.chunkSource.getChunkNow(SectionPos.blockToSectionCoord(offset.x), SectionPos.blockToSectionCoord(offset.z))

			if (chunk == null) {
				fullDiscovery = false
				set(dir, null)
				continue
			}

			val entity = chunk.getBlockEntity(offset)

			if (entity != null) {
				set(dir, nodeGetter(entity))
			} else {
				set(dir, null)
			}
		}

		return fullDiscovery
	}

	companion object {
		private var nextSeen = 0
		private val wrapped = Direction.entries.map { DirectionLink(it) }

		fun link(direction: Direction): Link {
			return wrapped[direction.ordinal]
		}

		val UP = link(Direction.UP)
		val DOWN = link(Direction.DOWN)
		val NORTH = link(Direction.NORTH)
		val SOUTH = link(Direction.SOUTH)
		val WEST = link(Direction.WEST)
		val EAST = link(Direction.EAST)

		private fun <N : GraphNode<N, G>, G : GraphNodeList<N, G>> removeNeighbours(a: N, b: N, key: Link) {
			val opposite = key.opposite

			// проверяем, действительно ли a соединён с b по key
			require(a.neighbours[key] === b) { "$a does not neighbour with $b at $key (a -> b)" }

			// проверяем обратную связь
			require(b.neighbours[opposite] === a) { "$b does not neighbour with $a at $opposite (b -> a)" }

			// находятся ли они в одном графе
			require(a.graph === b.graph) { "$a and $b belong to different graphs (${a.graph} / ${b.graph})" }

			a.neighbours.remove(key)
			b.neighbours.remove(opposite)

			if (a.isValid) a.onUnNeighbour(key)
			if (b.isValid) b.onUnNeighbour(key.opposite)
		}

		private fun <N : GraphNode<N, G>, G : GraphNodeList<N, G>> rebuildGraphs(nodes: Collection<N>) {
			if (nodes.isEmpty())
				return

			val seen = ++nextSeen
			val floods = ArrayList<ArrayList<N>>()

			for (node in nodes) {
				if (node.seen < seen) {
					floods.add(flood(node, seen))
				}
			}

			floods.sortedBy { it.first().graph.size }
			floods.first().first().graph.retain(ReferenceOpenHashSet(floods.first()))

			for (i in 1 until floods.size) {
				val graph = floods[i].first().graphFactory.invoke()

				for (node in floods[i]) {
					graph.addNode(node)
				}
			}
		}

		private fun <N : GraphNode<N, G>, G : GraphNodeList<N, G>> breakConnection(a: N, b: N, key: Link) {
			removeNeighbours(a, b, key)
			rebuildGraphs(listOf(a, b))
		}

		private fun <N : GraphNode<N, G>, G : GraphNodeList<N, G>> flood(startingNode: N, seen: Int): ArrayList<N> {
			val unopen = ArrayList<N>()
			val result = ArrayList<N>()
			unopen.add(startingNode)

			while (unopen.isNotEmpty()) {
				val last = unopen.removeLast()

				if (last.seen < seen) {
					result.add(last)
					last.seen = seen

					for (node in last.neighbours.values) {
						if (node.seen < seen) {
							unopen.add(node)
						}
					}
				}
			}

			return result
		}
	}
}
