package ru.dbotthepony.kstarbound.world

import it.unimi.dsi.fastutil.ints.IntArraySet
import org.lwjgl.BufferUtils
import ru.dbotthepony.kommons.arrays.Object2DArray
import ru.dbotthepony.kommons.util.IStruct4f
import ru.dbotthepony.kommons.math.RGBAColor
import ru.dbotthepony.kommons.math.linearInterpolation
import ru.dbotthepony.kstarbound.Starbound
import ru.dbotthepony.kstarbound.world.api.ICellAccess
import java.nio.ByteBuffer
import java.util.concurrent.Callable
import java.util.concurrent.CompletableFuture
import java.util.concurrent.ConcurrentLinkedQueue
import java.util.concurrent.ForkJoinPool
import java.util.concurrent.Future
import java.util.concurrent.locks.LockSupport
import java.util.function.Supplier
import kotlin.math.roundToInt
import kotlin.random.Random

// this implementation quite closely resembles original code, mostly
// because i found no other solution for light spreading
// however, code flow is HEAVILY altered, with many additions and changes
class LightCalculator(val parent: ICellAccess, val width: Int, val height: Int) {
	enum class Quality(val secondDiagonal: Boolean, val extraCell: Boolean) {
		HIGH(true, false),
		MEDIUM(false, true),
		LOW(false, false)
	}

	interface ICell : IStruct4f {
		val red: Float
		val green: Float
		val blue: Float

		val alpha: Float get() {
			return (red * red + green * green + blue * blue).coerceIn(0f, 1f)
		}

		override fun component1(): Float {
			return red
		}

		override fun component2(): Float {
			return green
		}

		override fun component3(): Float {
			return blue
		}

		override fun component4(): Float {
			return alpha
		}
	}

	private fun interface Getter {
		operator fun get(x: Int, y: Int): Grid.LightCell
	}

