package ru.dbotthepony.mc.otm.block.entity.cable

import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap
import it.unimi.dsi.fastutil.objects.ReferenceArraySet
import it.unimi.dsi.fastutil.objects.ReferenceLinkedOpenHashSet
import net.minecraft.core.Direction
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.mc.otm.SERVER_IS_LIVE
import ru.dbotthepony.mc.otm.UNIVERSE_TICKS
import ru.dbotthepony.mc.otm.capability.receiveEnergy
import ru.dbotthepony.mc.otm.config.CablesConfig
import ru.dbotthepony.mc.otm.core.addAll
import ru.dbotthepony.mc.otm.core.collect.map
import ru.dbotthepony.mc.otm.core.collect.reduce
import ru.dbotthepony.mc.otm.core.math.Decimal
import ru.dbotthepony.mc.otm.core.math.RelativeSide
import ru.dbotthepony.mc.otm.core.otmRandom
import ru.dbotthepony.mc.otm.core.shuffle
import ru.dbotthepony.mc.otm.graph.GraphNodeList
import ru.dbotthepony.mc.otm.onceServer
import java.util.*
import kotlin.collections.ArrayList
import kotlin.collections.HashSet
import kotlin.collections.LinkedHashSet
import kotlin.math.ln

private class LinkedPriorityQueue<T : Comparable<T>> {
	private class Entry<T>(val value: T) {
		var next: Entry<T>? = null
	}

	private var head: Entry<T>? = null

	fun isNotEmpty(): Boolean {
		return head != null
	}

	fun add(value: T) {
		if (head == null) {
			head = Entry(value)
		} else if (head!!.value >= value) {
			val entry = Entry(value)
			entry.next = head
			head = entry
		} else {
			var current = head
			var previous = current

			while (current != null) {
				if (current.value >= value) {
					val entry = Entry(value)
					entry.next = current
					previous!!.next = entry
					return
				}

				previous = current
				current = current.next
			}

			previous!!.next = Entry(value)
		}
	}

	fun remove(): T? {
		val head = head

		if (head != null) {
			this.head = head.next
			return head.value
		}

		return null
	}
}

class EnergyCableGraph : GraphNodeList<EnergyCableBlockEntity.Node, EnergyCableGraph>() {
	private val livelyNodes = HashSet<Pair<EnergyCableBlockEntity.Node, RelativeSide>>()
	private val livelyNodesList = ArrayList<Pair<EnergyCableBlockEntity.Node, RelativeSide>>()

	fun addLivelyNode(node: EnergyCableBlockEntity.Node) {
		when (contains(node)) {
			ContainsStatus.ABOUT_TO_BE_REMOVED, ContainsStatus.ABOUT_TO_BE_ADDED -> { } // do nothing
			ContainsStatus.DOES_NOT_BELONG -> throw IllegalArgumentException("$node does not belong to $this")
			ContainsStatus.CONTAINS -> {
				for (dir in RelativeSide.entries) {
					val pair = node to dir

					if (livelyNodes.add(pair)) {
						livelyNodesList.add(pair)
					}
				}
			}
		}
	}

	private class CacheEntry {
		val paths = ArrayList<SegmentPath>(1)
		var saturated = false
		var forceExploreAlternatives = false
	}

	// TODO: LRU cache?
	private val pathCache = HashMap<Pair<EnergyCableBlockEntity.Node, EnergyCableBlockEntity.Node>, CacheEntry>()

	private class SearchNode(
		val node: EnergyCableBlockEntity.Node,
		target: EnergyCableBlockEntity.Node,
		var parent: SearchNode? = null,
		previouslySeen: Boolean
	) : Comparable<SearchNode> {
		val heuristics: Double

		init {
			var heuristics = node.position.distSqr(target.position) * 0.0001

			if (!previouslySeen) {
				heuristics -= ln(node.segment.availableThroughput.coerceAtMost(Decimal.LONG_MAX_VALUE).toDouble())
			}

			this.heuristics = heuristics
		}

		override fun compareTo(other: SearchNode): Int {
			return heuristics.compareTo(other.heuristics)
		}
	}

