package ru.dbotthepony.kstarbound.defs.item

import com.google.common.collect.ImmutableList
import com.google.gson.Gson
import com.google.gson.JsonArray
import com.google.gson.JsonObject
import com.google.gson.JsonPrimitive
import com.google.gson.JsonSyntaxException
import com.google.gson.TypeAdapter
import com.google.gson.TypeAdapterFactory
import com.google.gson.annotations.JsonAdapter
import com.google.gson.internal.bind.JsonTreeReader
import com.google.gson.reflect.TypeToken
import com.google.gson.stream.JsonReader
import com.google.gson.stream.JsonWriter
import it.unimi.dsi.fastutil.objects.Object2IntArrayMap
import it.unimi.dsi.fastutil.objects.Object2IntMap
import ru.dbotthepony.kommons.util.Either
import ru.dbotthepony.kstarbound.Registry
import ru.dbotthepony.kommons.gson.consumeNull
import ru.dbotthepony.kommons.gson.get
import ru.dbotthepony.kstarbound.math.vector.Vector2d
import ru.dbotthepony.kstarbound.Starbound
import ru.dbotthepony.kstarbound.collect.WeightedList
import ru.dbotthepony.kstarbound.item.ItemStack
import ru.dbotthepony.kstarbound.json.jsonArrayOf
import ru.dbotthepony.kstarbound.json.stream
import ru.dbotthepony.kstarbound.util.WriteOnce
import ru.dbotthepony.kstarbound.util.random.nextRange
import ru.dbotthepony.kstarbound.util.random.random
import java.util.random.RandomGenerator

typealias ItemOrPool = Either<ItemDescriptor, Registry.Ref<TreasurePoolDefinition>>

@JsonAdapter(TreasurePoolDefinition.Adapter::class)
class TreasurePoolDefinition(pieces: List<Piece>) {
	var name: String by WriteOnce()

	init {
		require(pieces.isNotEmpty()) { "Treasure pool is empty" }
	}

	val pieces: ImmutableList<Piece> = pieces
		.stream()
		.sorted { o1, o2 -> o1.startingLevel.compareTo(o2.startingLevel) }
		.collect(ImmutableList.toImmutableList())

	fun evaluate(random: RandomGenerator, level: Double, visitedPools: Object2IntMap<String> = Object2IntArrayMap()): List<ItemStack> {
		require(level >= 0.0) { "Invalid loot level: $level" }

		val new = visitedPools.getInt(name) + 1

		if (new >= 50)
			throw IllegalStateException("Too deep treasure pool references, visited treasure pool $name 50 times (visited: $visitedPools)")

		visitedPools.put(name, new)

		try {
			if (pieces.isNotEmpty()) {
				return pieces.last { it.startingLevel <= level }.evaluate(random, level, visitedPools)
			}

			return emptyList()
		} finally {
			visitedPools.put(name, visitedPools.getInt(name) - 1)
		}
	}

	fun evaluate(seed: Long, level: Double): List<ItemStack> {
		return evaluate(random(seed), level)
	}

	data class Piece(
		val startingLevel: Double,
		val pool: WeightedList<ItemOrPool> = WeightedList(),
		val fill: ImmutableList<ItemOrPool> = ImmutableList.of(),
		val poolRounds: IPoolRounds = OneRound,
		val allowDuplication: Boolean = true,
		val levelVariance: Vector2d = Vector2d.ZERO
	) {
		fun evaluate(random: RandomGenerator, level: Double, visitedPools: Object2IntMap<String>): List<ItemStack> {
			val result = ArrayList<ItemStack>()
			val previousDescriptors = HashSet<ItemDescriptor>()

			for (entry in fill) {
				entry.map({
			          val stack = it.build(level = level + random.nextRange(levelVariance), random = random, seed = random.nextLong())

			          if (stack.isNotEmpty && (allowDuplication || previousDescriptors.add(stack.createDescriptor().copy(count = 1L)))) {
				          result.add(stack)
			          }
		          }, {
			          it.value?.evaluate(random, level + random.nextRange(levelVariance), visitedPools)?.forEach {
				          if (allowDuplication || previousDescriptors.add(it.createDescriptor().copy(count = 1L)))
					          result.add(it)
			          }
		          })
			}

			if (pool.isNotEmpty) {
				for (round in 0 until poolRounds.evaluate(random)) {
					pool.sample(random).orThrow { RuntimeException() }.map({
						val stack = it.build(level = level + random.nextRange(levelVariance), random = random, seed = random.nextLong())

						if (stack.isNotEmpty && (allowDuplication || previousDescriptors.add(stack.createDescriptor().copy(count = 1L)))) {
							result.add(stack)
						}
					}, {
						it.value?.evaluate(random, level + random.nextRange(levelVariance), visitedPools)?.forEach {
							if (allowDuplication || previousDescriptors.add(it.createDescriptor().copy(count = 1L)))
								result.add(it)
						}
					})
				}
			}

			return result
		}
	}

