package ru.dbotthepony.mc.otm.data

import com.google.gson.JsonObject
import com.google.gson.JsonParseException
import com.google.gson.JsonSyntaxException
import com.mojang.datafixers.util.Pair
import com.mojang.serialization.Codec
import com.mojang.serialization.DataResult
import com.mojang.serialization.DynamicOps
import com.mojang.serialization.JsonOps
import io.netty.buffer.Unpooled
import net.minecraft.data.recipes.FinishedRecipe
import net.minecraft.network.FriendlyByteBuf
import net.minecraft.resources.ResourceLocation
import net.minecraft.world.item.ItemStack
import net.minecraft.world.item.crafting.Ingredient
import net.minecraft.world.item.crafting.Recipe
import net.minecraft.world.item.crafting.RecipeSerializer
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.mc.otm.core.set
import ru.dbotthepony.mc.otm.core.util.readBinaryJsonWithCodecIndirect
import ru.dbotthepony.mc.otm.core.util.writeBinaryJsonWithCodec
import kotlin.collections.ArrayDeque
import kotlin.concurrent.getOrSet

class Codec2RecipeSerializer<S : Recipe<*>>(
	val empty: S?,
	codec: (Codec2RecipeSerializer<S>.Context) -> Codec<S>,
) : Codec<S>, RecipeSerializer<S> {
	constructor(supplier: (Codec2RecipeSerializer<S>.Context) -> Codec<S>) : this(null, supplier)

	private class CurrentContext {
		val idStack = ArrayDeque<ResourceLocation>()
		var isNetwork = 0
	}

	inner class Context {
		val id: ResourceLocation
			get() = checkNotNull(context.idStack.lastOrNull()) { "Not currently deserializing recipe" }

		val ingredients: Codec<Ingredient> get() = ActualIngredientCodec

		val isNetwork: Boolean
			get() = context.isNetwork > 0
	}

	private val codec = codec.invoke(Context())

	override fun <T : Any> encode(input: S, ops: DynamicOps<T>, prefix: T): DataResult<T> {
		return codec.encode(input, ops, prefix)
	}

	override fun <T : Any> decode(ops: DynamicOps<T>, input: T): DataResult<Pair<S, T>> {
		return codec.decode(ops, input)
	}

	fun <O : Recipe<*>> xmap(to: (S) -> O, from: (O) -> S): Codec2RecipeSerializer<O> {
		return Codec2RecipeSerializer(empty?.let(to)) { _ ->
			codec.xmap(to, from)
		}
	}

	override fun fromJson(id: ResourceLocation, data: JsonObject): S {
		try {
			context.idStack.addLast(id)

			return decode(JsonOps.INSTANCE, data).get().map(
				{ it.first },
				{ empty ?: throw JsonSyntaxException("Failed to deserialize recipe from JSON: ${it.message()}") }
			)
		} finally {
			context.idStack.removeLast()
		}
	}

	override fun fromNetwork(id: ResourceLocation, data: FriendlyByteBuf): S? {
		try {
			context.idStack.addLast(id)
			context.isNetwork++

			return data.readBinaryJsonWithCodecIndirect(this)
				.resultOrPartial { LOGGER.error("Failed to read recipe $id from network: $it") }.orElse(null)
		} finally {
			context.isNetwork--
			context.idStack.removeLast()
		}
	}

	override fun toNetwork(data: FriendlyByteBuf, recipe: S) {
		try {
			context.idStack.addLast(recipe.id)
			context.isNetwork++

			data.writeBinaryJsonWithCodec(this, recipe)
		} finally {
			context.isNetwork--
			context.idStack.removeLast()
		}
	}

	fun toFinished(recipe: S): FinishedRecipe {
		return object : FinishedRecipe {
			override fun serializeRecipeData(p_125967_: JsonObject) {
				encode(recipe, JsonOps.INSTANCE, p_125967_).get().map(
					{
						it as JsonObject

						for ((k, v) in it.entrySet()) {
							p_125967_[k] = v
						}
					},
					{
						throw JsonParseException("Failed to serialize recipe: ${it.message()}")
					}
				)
			}

			override fun getId(): ResourceLocation {
				return recipe.id
			}

			override fun getType(): RecipeSerializer<*> {
				return this@Codec2RecipeSerializer
			}

			override fun serializeAdvancement(): JsonObject? {
				return null
			}

			override fun getAdvancementId(): ResourceLocation? {
				return null
			}
		}
	}

	private object ActualIngredientCodec : Codec<Ingredient> {
		override fun <T : Any> encode(input: Ingredient, ops: DynamicOps<T>, prefix: T): DataResult<T> {
			return if (context.isNetwork > 0) {
				networkIngredientCodec.encode(input, ops, prefix)
			} else {
				IngredientCodec.encode(input, ops, prefix)
			}
		}

		override fun <T : Any> decode(ops: DynamicOps<T>, input: T): DataResult<Pair<Ingredient, T>> {
			return if (context.isNetwork > 0) {
				networkIngredientCodec.decode(ops, input)
			} else {
				IngredientCodec.decode(ops, input)
			}
		}
	}

	companion object {
		private val LOGGER = LogManager.getLogger()
		private val networkIngredientCodec = Codec.list(ItemStack.CODEC).xmap({ Ingredient.of(it.stream()) }, { it.items.toMutableList() })

		/**
		 * [ThreadLocal] because optimization mods can (and probably should) parallelize recipe deserialization,
		 * since RecipeSerializers are expected to be stateless. [Codec2RecipeSerializer], however, is stateful (threading PoV).
		 * To make it stateless, [ThreadLocal] is used.
		 */
		private val contextHolder = ThreadLocal<CurrentContext>()
		private val context: CurrentContext
			get() = contextHolder.getOrSet { CurrentContext() }
	}
}