	class SegmentPath(val a: EnergyCableBlockEntity.Node, val b: EnergyCableBlockEntity.Node) {
		val segments = LinkedHashSet<Segment>()
		val nodes = HashSet<EnergyCableBlockEntity.Node>()
		var shortCircuit = false
		private var lastTickTransfers = 0
		private var lastTick = 0

		val availableThroughput: Decimal
			get() = segments.minOfOrNull { it.availableThroughput } ?: Decimal.ZERO

		fun availableThroughput(instantSnapshot: Reference2ObjectOpenHashMap<Segment, Decimal>): Decimal {
			return segments.minOfOrNull { it.throughput(instantSnapshot) } ?: Decimal.ZERO
		}

		operator fun contains(node: EnergyCableBlockEntity.Node): Boolean {
			return node in nodes
		}

		override fun equals(other: Any?): Boolean {
			return this === other || other is SegmentPath && segments == other.segments
		}

		fun transfer(amount: Decimal, simulate: Boolean, instantSnapshot: Reference2ObjectOpenHashMap<Segment, Decimal>): Decimal {
			if (!amount.isPositive || shortCircuit) {
				return Decimal.ZERO
			}

			if (lastTickTransfers++ >= 40000) {
				if (lastTick == UNIVERSE_TICKS) {
					shortCircuit = true
					LOGGER.warn("Cable path from $a to $b appears to be a short circuit, disabling.")
				} else {
					lastTick = UNIVERSE_TICKS
					lastTickTransfers = 0
				}
			}

			val min = minOf(if (simulate) segments.minOf { it.throughput(instantSnapshot) } else segments.minOf { it.availableThroughput }, amount)

			if (min.isPositive) {
				if (simulate)
					segments.forEach { it.transfer(min, instantSnapshot) }
				else
					segments.forEach { it.transfer(min) }
			}

			return min
		}

		fun triggerBlockstateUpdates() {
			segments.forEach { it.triggerBlockstateUpdates() }
		}

		fun refund(amount: Decimal, simulate: Boolean, instantSnapshot: MutableMap<Segment, Decimal>) {
			segments.forEach { it.refund(amount, simulate, instantSnapshot) }
		}

		fun remove() {
			segments.toTypedArray().forEach { it.remove(this) }
		}
	}

	class Segment() {
		constructor(node: EnergyCableBlockEntity.Node) : this() {
			nodes.add(node)
		}

		val nodes = LinkedHashSet<EnergyCableBlockEntity.Node>()
		var paths = ReferenceLinkedOpenHashSet<SegmentPath>()

		operator fun contains(node: EnergyCableBlockEntity.Node): Boolean {
			return node in nodes
		}

		fun missingNodes(from: Collection<EnergyCableBlockEntity.Node>): List<EnergyCableBlockEntity.Node> {
			val missing = ArrayList<EnergyCableBlockEntity.Node>()

			this.nodes.forEach {
				if (it !in from) {
					missing.add(it)
				}
			}

			return missing
		}

		var throughput = Decimal.ZERO
			private set

		private fun checkThroughput() {
			if (!throughputKnown) {
				throughput = nodes.minOfOrNull { it.energyThroughput } ?: Decimal.ZERO
				throughputKnown = true
			}
		}

		var transferredLastTick = Decimal.ZERO
			private set

		val availableThroughput: Decimal get() {
			if (lastTick != UNIVERSE_TICKS) {
				transferredLastTick = Decimal.ZERO
				lastTick = UNIVERSE_TICKS
			}

			checkThroughput()
			return throughput - transferredLastTick
		}

		private var throughputKnown = false
		private var lastTick = 0
		private var lastPoweredStatus = 0
		private var lastPoweredSwitchFrequency = 0
		private var poweredStatusFlightTicks = 4

