package ru.dbotthepony.mc.otm.core.util

import com.google.common.collect.ImmutableList
import com.mojang.datafixers.util.Pair
import com.mojang.serialization.Codec
import com.mojang.serialization.DataResult
import com.mojang.serialization.DynamicOps
import net.minecraft.nbt.NbtAccounter
import net.minecraft.world.item.ItemStack
import net.minecraftforge.fluids.FluidStack
import ru.dbotthepony.mc.otm.core.immutableMap
import ru.dbotthepony.mc.otm.core.math.RGBAColor
import ru.dbotthepony.mc.otm.core.math.readDecimal
import ru.dbotthepony.mc.otm.core.math.writeDecimal
import ru.dbotthepony.mc.otm.core.readItemType
import ru.dbotthepony.mc.otm.core.writeItemType
import java.io.DataInput
import java.io.DataInputStream
import java.io.DataOutput
import java.io.DataOutputStream
import java.io.InputStream
import java.io.OutputStream
import java.util.*
import java.util.function.Predicate
import kotlin.NoSuchElementException
import kotlin.math.absoluteValue
import kotlin.reflect.KClass

/**
 * Represents value which can be encoded onto or decoded from stream.
 *
 * Also provides [copy] and [compare] methods
 */
interface IStreamCodec<V> {
	fun read(stream: DataInputStream): V
	fun write(stream: DataOutputStream, value: V)

	/**
	 * if value is immutable, return it right away
	 */
	fun copy(value: V): V

	/**
	 * Optional equality check override. Utilized to determine whenever e.g. network value is different from new value
	 *
	 * By default uses [Any.equals]
	 */
	fun compare(a: V, b: V): Boolean {
		return a == b
	}
}

class StreamCodec<V>(
	private val reader: (stream: DataInputStream) -> V,
	private val writer: (stream: DataOutputStream, value: V) -> Unit,
	private val copier: ((value: V) -> V) = { it },
	private val comparator: ((a: V, b: V) -> Boolean) = { a, b -> a == b }
) : IStreamCodec<V> {
	constructor(
		reader: (stream: DataInputStream) -> V,
		payloadSize: Long,
		writer: (stream: DataOutputStream, value: V) -> Unit,
		copier: ((value: V) -> V) = { it },
		comparator: ((a: V, b: V) -> Boolean) = { a, b -> a == b }
	) : this({ stream -> reader.invoke(stream) }, writer, copier, comparator)

	val nullable = object : IStreamCodec<V?> {
		override fun read(stream: DataInputStream): V? {
			return if (stream.read() == 0) null else reader.invoke(stream)
		}

		override fun write(stream: DataOutputStream, value: V?) {
			if (value === null) stream.write(0) else {
				stream.write(1)
				writer.invoke(stream, value)
			}
		}

		override fun copy(value: V?): V? {
			return if (value === null) null else copier.invoke(value)
		}

		override fun compare(a: V?, b: V?): Boolean {
			if (a === null && b === null) return true
			if (a === null || b === null) return false
			return comparator.invoke(a, b)
		}
	}

	override fun read(stream: DataInputStream): V {
		return reader.invoke(stream)
	}

	override fun write(stream: DataOutputStream, value: V) {
		writer.invoke(stream, value)
	}

	override fun copy(value: V): V {
		return copier.invoke(value)
	}

	override fun compare(a: V, b: V): Boolean {
		return comparator.invoke(a, b)
	}
}

class CollectionStreamCodec<E, C : MutableCollection<E>>(val elementCodec: IStreamCodec<E>, val collectionFactory: (Int) -> C) : IStreamCodec<C> {
	override fun read(stream: DataInputStream): C {
		val size = stream.readVarIntLE()

		if (size <= 0) {
			return collectionFactory.invoke(0)
		}

		val collection = collectionFactory.invoke(size)

		for (i in 0 until size) {
			collection.add(elementCodec.read(stream))
		}

		return collection
	}

	override fun write(stream: DataOutputStream, value: C) {
		stream.writeVarIntLE(value.size)
		value.forEach { elementCodec.write(stream, it) }
	}

	override fun copy(value: C): C {
		val new = collectionFactory.invoke(value.size)
		value.forEach { new.add(elementCodec.copy(it)) }
		return new
	}
}