	private inner class Grid {
		inner class LightCell(val x: Int, val y: Int) : ICell {
			val actualDropoff by lazy(LazyThreadSafetyMode.NONE) {
				val parent = this@LightCalculator.parent.getCell(x, y) ?: return@lazy 0f
				val lightBlockStrength: Float

				if (parent.foreground.material.value.renderParameters.lightTransparent) {
					lightBlockStrength = 0f
				} else {
					lightBlockStrength = 1f
				}

				linearInterpolation(lightBlockStrength, invMaxAirSpread, invMaxObstacleSpread)
			}

			private var empty = true
			override var red: Float = 0f
			override var green: Float = 0f
			override var blue: Float = 0f

			fun spreadInto(target: LightCell, drop: Float, alreadyHadChanges: Boolean): Boolean {
				val max = red.coerceAtLeast(green).coerceAtLeast(blue)
				if (max <= 0f) return alreadyHadChanges
				val newDrop = 1f - actualDropoff / max * drop

				@Suppress("name_shadowing")
				var alreadyHadChanges = alreadyHadChanges
				val nred = target.red.coerceAtLeast(red * newDrop)
				val ngreen = target.green.coerceAtLeast(green * newDrop)
				val nblue = target.blue.coerceAtLeast(blue * newDrop)

				if (!alreadyHadChanges)
					alreadyHadChanges = nred != target.red || ngreen != target.green || nblue != target.blue

				target.red = nred
				target.green = ngreen
				target.blue = nblue

				if (target.empty && (target.red >= epsilon || target.blue >= epsilon || target.green >= epsilon)) {
					minX = minX.coerceAtMost(target.x - 1)
					minY = minY.coerceAtMost(target.y - 1)
					maxX = maxX.coerceAtLeast(target.x + 1)
					maxY = maxY.coerceAtLeast(target.y + 1)
					target.empty = false
					clampRect()
				}

				return alreadyHadChanges
			}
		}

		var minX = width
		var maxX = 0
		var minY = height
		var maxY = 0

		fun clampRect() {
			minX = minX.coerceIn(1, width - 2)
			maxX = maxX.coerceIn(1, width - 2)
			minY = minY.coerceIn(1, height - 2)
			maxY = maxY.coerceIn(1, height - 2)
		}

		init {
			clampRect()
		}

		private val mem by lazy(LazyThreadSafetyMode.NONE) {
			Object2DArray.nulls<LightCell>(width, height)
		}

		operator fun get(x: Int, y: Int): LightCell {
			var get = mem[x, y]

			if (get == null) {
				get = LightCell(x, y)
				mem[x, y] = get
			}

			return get
		}

		fun get0(x: Int, y: Int) = mem[x, y]

		fun safeGet(x: Int, y: Int): LightCell? {
			if (x in 0 until width && y in 0 until height) {
				return this[x, y]
			} else {
				return null
			}
		}

		/*
		private val mem = Int2ObjectOpenHashMap<LightCell>(2048)

		operator fun get(x: Int, y: Int): LightCell {
			var get = mem[x shl 16 or y]

			if (get == null) {
				get = LightCell(x, y)
				mem[x shl 16 or y] = get
			}

			return get
		}

		fun get0(x: Int, y: Int) = mem[x shl 16 or y]

		fun safeGet(x: Int, y: Int): LightCell? {
			if (x in 0 until width && y in 0 until height) {
				return this[x, y]
			} else {
				return null
			}
		}*/

		private inner class Copy : Getter {
			private val minX = this@Grid.minX
			private val maxX = this@Grid.maxX
			private val minY = this@Grid.minY
			private val maxY = this@Grid.maxY

			private val mem = Object2DArray(maxX - minX + 3, maxY - minY + 3) { a, b -> this@Grid[a + minX - 1, b + minY - 1] }

			override fun get(x: Int, y: Int): LightCell {
				return mem[x - minX + 1, y - minY + 1]
			}
		}

		private val passthrough = Getter { x, y -> this@Grid[x, y] }

		fun calculateSpread() {
			if (minX > maxX || minY > maxY) return

			// spread light in several passes
			var repeats = passes

			var copy: Getter? = null

			while (repeats-- >= 0) {
				val minX = minX
				val maxX = maxX
				val minY = minY
				val maxY = maxY
				var changes = false

				if (copy == null) {
					copy = if (maxX - minX >= width / 2 || maxY - minY >= height / 2) {
						passthrough
					} else {
						Copy()
					}
				}

				// bottom to top
				for (y in minY .. maxY) {
					// left to right
					for (x in minX .. maxX) {
						val current = copy[x, y]

						changes = current.spreadInto(copy[x, y + 1], 1f, changes)
						changes = current.spreadInto(copy[x + 1, y], 1f, changes)
						changes = current.spreadInto(copy[x + 1, y + 1], 1.4142135f, changes)

						// original code performs this spread to camouflage prism shape of light spreading
						// we instead gonna do light pass on different diagonal
						if (quality.extraCell)
							changes = current.spreadInto(copy[x + 1, y - 1], 1.4142135f, changes)
					}

					// right to left
					if (quality.secondDiagonal) {
						for (x in maxX downTo minX) {
							val current = copy[x, y]

							changes = current.spreadInto(copy[x, y + 1], 1f, changes)
							changes = current.spreadInto(copy[x - 1, y], 1f, changes)
							changes = current.spreadInto(copy[x - 1, y + 1], 1.4142135f, changes)
						}
					}
				}

				// top to bottom
				for (y in maxY downTo minY) {
					// right to left
					for (x in maxX downTo minX) {
						val current = copy[x, y]

						changes = current.spreadInto(copy[x, y - 1], 1f, changes)
						changes = current.spreadInto(copy[x - 1, y], 1f, changes)
						changes = current.spreadInto(copy[x - 1, y - 1], 1.4142135f, changes)

						// original code performs this spread to camouflage prism shape of light spreading
						// we instead gonna do light pass on different diagonal
						if (quality.extraCell)
							current.spreadInto(copy[x - 1, y + 1], 1.4142135f, changes)
					}

					// left to right
					if (quality.secondDiagonal) {
						for (x in minX .. maxX) {
							val current = this[x, y]

							changes = current.spreadInto(copy[x, y - 1], 1f, changes)
							changes = current.spreadInto(copy[x + 1, y], 1f, changes)
							changes = current.spreadInto(copy[x + 1, y - 1], 1.4142135f, changes)
						}
					}
				}

				// if our boundaries have updated, re-spread light
				if (
					minX != this.minX ||
					maxX != this.maxX ||
					minY != this.minY ||
					maxY != this.maxY
				) {
					repeats++

					copy = if (this.maxX - this.minX >= width / 2 || this.maxY - this.minY >= height / 2) {
						passthrough
					} else {
						Copy()
					}
				} else if (!changes)
					break
			}
		}
	}

