package ru.dbotthepony.mc.otm.graph

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.math.unaryMinus
import ru.dbotthepony.mc.otm.core.orNull
import java.util.EnumMap
import java.util.concurrent.atomic.AtomicInteger

open class Graph6Node<N : Graph6Node<N, G>, G : Abstract6Graph<N, G>>(val graphFactory: () -> G) {
	private val neighbours = EnumMap<Direction, N>(Direction::class.java)

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

	init {
		graph.addNode(this as N)
	}

	private var seen: Int = 0

	operator fun get(direction: Direction): N? = neighbours[direction]

	operator fun set(direction: Direction, node: N?) {
		set(direction, node, false)
	}

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

		val old = neighbours[direction]
		if (old === node) return

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

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

			val opposite = -direction

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

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

			node.neighbours[opposite] = this as N
			neighbours[direction] = node
			node.graph.merge(graph, setter)
			node.onNeighbour(opposite)
			onNeighbour(direction)
		} else {
			neighbours.remove(direction)
		}
	}

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

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

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

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

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

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

	open fun onNeighbour(direction: Direction) {}
	open fun onUnNeighbour(direction: Direction) {}

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

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

			if (!value) {
				val neighbours = ArrayList(neighbours.entries)

				for ((dir, node) in neighbours) {
					breakConnection(this as N, node, dir)
				}

				graph.removeNode(this as N)
				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 directions) {
			val offset = blockPos + dir
			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 val setter = Graph6Node<*, *>::graph::set
		private val nextSeen = AtomicInteger()
		private val directions = Direction.values()

		private fun <N : Graph6Node<N, G>, G : Abstract6Graph<N, G>> breakConnection(a: N, b: N, direction: Direction) {
			val opposite = -direction

			require(a.neighbours[direction] === b) { "$a does not neighbour with $b at direction $direction (forward)" }
			require(b.neighbours[opposite] === a) { "$b does not neighbour with $a at direction $opposite (backward)" }
			require(a.graph === b.graph) { "$a and $b belong to different graphs (${a.graph} vs ${b.graph})" }

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

			val seen = nextSeen.incrementAndGet()
			val flood1 = flood(a, seen)

			if (b.seen != seen) {
				val flood2 = flood(b, seen)

				val big: ArrayList<N>
				val small: ArrayList<N>

				if (flood1.size >= flood2.size) {
					big = flood1
					small = flood2
				} else {
					big = flood2
					small = flood1
				}

				a.graph.retain(ReferenceOpenHashSet(big))
				val newGraph = a.graphFactory.invoke()

				for (node in small) {
					node.graph = newGraph
				}

				for (node in small) {
					if (node.isValid) {
						newGraph.addNode(node)
					}
				}
			}

			if (a.isValid)
				a.onUnNeighbour(direction)

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

		private fun <N : Graph6Node<N, G>, G : Abstract6Graph<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
		}
	}
}
