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, sizeLimit: NbtAccounter = NbtAccounter(1L shl 18 /* 256 KiB */)): 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, sizeLimit: NbtAccounter) -> 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, sizeLimit -> sizeLimit.accountBytes(payloadSize); reader.invoke(stream) }, writer, copier, comparator)

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

		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, sizeLimit: NbtAccounter): V {
		return reader.invoke(stream, sizeLimit)
	}

	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)
	}
}

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, a -> a.accountBytes(8L); 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 RGBCodec: StreamCodec<RGBAColor> = StreamCodec(
	{ s, a -> a.accountBytes(12L); 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, a -> a.accountBytes(16L); 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, sizeLimit: NbtAccounter): V {
		val id = stream.readVarIntLE(sizeLimit)
		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)

fun OutputStream.writeInt(value: Int) {
	if (this is DataOutput) {
		writeInt(value)
		return
	}

	write(value ushr 24)
	write(value ushr 16)
	write(value ushr 8)
	write(value)
}

fun InputStream.readInt(): Int {
	if (this is DataInput) {
		return readInt()
	}

	return (read() shl 24) or (read() shl 16) or (read() shl 8) or read()
}

fun OutputStream.writeLong(value: Long) {
	if (this is DataOutput) {
		writeLong(value)
		return
	}

	write((value ushr 48).toInt())
	write((value ushr 40).toInt())
	write((value ushr 32).toInt())
	write((value ushr 24).toInt())
	write((value ushr 16).toInt())
	write((value ushr 8).toInt())
	write(value.toInt())
}

fun InputStream.readLong(): Long {
	if (this is DataInput) {
		return readLong()
	}

	return (read().toLong() shl 48) or
		(read().toLong() shl 40) or
		(read().toLong() shl 32) or
		(read().toLong() shl 24) or
		(read().toLong() shl 16) or
		(read().toLong() shl 8) or
		read().toLong()
}

fun OutputStream.writeFloat(value: Float) = writeInt(value.toBits())
fun InputStream.readFloat() = Float.fromBits(readInt())
fun OutputStream.writeDouble(value: Double) = writeLong(value.toBits())
fun InputStream.readDouble() = Double.fromBits(readLong())

fun InputStream.readVarIntLE(sizeLimit: NbtAccounter? = null): Int {
	sizeLimit?.accountBytes(1L)
	val readFirst = read()

	if (readFirst < 0) {
		throw NoSuchElementException("Reached end of stream")
	}

	if (readFirst and 64 == 0) {
		return if (readFirst and 128 != 0) -(readFirst and 63) else readFirst and 63
	}

	var result = 0
	var nextBit = readFirst and 64
	var read = readFirst and 63
	var i = 0

	while (nextBit != 0) {
		result = result or (read shl i)
		sizeLimit?.accountBytes(1L)
		read = read()

		if (read < 0) {
			throw NoSuchElementException("Reached end of stream")
		}

		nextBit = read and 128
		read = read and 127

		if (i == 0)
			i = 6
		else
			i += 7
	}

	result = result or (read shl i)
	return if (readFirst and 128 != 0) -result else result
}

fun OutputStream.writeVarIntLE(value: Int) {
	write((if (value < 0) 128 else 0) or (if (value in -63 .. 63) 0 else 64) or (value.absoluteValue and 63))
	var written = value.absoluteValue ushr 6

	while (written != 0) {
		write((written and 127) or (if (written >= 128) 128 else 0))
		written = written ushr 7
	}
}

fun InputStream.readVarLongLE(sizeLimit: NbtAccounter? = null): Long {
	sizeLimit?.accountBytes(1L)

	val readFirst = read()

	if (readFirst < 0) {
		throw NoSuchElementException("Reached end of stream")
	}

	if (readFirst and 64 == 0) {
		return if (readFirst and 128 != 0) -(readFirst and 63).toLong() else (readFirst and 63).toLong()
	}

	var result = 0L
	var nextBit = readFirst and 64
	var read = readFirst and 63
	var i = 0

	while (nextBit != 0) {
		result = result or (read shl i).toLong()
		sizeLimit?.accountBytes(1L)
		read = read()

		if (read < 0) {
			throw NoSuchElementException("Reached end of stream")
		}

		nextBit = read and 128
		read = read and 127

		if (i == 0)
			i = 6
		else
			i += 7
	}

	result = result or (read shl i).toLong()
	return if (readFirst and 128 != 0) -result else result
}

fun OutputStream.writeVarLongLE(value: Long) {
	write((if (value < 0L) 128 else 0) or (if (value in -63 .. 63) 0 else 64) or (value.absoluteValue and 63).toInt())
	var written = value.absoluteValue ushr 6

	while (written != 0L) {
		write((written and 127).toInt() or (if (written >= 128) 128 else 0))
		written = written ushr 7
	}
}

fun InputStream.readBinaryString(sizeLimit: NbtAccounter = NbtAccounter(1L shl 18 /* 256 KiB */)): String {
	val size = readVarIntLE()
	require(size >= 0) { "Negative payload size: $size" }
	sizeLimit.accountBytes(size.toLong())
	val bytes = ByteArray(size)
	read(bytes)
	return bytes.decodeToString()
}

fun OutputStream.writeBinaryString(input: String) {
	val bytes = input.encodeToByteArray()
	writeVarIntLE(bytes.size)
	write(bytes)
}

private data class IndexedStreamCodec<T>(
	val condition: Predicate<Any?>,
	val id: Int,
	val codec: IStreamCodec<T>
) {
	fun read(stream: DataInputStream, sizeLimit: NbtAccounter = NbtAccounter(1L shl 18 /* 256 KiB */)): T {
		return codec.read(stream, sizeLimit)
	}

	fun write(stream: DataOutputStream, value: Any?) {
		codec.write(stream, value as T)
	}
}

private inline fun <reified T> IndexedStreamCodec(id: Int, codec: StreamCodec<T>): IndexedStreamCodec<T> {
	return IndexedStreamCodec({ it is T }, id, codec)
}

private val codecs: List<IndexedStreamCodec<*>> = com.google.common.collect.ImmutableList.builder<IndexedStreamCodec<*>>().let {
	var codecID = 0

	it.add(IndexedStreamCodec({ it == null }, codecID++, NullValueCodec))

	it.add(IndexedStreamCodec(codecID++, BooleanValueCodec))
	it.add(IndexedStreamCodec(codecID++, ByteValueCodec))
	it.add(IndexedStreamCodec(codecID++, ShortValueCodec))
	it.add(IndexedStreamCodec(codecID++, VarIntValueCodec))
	it.add(IndexedStreamCodec(codecID++, VarLongValueCodec))
	it.add(IndexedStreamCodec(codecID++, FloatValueCodec))
	it.add(IndexedStreamCodec(codecID++, DoubleValueCodec))
	it.add(IndexedStreamCodec(codecID++, ItemStackValueCodec))
	it.add(IndexedStreamCodec(codecID++, DecimalValueCodec))
	it.add(IndexedStreamCodec(codecID++, BigDecimalValueCodec))
	it.add(IndexedStreamCodec(codecID++, UUIDValueCodec))
	it.add(IndexedStreamCodec(codecID++, BinaryStringCodec))

	it.build()
}

class NotSerializableValueException(message: String) : Exception(message)

/**
 * Write arbitrary data to this stream, in exploit-free way
 */
fun DataOutputStream.writeType(value: Any?) {
	for (codec in codecs) {
		if (codec.condition.test(value)) {
			write(codec.id)
			codec.write(this, value)
			return
		}
	}

	throw NotSerializableValueException("Value $value <${value?.let { it::class.qualifiedName }}> can not be serialized")
}

/**
 * Read arbitrary data from this stream, in exploit-free way
 */
fun DataInputStream.readType(sizeLimit: NbtAccounter = NbtAccounter(1L shl 18 /* 256 KiB */)): Any? {
	sizeLimit.accountBytes(1L)
	val id = read()

	if (id >= codecs.size) {
		throw IndexOutOfBoundsException("No codec for network type $id")
	}

	return codecs[id].read(this, sizeLimit)
}