		private fun updateBlockstates(): Boolean {
			if (nodes.isEmpty()) return false

			if (lastTick + 2 < UNIVERSE_TICKS) {
				transferredLastTick = Decimal.ZERO
				lastTick = UNIVERSE_TICKS
			}

			val newState = transferredLastTick.isPositive
			val numState = if (newState) 1 else -1

			if (numState != lastPoweredStatus || --poweredStatusFlightTicks > 0) {
				lastPoweredStatus = numState
				nodes.forEach { if (it.segment === this) it.updatePoweredState(newState) }
			}

			return newState
		}

		private var hasBlockstateTimer = false

		private fun blockstateUpdateStep() {
			val new = updateBlockstates()
			val changed = hasBlockstateTimer != new
			hasBlockstateTimer = new
			lastPoweredSwitchFrequency = (lastPoweredSwitchFrequency shl 1) or (if (changed) 1 else 0)

			if (hasBlockstateTimer && SERVER_IS_LIVE) {
				var i = 30

				for (index in 0 .. 26) {
					if (lastPoweredSwitchFrequency and (1 shl index) != 0) {
						i--
					}
				}

				onceServer(i) {
					blockstateUpdateStep()
				}
			}
		}

		fun triggerBlockstateUpdates() {
			if (!hasBlockstateTimer && transferredLastTick > Decimal.ZERO) {
				blockstateUpdateStep()
			}
		}

		fun throughput(instantSnapshot: Reference2ObjectOpenHashMap<Segment, Decimal>): Decimal {
			if (lastTick != UNIVERSE_TICKS) {
				transferredLastTick = Decimal.ZERO
				lastTick = UNIVERSE_TICKS
			}

			checkThroughput()
			val currentTransferred = instantSnapshot[this] ?: transferredLastTick
			return throughput - currentTransferred
		}

		fun transfer(amount: Decimal, instantSnapshot: Reference2ObjectOpenHashMap<Segment, Decimal>): Decimal {
			if (lastTick != UNIVERSE_TICKS) {
				transferredLastTick = Decimal.ZERO
				lastTick = UNIVERSE_TICKS
			}

			checkThroughput()
			val currentTransferred = instantSnapshot[this] ?: transferredLastTick

			if (currentTransferred >= throughput || !amount.isPositive) {
				return Decimal.ZERO
			}

			val new = minOf(currentTransferred + amount, throughput)
			val diff = new - currentTransferred
			instantSnapshot[this] = new
			return diff
		}

		fun transfer(amount: Decimal): Decimal {
			if (lastTick != UNIVERSE_TICKS) {
				transferredLastTick = Decimal.ZERO
				lastTick = UNIVERSE_TICKS
			}

			checkThroughput()
			if (transferredLastTick >= throughput || !amount.isPositive) {
				return Decimal.ZERO
			}

			val new = minOf(transferredLastTick + amount, throughput)
			val diff = new - transferredLastTick
			transferredLastTick = new

			return diff
		}

		fun refund(amount: Decimal, simulate: Boolean, instantSnapshot: MutableMap<Segment, Decimal>) {
			if (simulate) {
				if (this in instantSnapshot) {
					instantSnapshot[this] = maxOf((instantSnapshot[this] ?: Decimal.ZERO) - amount, Decimal.ZERO)
				}
			} else {
				transferredLastTick = maxOf(transferredLastTick - amount, Decimal.ZERO)
			}
		}

		fun remove(node: EnergyCableBlockEntity.Node) {
			check(nodes.remove(node)) { "Tried to remove node $node from segment $this, but that node does not belong to this segment" }

			paths.forEach { it.nodes.remove(node) }
			throughputKnown = false
		}

		fun add(node: EnergyCableBlockEntity.Node) {
			check(nodes.add(node)) { "Tried to add node $node to segment $this, but we already have that node added" }
			paths.forEach { it.nodes.add(node) }

			if (throughputKnown)
				throughput = maxOf(throughput, node.energyThroughput)
		}