	interface IPoolRounds {
		fun evaluate(random: RandomGenerator): Int
	}

	object OneRound : IPoolRounds {
		override fun evaluate(random: RandomGenerator): Int {
			return 1
		}
	}

	data class ConstantRound(val amount: Int) : IPoolRounds {
		init {
			check(amount >= 1) { "Invalid treasure pool rounds: $amount" }
		}

		override fun evaluate(random: RandomGenerator): Int {
			return amount
		}
	}

	data class WeightedRounds(val edges: WeightedList<Int>) : IPoolRounds {
		override fun evaluate(random: RandomGenerator): Int {
			return edges.sample(random).orElse(1)
		}
	}

	class Adapter(gson: Gson) : TypeAdapter<TreasurePoolDefinition>() {
		private val itemAdapter = gson.getAdapter(ItemDescriptor::class.java)
		private val poolAdapter = gson.getAdapter(object : TypeToken<Registry.Ref<TreasurePoolDefinition>>() {})
		private val vectors = gson.getAdapter(Vector2d::class.java)

		override fun write(out: JsonWriter, value: TreasurePoolDefinition?) {
			if (value == null) {
				out.nullValue()
			} else {
				TODO()
			}
		}

		override fun read(`in`: JsonReader): TreasurePoolDefinition? {
			if (`in`.consumeNull()) {
				return null
			}

			val pieces = ArrayList<Piece>()

			`in`.beginArray()

			while (`in`.hasNext()) {
				`in`.beginArray()

				val level = `in`.nextDouble()
				val things = Starbound.ELEMENTS_ADAPTER.objects.read(`in`)

				val pool = ImmutableList.Builder<Pair<Double, ItemOrPool>>()
				val fill = ImmutableList.Builder<Either<ItemDescriptor, Registry.Ref<TreasurePoolDefinition>>>()
				var poolRounds: IPoolRounds = OneRound
				val allowDuplication = things["allowDuplication"]?.asBoolean ?: true

				things["poolRounds"]?.let {
					if (it is JsonPrimitive) {
						poolRounds = ConstantRound(it.asInt)
					} else if (it is JsonArray) {
						poolRounds = WeightedRounds(WeightedList(Starbound.gson.getAdapter(object : TypeToken<ImmutableList<Pair<Double, Int>>>() {}).fromJsonTree(it)))
					} else {
						throw JsonSyntaxException("Expected either a number or an array for poolRounds, but got ${it::class.simpleName}")
					}
				}

				things["pool"]?.let {
					it as? JsonArray ?: throw JsonSyntaxException("pool must be an array")

					for ((i, elem) in it.withIndex()) {
						elem as? JsonObject ?: throw JsonSyntaxException("Pool element at $i is not an object")
						val weight = (elem["weight"] as? JsonPrimitive)?.asDouble ?: throw JsonSyntaxException("Pool element at $i is missing weight")

						if (elem.has("item")) {
							pool.add(weight to Either.left(itemAdapter.read(JsonTreeReader(elem["item"]!!))))
						} else if (elem.has("pool")) {
							pool.add(weight to Either.right(poolAdapter.read(JsonTreeReader(elem["pool"]!!))))
						} else {
							throw JsonSyntaxException("Pool element at $i is missing both 'item' and 'pool' entries")
						}
					}
				}

				things["fill"]?.let {
					it as? JsonArray ?: throw JsonSyntaxException("fill must be an array")

					for ((i, elem) in it.withIndex()) {
						elem as? JsonObject ?: throw JsonSyntaxException("Fill element at $i is not an object")

						if (elem.has("item")) {
							fill.add(Either.left(itemAdapter.read(JsonTreeReader(elem["item"]!!))))
						} else if (elem.has("pool")) {
							fill.add(Either.right(poolAdapter.read(JsonTreeReader(elem["pool"]!!))))
						} else {
							throw JsonSyntaxException("Fill element at $i is missing both 'item' and 'pool' entries")
						}
					}
				}

				pieces.add(Piece(
					pool = WeightedList(pool.build()),
					fill = fill.build(),
					poolRounds = poolRounds,
					startingLevel = level,
					allowDuplication = allowDuplication,
					levelVariance = vectors.fromJsonTree(things.get("levelVariance", jsonArrayOf(0.0, 0.0)))
				))

				`in`.endArray()
			}

			`in`.endArray()

			return TreasurePoolDefinition(pieces)
		}
	}
}