val NullValueCodec = StreamCodec({ _ -> null }, { _, _ -> })
val BooleanValueCodec = StreamCodec(DataInputStream::readBoolean, 1L, DataOutputStream::writeBoolean) { a, b -> a == b }
val ByteValueCodec = StreamCodec(DataInputStream::readByte, 1L, { s, v -> s.writeByte(v.toInt()) }) { a, b -> a == b }
val ShortValueCodec = StreamCodec(DataInputStream::readShort, 2L, { s, v -> s.writeShort(v.toInt()) }) { a, b -> a == b }
val IntValueCodec = StreamCodec(DataInputStream::readInt, 4L, DataOutputStream::writeInt) { a, b -> a == b }
val LongValueCodec = StreamCodec(DataInputStream::readLong, 8L, DataOutputStream::writeLong) { a, b -> a == b }
val FloatValueCodec = StreamCodec(DataInputStream::readFloat, 4L, DataOutputStream::writeFloat) { a, b -> a == b }
val DoubleValueCodec = StreamCodec(DataInputStream::readDouble, 8L, DataOutputStream::writeDouble) { a, b -> a == b }
val ItemStackValueCodec = StreamCodec(DataInputStream::readItem, DataOutputStream::writeItem, ItemStack::copy) { a, b -> a.equals(b, true) }
val FluidStackValueCodec = StreamCodec(DataInputStream::readFluidStack, DataOutputStream::writeFluidStack, FluidStack::copy) { a, b -> a == b && a.amount == b.amount }
val ItemValueCodec = StreamCodec(DataInputStream::readItemType, DataOutputStream::writeItemType) { a, b -> a === b }
val DecimalValueCodec = StreamCodec(DataInputStream::readDecimal, DataOutputStream::writeDecimal)
val BigDecimalValueCodec = StreamCodec(DataInputStream::readBigDecimal, DataOutputStream::writeBigDecimal)
val UUIDValueCodec = StreamCodec({ s -> UUID(s.readLong(), s.readLong()) }, { s, v -> s.writeLong(v.mostSignificantBits); s.writeLong(v.leastSignificantBits) })
val VarIntValueCodec = StreamCodec(DataInputStream::readVarIntLE, DataOutputStream::writeVarIntLE) { a, b -> a == b }
val VarLongValueCodec = StreamCodec(DataInputStream::readVarLongLE, DataOutputStream::writeVarLongLE) { a, b -> a == b }
val BinaryStringCodec = StreamCodec(DataInputStream::readBinaryString, DataOutputStream::writeBinaryString)
val ResourceLocationValueCodec = StreamCodec(DataInputStream::readResourceLocation, DataOutputStream::writeResourceLocation)

val RGBCodec: StreamCodec<RGBAColor> = StreamCodec(
	{ s -> RGBAColor(s.readFloat(), s.readFloat(), s.readFloat()) },
	{ s, v -> s.writeFloat(v.red); s.writeFloat(v.green); s.writeFloat(v.blue) })

val RGBACodec: StreamCodec<RGBAColor> = StreamCodec(
	{ s -> RGBAColor(s.readFloat(), s.readFloat(), s.readFloat(), s.readFloat()) },
	{ s, v -> s.writeFloat(v.red); s.writeFloat(v.green); s.writeFloat(v.blue); s.writeFloat(v.alpha) })

class EnumValueCodec<V : Enum<V>>(clazz: Class<out V>) : IStreamCodec<V>, Codec<V> {
	val clazz = searchClass(clazz)
	val values: ImmutableList<V> = ImmutableList.copyOf(this.clazz.enumConstants!!)
	val valuesMap = immutableMap {
		for (v in values) put(v.name, v)
	}

	override fun read(stream: DataInputStream): V {
		val id = stream.readVarIntLE()
		return values.getOrNull(id) ?: throw NoSuchElementException("No such enum with index $id")
	}

	override fun write(stream: DataOutputStream, value: V) {
		stream.writeVarIntLE(value.ordinal)
	}

	override fun copy(value: V): V {
		return value
	}

	override fun compare(a: V, b: V): Boolean {
		return a === b
	}

	override fun <T : Any> encode(input: V, ops: DynamicOps<T>, prefix: T): DataResult<T> {
		if (ops.compressMaps()) {
			return DataResult.success(ops.createInt(input.ordinal))
		}

		return DataResult.success(ops.createString(input.name))
	}

	override fun <T : Any> decode(ops: DynamicOps<T>, input: T): DataResult<Pair<V, T>> {
		if (ops.compressMaps()) {
			return ops.getNumberValue(input)
				.flatMap { values.getOrNull(it.toInt())?.let { DataResult.success(Pair(it, ops.empty())) } ?: DataResult.error("No such enum with ordinal index $it") }
		}

		return ops.getStringValue(input)
			.flatMap { valuesMap[it]?.let { DataResult.success(Pair(it, ops.empty())) } ?: DataResult.error("No such enum value $it") }
	}

	companion object {
		/**
		 * FIXME: enums with abstract methods which get compiled to subclasses, whose DO NOT expose "parent's" enum constants array
		 *
		 * is there an already existing solution?
		 */
		fun <V : Enum<V>> searchClass(clazz: Class<out V>): Class<out V> {
			var search: Class<*> = clazz

			while (search.enumConstants == null && search.superclass != null) {
				search = search.superclass
			}

			if (search.enumConstants == null) {
				throw ClassCastException("$clazz does not represent an enum or enum subclass")
			}

			return search as Class<out V>
		}
	}
}

fun <E : Enum<E>> Class<E>.codec() = EnumValueCodec(this)
fun <E : Enum<E>> KClass<E>.codec() = EnumValueCodec(this.java)