		fun add(path: SegmentPath) {
			if (paths.add(path)) {
				check(path.segments.add(this)) { "Path set and Segment disagree whenever $this is absent from $path" }
				path.nodes.addAll(nodes)
			}
		}

		fun remove(path: SegmentPath) {
			check(paths.remove(path)) { "Path $path should already contain $this" }
			check(path.segments.remove(this)) { "Path set and Segment disagree whenever $this is present in $path" }
			path.nodes.removeAll(nodes)
		}

		// breaks "instant snapshots" of segments atm
		// shouldn't cause major gameplay issues though
		fun split(nodes: Collection<EnergyCableBlockEntity.Node>): List<Segment> {
			for (node in nodes)
				require(node in this.nodes) { "$node does not belong to $this" }

			if (this.nodes.size == 1) {
				return listOf(this)
			} else {
				val result = ArrayList<Segment>()
				result.add(this)

				for (v in nodes) {
					v.segment = Segment()
					v.segment.throughput = v.energyThroughput
					v.segment.throughputKnown = true
					v.segment.lastTick = lastTick
					v.segment.transferredLastTick = transferredLastTick
					v.segment.paths = paths.clone()
					v.updatePoweredState(false)

					paths.forEach { it.segments.add(v.segment) }

					result.add(v.segment)
				}

				lastPoweredStatus = 0
				poweredStatusFlightTicks = 4
				paths.forEach { it.nodes.addAll(nodes) }
				return result
			}
		}

		fun split(): List<Segment> {
			val result = ArrayList<Segment>()
			result.add(this)

			for (v in nodes) {
				val segment = Segment()
				v._segment = segment
				segment.nodes.add(v)
				segment.lastTick = lastTick
				segment.throughput = v.energyThroughput
				segment.throughputKnown = true
				segment.transferredLastTick = transferredLastTick
				result.add(segment)
				v.updatePoweredState(false)
			}

			// mark this segment as dead
			this.nodes.clear()
			return result
		}

		private fun combineInto(segment: Segment) {
			nodes.forEach { it._segment = segment }
			segment.nodes.addAll(nodes)
			// this should match on both segments
			// segment.transferredLastTick = maxOf(segment.transferredLastTick, transferredLastTick)
			paths.forEach { it.segments.remove(this) }
			segment.throughputKnown = false
			nodes.clear() // mark segment as dead
		}

		private fun tryCombineWith(segment: Segment): Boolean {
			if (segment === this || segment.nodes.isEmpty() || segment.paths.isEmpty()) return false

			// if nodes have same set of paths then it is safe to combine this with target segment
			if (segment.paths == paths) {
				if (nodes.size >= segment.nodes.size) {
					segment.combineInto(this)
				} else {
					combineInto(segment)
				}

				return true
			}

			return false
		}

		fun tryCombine() {
			if (nodes.isEmpty()) return

			// TODO: This is inefficient as fk
			for (path in paths) {
				for (segment in path.segments) {
					if (tryCombineWith(segment)) return
				}
			}
		}

		fun tryCombine(with: Collection<Segment>) {
			if (nodes.isEmpty() || paths.isEmpty()) return

			for (segment in with) {
				if (tryCombineWith(segment)) return
			}
		}
	}