	var quality = Quality.HIGH

	// light values below this are considered too small to bother with
	var epsilon = 0.01f

	// values below are specified by lighting.config

	// Number of ambient spread passes.  Needs to be at least spreadMaxAir /
	// spreadMaxObstacle big, but sometimes it can stand to be a bit less and
	// you won't notice.
	var passes = 4

	// Maximum distance through empty space that 100% ambient light can pass through
	var maxAirSpread = 32f
		set(value) {
			field = value
			invMaxAirSpread = 1f / field
		}

	// Maximum distance through rock that 100% ambient light can pass through
	var maxObstacleSpread = 8f
		set(value) {
			field = value
			invMaxObstacleSpread = 1f / field
		}

	private var invMaxAirSpread = 1f / maxAirSpread
	private var invMaxObstacleSpread = 1f / maxObstacleSpread

	// Maximum distance through emtpy space that 100% point light can pass through
	var pointMaxAir = 48f
	// Maximum distance through rock that 100% point light can pass through
	var pointMaxObstacle = 9f

	private data class Cell(override var red: Float = 0f, override var green: Float = 0f, override var blue: Float = 0f) : ICell

	private var mainGrid = lazy {
		Object2DArray.nulls<Cell>(width, height)
	}

	private object Empty : ICell {
		override val red: Float = 0f
		override val green: Float = 0f
		override val blue: Float = 0f
		override val alpha: Float = 0f
	}

	private object Fullbright : ICell {
		override val red: Float = 1f
		override val green: Float = 1f
		override val blue: Float = 1f
		override val alpha: Float = 1f
	}

	private class PointLight(val x: Int, val y: Int, val red: Float, val green: Float, val blue: Float) {
		var assignedTo: TaskCluster? = null

		fun error(to: TaskCluster): Double {
			val dx = x - to.x
			val dy = y - to.y
			return dx * dx + dy * dy
		}
	}

	private val pointLights = ArrayList<PointLight>()

	fun addPointLight(x: Int, y: Int, red: Float, green: Float, blue: Float) {
		if (x !in 0 until width || y !in 0 until height) return
		pointLights.add(PointLight(x, y, red, green, blue))
	}

	fun addPointLight(x: Int, y: Int, color: RGBAColor) {
		val (red, green, blue) = color
		addPointLight(x, y, red, green, blue)
	}

	operator fun get(x: Int, y: Int): ICell {
		if (fullbright) return Fullbright
		if (!mainGrid.isInitialized()) return Empty
		return mainGrid.value[x, y] ?: Empty
	}

	var multithreaded = false
	var fullbright = false