	// isn't exactly A*, but greedy algorithm, which searched for locally optimal solutions because they lead to globally optimal ones
	private fun findPath(a: EnergyCableBlockEntity.Node, b: EnergyCableBlockEntity.Node, existing: Collection<SegmentPath>, threshold: Decimal): SegmentPath? {
		val seenTop = existing.any { a in it }

		if (seenTop && a.energyThroughput <= threshold) {
			return null
		}

		val openNodes = LinkedPriorityQueue<SearchNode>()
		val seenNodes = HashSet<EnergyCableBlockEntity.Node>()

		openNodes.add(SearchNode(a, b, null, seenTop))

		val existingNodes = HashSet<EnergyCableBlockEntity.Node>()
		existing.forEach { existingNodes.addAll(it.nodes) }

		while (openNodes.isNotEmpty()) {
			val first = openNodes.remove()!!

			if (first.node === b) {
				// solution found
				val solution = LinkedHashSet<EnergyCableBlockEntity.Node>()

				var last = first.parent
				solution.add(first.node)

				while (last != null) {
					solution.add(last.node)
					last = last.parent
				}

				val touchedSegments = ReferenceArraySet<Segment>()

				for (it in solution) {
					if (it.segment in touchedSegments)
						continue

					val diff = it.segment.missingNodes(solution)

					// segment satisfies new constraints
					if (diff.isEmpty()) {
						touchedSegments.add(it.segment)
					} else {
						touchedSegments.addAll(it.segment.split(diff))
					}
				}

				val path = SegmentPath(a, b)
				solution.forEach { it.segment.add(path) }
				touchedSegments.forEach { it.tryCombine(touchedSegments) }

				return path
			} else {
				for (neighbour in first.node.neighboursView.values) {
					if (!seenNodes.add(neighbour) || !neighbour.canTraverse) continue
					val seen = neighbour in existingNodes
					if (seen && neighbour.energyThroughput <= threshold) continue
					openNodes.add(SearchNode(neighbour, b, first, seen))
				}
			}
		}

		return null
	}

	private var lastTickSearches = 0
	private var searchedOnTick = 0

	private fun getPath(a: EnergyCableBlockEntity.Node, b: EnergyCableBlockEntity.Node, energyToTransfer: Decimal, sort: Boolean): List<SegmentPath> {
		if (!a.canTraverse || !b.canTraverse || energyToTransfer <= Decimal.ZERO)
			return listOf()

		if (searchedOnTick != UNIVERSE_TICKS) {
			searchedOnTick = UNIVERSE_TICKS
			lastTickSearches = 0
		}

		val list = pathCache.computeIfAbsent(a to b) { CacheEntry() }

		val maxSearches = CablesConfig.SEARCHES_PER_TICK

		if (!list.saturated && lastTickSearches <= maxSearches) {
			var maxThroughput = list.paths.maxOfOrNull { it.availableThroughput } ?: Decimal.ZERO
			var combinedThroughput = list.paths.iterator().map { it.availableThroughput }.reduce(Decimal.ZERO, Decimal::plus)

			while ((list.forceExploreAlternatives || combinedThroughput < energyToTransfer) && !list.saturated && ++lastTickSearches <= maxSearches) {
				list.forceExploreAlternatives = false
				val find = findPath(a, b, list.paths, maxThroughput)

				if (find == null || find in list.paths) {
					list.saturated = true
				} else {
					list.paths.add(find)
					combinedThroughput += find.availableThroughput
					maxThroughput = maxOf(maxThroughput, find.availableThroughput)
				}
			}
		}

		if (sort)
			list.paths.sortByDescending { it.availableThroughput }

		return list.paths
	}

	fun notifyThroughputsChanged() {
		pathCache.values.forEach {
			it.saturated = false
			it.forceExploreAlternatives = true
		}
	}

	override fun onNodeRemoved(node: EnergyCableBlockEntity.Node) {
		for (dir in RelativeSide.entries) {
			val pair = node to dir

			if (livelyNodes.remove(pair)) {
				check(livelyNodesList.remove(pair))
			}
		}

		val touchedSegments = HashSet<Segment>()

		node.segment.paths.toTypedArray().forEach {
			pathCache.remove(it.a to it.b)

			it.segments.forEach { s ->
				touchedSegments.add(s)
				s.paths.remove(it)

				// this doesn't cover 100% of cases where previously split path segments
				// due to "this" path get recombined, but it should be good enough
				s.paths.forEach { touchedSegments.addAll(it.segments) }
			}
		}

		touchedSegments.addAll(node.segment.split())
		touchedSegments.forEach { it.tryCombine(touchedSegments) }
	}