	fun calculate(targetMem: ByteBuffer? = null, targetWidth: Int = 0, targetHeight: Int = 0) {
		if (pointLights.isEmpty() || fullbright) return

		val initialPos = targetMem?.position() ?: 0
		require(targetWidth >= 0) { "Invalid target width" }
		require(targetHeight >= 0) { "Invalid target height" }

		if (targetMem != null && (targetWidth == 0 || targetHeight == 0)) return

		// perform multithreaded calculation only when it makes sense
		if (multithreaded && pointLights.size > 1) {
			val thread = Thread.currentThread()
			// calculate k-means clusters of point lights
			// to effectively utilize CPU cores
			val clusterCount = Starbound.EXECUTOR.parallelism.coerceAtMost(pointLights.size)
			val clusters = ArrayList<TaskCluster>(clusterCount)
			val startingPoints = IntArraySet()
			// val rand = Random(System.nanoTime())

			while (startingPoints.size < clusterCount) {
				// startingPoints.add(rand.nextInt(0, pointLights.size))
				startingPoints.add(startingPoints.size)
			}

			for (index in startingPoints.intIterator()) {
				clusters.add(TaskCluster(pointLights[index].x.toDouble(), pointLights[index].y.toDouble(), this))
			}

			var converged = false

			while (!converged) {
				converged = true

				// assign
				for (light in pointLights) {
					val oldCluster = light.assignedTo

					// do selection sort since it is faster here
					for (cluster in clusters) {
						if (cluster === light.assignedTo) continue
						val old = light.assignedTo?.let { light.error(it) } ?: Double.MAX_VALUE
						val new = light.error(cluster)

						if (new < old) {
							light.assignedTo = cluster
						}
					}

					if (oldCluster != light.assignedTo) {
						oldCluster?.lights?.remove(light)
						light.assignedTo!!.lights.add(light)
						converged = false
					}
				}

				// update center of mass
				for (cluster in clusters) {
					cluster.updateCenter()
				}

				// we settled on clusters
				// check their centres of mass, and probably
				// merge clusters which are too close to each other,
				// to avoid excess work
				// if we merge something, re-run k-means algorithm
				if (converged) {
					for (cluster1 in clusters) {
						for (cluster2 in clusters) {
							if (cluster1 === cluster2) continue
							// don't create big clusters
							if (cluster1.lights.size >= 4 && cluster2.lights.size >= 4) continue

							val dx = cluster1.x - cluster2.x
							val dy = cluster1.y - cluster2.y

							val distance = dx * dx + dy * dy

							if (distance <= 64.0) {
								cluster1.lights.addAll(cluster2.lights)
								clusters.remove(cluster2)
								converged = false
								break
							}
						}

						if (!converged) break
					}
				}
			}

			val tasks = ArrayList<Future<Grid>>()
			clusters.forEach { tasks.add(CompletableFuture.supplyAsync(it, Starbound.EXECUTOR).also { it.thenApply { LockSupport.unpark(thread); it } }) }

			while (tasks.isNotEmpty()) {
				tasks.removeIf {
					if (it.isDone) {
						val grid = it.get()

						if (targetMem != null) {
							for (x in grid.minX - 1 .. grid.maxX) {
								for (y in grid.minY - 1 .. grid.maxY) {
									if (x !in 0 until targetWidth || y !in 0 until targetHeight) continue
									val a = grid.get0(x, y) ?: continue
									val pos = initialPos + (targetWidth * y + x) * 3
									targetMem.position(pos)

									var bRed = (targetMem.get().toInt() and 0xFF) / 255f
									var bGreen = (targetMem.get().toInt() and 0xFF) / 255f
									var bBlue = (targetMem.get().toInt() and 0xFF) / 255f

									bRed = a.red.coerceAtLeast(bRed)
									bGreen = a.green.coerceAtLeast(bGreen)
									bBlue = a.blue.coerceAtLeast(bBlue)

									targetMem.position(pos)

									targetMem.put((bRed * 255f).coerceIn(0f, 255f).roundToInt().toByte())
									targetMem.put((bGreen * 255f).coerceIn(0f, 255f).roundToInt().toByte())
									targetMem.put((bBlue * 255f).coerceIn(0f, 255f).roundToInt().toByte())
								}
							}
						} else {
							val mainGrid = mainGrid.value

							for (x in grid.minX - 1 .. grid.maxX) {
								for (y in grid.minY - 1 .. grid.maxY) {
									val a = grid.get0(x, y) ?: continue
									var b = mainGrid[x, y]

									if (b == null) {
										b = Cell()
										mainGrid[x, y] = b
									}

									b.red = a.red.coerceAtLeast(b.red)
									b.green = a.green.coerceAtLeast(b.green)
									b.blue = a.blue.coerceAtLeast(b.blue)
								}
							}
						}
					}

					it.isDone
				}

				LockSupport.parkNanos(500_000L)
			}
		} else {
			val grid = Grid()

			for (light in pointLights) {
				val cell = grid.safeGet(light.x, light.y) ?: continue
				val speculatedSpread = (maxAirSpread * light.red.coerceAtLeast(light.green).coerceAtLeast(light.blue)).roundToInt()

				cell.red = light.red
				cell.green = light.green
				cell.blue = light.blue

				grid.minX = grid.minX.coerceAtMost(light.x - speculatedSpread)
				grid.minY = grid.minY.coerceAtMost(light.y - speculatedSpread)

				grid.maxX = grid.maxX.coerceAtLeast(light.x + speculatedSpread)
				grid.maxY = grid.maxY.coerceAtLeast(light.y + speculatedSpread)

				grid.clampRect()
			}

			grid.calculateSpread()

			if (targetMem != null) {
				for (x in grid.minX - 1 .. grid.maxX) {
					for (y in grid.minY - 1 .. grid.maxY) {
						if (x !in 0 until targetWidth || y !in 0 until targetHeight) continue
						val a = grid.get0(x, y) ?: continue
						val (bRed, bGreen, bBlue) = a

						targetMem.position(initialPos + (targetWidth * y + x) * 3)
						targetMem.put((bRed * 255f).coerceIn(0f, 255f).roundToInt().toByte())
						targetMem.put((bGreen * 255f).coerceIn(0f, 255f).roundToInt().toByte())
						targetMem.put((bBlue * 255f).coerceIn(0f, 255f).roundToInt().toByte())
					}
				}
			} else {
				val mainGrid = mainGrid.value

				for (x in grid.minX - 1 .. grid.maxX) {
					for (y in grid.minY - 1 .. grid.maxY) {
						val a = grid.get0(x, y) ?: continue

						mainGrid[x, y] = Cell(
							a.red,
							a.green,
							a.blue,
						)
					}
				}
			}
		}
	}