	override fun onNodeAdded(node: EnergyCableBlockEntity.Node) {
		for (dir in RelativeSide.entries) {
			val pair = node to dir
			check(livelyNodes.add(pair))
			livelyNodesList.add(pair)
		}

		notifyThroughputsChanged()
	}

	fun receiveEnergy(howMuch: Decimal, simulate: Boolean, fromNode: EnergyCableBlockEntity.Node, fromSide: RelativeSide): Decimal {
		livelyNodesList.shuffle(fromNode.blockEntity.level!!.otmRandom)

		val itr = livelyNodesList.iterator()
		var received = Decimal.ZERO
		var residue = howMuch.coerceAtMost(fromNode.energyThroughput)
		val snapshot = Reference2ObjectOpenHashMap<Segment, Decimal>()

		for (pair in itr) {
			val (node, relSide) = pair
			val side = node.sides[relSide]!!

			if (!side.isEnabled) {
				itr.remove()
				check(livelyNodes.remove(pair)) { "Lively nodes Set does not contain $pair" }
				continue
			} else if (fromNode === node && side.side === fromSide) {
				continue
			}

			val it = side.neighbour.get()

			if (it == null || it is EnergyCableBlockEntity.CableSide) {
				itr.remove()
				check(livelyNodes.remove(pair)) { "Lively nodes Set does not contain $pair" }
				continue
			}

			val paths = getPath(fromNode, node, it.receiveEnergy(residue, true), !simulate)

			if (paths.size == 1) {
				// Single path, fast scenario
				val path = paths[0]
				val pathTransferred = path.transfer(residue, simulate, snapshot)
				if (pathTransferred <= Decimal.ZERO) continue
				val thisReceived = it.receiveEnergy(pathTransferred, simulate)

				// If cable transferred more than machine accepted, then "refund" energy
				// so cables record actual value transferred through them
				if (thisReceived != pathTransferred) {
					path.refund(pathTransferred - thisReceived, simulate, snapshot)

					if (!simulate && thisReceived != Decimal.ZERO) {
						path.triggerBlockstateUpdates()
					}
				} else if (!simulate) {
					path.triggerBlockstateUpdates()
				}

				received += thisReceived
				residue -= thisReceived
				if (!residue.isPositive) return received
			} else if (paths.size >= 2) {
				// Multiple paths, a bit more complicated
				// Determine how much machine is likely to accept
				val potentiallyAccepted = it.receiveEnergy(residue, true)

				// Won't accept anything
				if (potentiallyAccepted <= Decimal.ZERO) continue

				// Now determine combined available throughput
				// Make a copy of snapshot, so we can freely write into it
				val copy = snapshot.clone()
				var calcResidue = potentiallyAccepted

				// TODO: Currently, all transfers cause Braess's paradox, because of greedy selection of "fastest" cable
				//  Need to implement heuristics to better distribute load across different paths/segments
				for (path in paths) {
					val passed = path.transfer(calcResidue, true, copy)
					calcResidue -= passed
					if (calcResidue <= Decimal.ZERO) break
				}

				if (calcResidue == potentiallyAccepted) {
					// мда
					continue
				}

				var thisReceived = it.receiveEnergy(potentiallyAccepted - calcResidue, simulate)
				received += thisReceived
				residue -= thisReceived

				for (path in paths) {
					val passed = path.transfer(thisReceived, simulate, snapshot)

					if (!simulate)
						path.triggerBlockstateUpdates()

					thisReceived -= passed
					if (thisReceived <= Decimal.ZERO) break
				}

				if (!residue.isPositive) return received
				//check(thisReceived <= Decimal.ZERO) { "Путом, алло, Путом, какого чёрта Путом? Путом почему ты заблокировал логику, а Путом?" }

				if (thisReceived > Decimal.ZERO) {
					LOGGER.warn("Cable path from $fromNode to $node doesn't follow common sense, disabling.")
					paths.forEach { it.shortCircuit = true }
				}
			}
		}

		return received
	}

	companion object {
		private val LOGGER = LogManager.getLogger()
	}
}