	fun clear() {
		pointLights.clear()
		mainGrid = lazy { Object2DArray.nulls(width, height) }
	}

	private class TaskCluster(
		var x: Double,
		var y: Double,
		val parent: LightCalculator,
	) : Supplier<Grid> {
		val lights = ArrayList<PointLight>()

		fun updateCenter() {
			if (lights.isEmpty()) return

			x = lights.first().x.toDouble()
			y = lights.first().y.toDouble()

			for (i in 1 until lights.size) {
				x += lights[i].x.toDouble()
				y += lights[i].y.toDouble()
			}

			x /= lights.size
			y /= lights.size
		}

		override fun get(): Grid {
			val grid = parent.Grid()

			for (light in lights) {
				val speculatedSpread = (parent.maxAirSpread * light.red.coerceAtLeast(light.green).coerceAtLeast(light.blue)).roundToInt()
				val cell = grid.safeGet(light.x, light.y) ?: continue

				cell.red = light.red
				cell.green = light.green
				cell.blue = light.blue

				grid.minX = grid.minX.coerceAtMost(light.x - speculatedSpread)
				grid.minY = grid.minY.coerceAtMost(light.y - speculatedSpread)

				grid.maxX = grid.maxX.coerceAtLeast(light.x + speculatedSpread)
				grid.maxY = grid.maxY.coerceAtLeast(light.y + speculatedSpread)

				grid.clampRect()
			}

			grid.calculateSpread()
			return grid
		}
	}
}
