package dev.latvian.mods.kubejs.recipe.schema;

import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.mojang.serialization.DynamicOps;
import dev.latvian.mods.kubejs.recipe.RecipeKey;
import dev.latvian.mods.kubejs.recipe.schema.function.RecipeFunctionInstance;
import dev.latvian.mods.kubejs.recipe.schema.function.RecipeSchemaFunction;
import dev.latvian.mods.kubejs.recipe.schema.postprocessing.RecipePostProcessor;
import dev.latvian.mods.kubejs.script.ConsoleJS;
import dev.latvian.mods.kubejs.util.Cast;
import dev.latvian.mods.kubejs.util.JsonUtils;
import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap;
import net.minecraft.resources.Identifier;
import net.minecraft.server.packs.resources.ResourceManager;
import org.jspecify.annotations.NullUnmarked;
import org.jspecify.annotations.Nullable;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.SequencedMap;

import static com.mojang.serialization.DataResult.Error;
import static com.mojang.serialization.DataResult.Success;

public class JsonRecipeSchemaLoader {
	@NullUnmarked
	private static final class RecipeSchemaBuilder {
		private final Identifier id;
		private final RecipeSchemaData data;
		private RecipeSchema schema;

		private RecipeSchemaBuilder parent;
		private Identifier overrideType;
		private List<RecipeKey<?>> keys;
		private List<RecipeSchemaData.ConstructorData> constructors;
		private Map<String, RecipeSchemaFunction> functions;
		private KubeRecipeFactory recipeFactory;
		private List<String> unique;
		private Boolean hidden;
		private Map<String, JsonElement> overrideKeys;
		private List<RecipePostProcessor> postProcessors;

		private RecipeSchemaBuilder(Identifier id, RecipeSchemaData data) {
			this.id = id;
			this.data = data;
		}

		private List<RecipeKey<?>> getKeys() {
			if (keys != null) {
				if (data.mergeKeys()) {
					var merged = new LinkedHashMap<String, RecipeKey<?>>();

					if (parent != null) {
						for (var key : parent.getKeys()) {
							merged.put(key.name, key);
						}
					}

					var newOptionals = 0;
					for (var key : keys) {
						if (newOptionals != 0 && !key.optional()) {
							throw new IllegalArgumentException("Required key '%s' must be before optional keys %s".formatted(key.name,
								requiredFirst(merged).stream().skip(merged.size() - newOptionals).map(k -> k.name).toList()));
						}


						var oldKeyOptional = merged.containsKey(key.name) && merged.get(key.name).optional();
						if (key.optional() && !oldKeyOptional) {
							newOptionals++;
						} else if (!key.optional() && oldKeyOptional) {
							throw new IllegalArgumentException("Optional key '%s' from parent may not be replaced by required key!".formatted(key.name));
						}

						merged.put(key.name, key);
					}

					return requiredFirst(merged);
				}

				return keys;
			} else if (parent != null) {
				return parent.getKeys();
			} else {
				return List.of();
			}
		}

		private List<RecipeKey<?>> requiredFirst(SequencedMap<String, RecipeKey<?>> map) {
			var required = new ArrayList<RecipeKey<?>>(map.size());
			var optional = new ArrayList<RecipeKey<?>>(map.size());

			for (RecipeKey<?> key : map.sequencedValues()) {
				(key.optional() ? optional : required).add(key);
			}

			required.addAll(optional);
			return required;
		}

		private List<RecipeSchemaData.ConstructorData> getConstructors() {
			if (constructors != null) {
				if (data.mergeConstructors()) {
					var list = new ArrayList<RecipeSchemaData.ConstructorData>();

					if (parent != null) {
						list.addAll(parent.getConstructors());
					}

					list.addAll(constructors);
					return list;
				}

				return constructors;
			} else if (parent != null) {
				return parent.getConstructors();
			} else {
				return List.of();
			}
		}

		private void gatherFunctions(Map<String, RecipeSchemaFunction> list) {
			if (parent != null) {
				parent.gatherFunctions(list);
			}

			if (functions != null) {
				list.putAll(functions);
			}
		}

		@Nullable
		private KubeRecipeFactory getRecipeFactory() {
			if (recipeFactory != null) {
				return recipeFactory;
			} else if (parent != null) {
				return parent.getRecipeFactory();
			} else {
				return null;
			}
		}

		private List<String> getUnique() {
			if (unique != null) {
				if (data.mergeUnique()) {
					var u = new LinkedHashSet<String>();

					if (parent != null) {
						u.addAll(parent.getUnique());
					}

					u.addAll(unique);
					return List.copyOf(u);
				}

				return unique;
			} else if (parent != null) {
				return parent.getUnique();
			} else {
				return List.of();
			}
		}

		private boolean isHidden() {
			if (hidden != null) {
				return hidden;
			} else if (parent != null) {
				return parent.isHidden();
			} else {
				return false;
			}
		}

		private List<RecipePostProcessor> getPostProcessors() {
			if (postProcessors != null) {
				if (data.mergePostProcessors()) {
					var list = new ArrayList<RecipePostProcessor>();

					if (parent != null) {
						list.addAll(parent.getPostProcessors());
					}

					list.addAll(postProcessors);
					return list;
				}

				return postProcessors;
			} else if (parent != null) {
				return parent.getPostProcessors();
			} else {
				return List.of();
			}
		}

		private RecipeSchema getSchema(DynamicOps<JsonElement> jsonOps) {
			if (schema == null) {
				if (overrideType != null || keys != null || constructors != null || functions != null || recipeFactory != null || unique != null || overrideKeys != null) {
					var keys = getKeys();
					var keyMap = new HashMap<String, RecipeKey<?>>();

					for (var key : keys) {
						keyMap.put(key.name, key);
					}

					var functionMap = new HashMap<String, RecipeSchemaFunction>();
					gatherFunctions(functionMap);

					Map<RecipeKey<?>, RecipeOptional<?>> keyOverrides = Map.of();

					if (overrideKeys != null) {
						keyOverrides = new Reference2ObjectOpenHashMap<>(overrideKeys.size());

						for (var entry : overrideKeys.entrySet()) {
							var key = keyMap.get(entry.getKey());

							if (key != null) {
								try {
									keyOverrides.put(key, RecipeOptional.unit(key.codec.decode(jsonOps, entry.getValue()).getOrThrow().getFirst()));
								} catch (Exception ex) {
									throw new IllegalArgumentException("Failed to create optional value for key '" + key + "' of '" + id + "' from " + entry.getValue(), ex);
								}
							} else {
								throw new NullPointerException("Key '" + entry.getKey() + "' not found in key overrides of recipe schema '" + id + "'");
							}
						}
					}

					var schema = new RecipeSchema(keyOverrides, getKeys());
					this.schema = schema;

					schema.typeOverride = overrideType;

					var rf = getRecipeFactory();

					if (rf != null) {
						schema.recipeFactory = rf;
					}

					var constructors = getConstructors();

					if (!constructors.isEmpty()) {
						for (var c : constructors) {
							var cKeys = new ArrayList<RecipeKey<?>>();

							for (var keyName : c.keys()) {
								var key = keyMap.get(keyName);

								if (key != null) {
									cKeys.add(key);
								} else {
									throw new NullPointerException("Key '" + keyName + "' not found in constructor of recipe schema '" + id + "'");
								}
							}

							var constructor = new RecipeConstructor(List.copyOf(cKeys));

							if (!c.overrides().isEmpty()) {
								for (var entry : c.overrides().entrySet()) {
									var key = keyMap.get(entry.getKey());

									if (key != null) {
										try {
											constructor.overrideValue(key, Cast.to(key.codec.parse(jsonOps, entry.getValue()).getOrThrow()));
										} catch (Exception ex) {
											throw new IllegalArgumentException("Failed to create optional value for key '" + key + "' of '" + id + "' from " + entry.getValue(), ex);
										}
									} else {
										throw new NullPointerException("Key '" + entry.getKey() + "' not found in overrides of constructor of recipe schema '" + id + "'");
									}
								}
							}

							schema.constructor(constructor);
						}
					}

					for (var entry : functionMap.entrySet()) {
						var func = entry.getValue().resolve(jsonOps, schema);

						if (func.isSuccess()) {
							schema.function(new RecipeFunctionInstance(entry.getKey(), func.getOrThrow()));
						} else {
							throw new NullPointerException("Failed to parse function '" + entry.getKey() + "' of recipe schema '" + id + "': " + func.error().map(Error::message).orElse("Unknown Error"));
						}
					}

					var uniqueKeyNames = getUnique();

					if (!uniqueKeyNames.isEmpty()) {
						var uniqueKeys = new ArrayList<RecipeKey<?>>();

						for (var keyName : uniqueKeyNames) {
							var key = keyMap.get(keyName);

							if (key != null) {
								uniqueKeys.add(key);
							} else {
								throw new NullPointerException("Key '" + keyName + "' not found in unique keys of recipe schema '" + id + "'");
							}
						}

						schema.uniqueIds(uniqueKeys);
					}

					schema.hidden = isHidden();

					for (var postProcessor : getPostProcessors()) {
						schema.postProcessor(postProcessor);
					}
				} else if (parent != null) {
					schema = parent.getSchema(jsonOps);
				} else {
					schema = new RecipeSchema(Map.of(), List.of());
					schema.constructor();
				}
			}

			return schema;
		}
	}

	static void load(DynamicOps<JsonElement> ops, RecipeSchemaStorage storage, RecipeSchemaRegistry event, ResourceManager resourceManager) {
		var map = new HashMap<Identifier, RecipeSchemaBuilder>();
		var recipeSchemaDataCodec = RecipeSchemaData.CODEC;

		for (var entry : resourceManager.listResources("kubejs/recipe_schema", path -> path.getPath().endsWith(".json")).entrySet()) {
			try (var reader = entry.getValue().openAsReader()) {
				var json = JsonUtils.GSON.fromJson(reader, JsonObject.class);
				var id = entry.getKey().withPath(entry.getKey().getPath().substring("kubejs/recipe_schema/".length(), entry.getKey().getPath().length() - ".json".length()));

				switch (recipeSchemaDataCodec.parse(ops, json)) {
					case Success(var value, _) -> map.put(id, new RecipeSchemaBuilder(id, value));
					case Error(var msg, _, _) -> ConsoleJS.SERVER.error("Error parsing recipe schema json %s: %s".formatted(entry.getKey(), msg.get()));
				}
			} catch (Exception ex) {
				ConsoleJS.SERVER.error("Error reading recipe schema json " + entry.getKey(), ex);
			}
		}

		for (var builder : map.values()) {
			for (var m : builder.data.mappings()) {
				storage.mappings.put(m, builder.id);
			}
		}

		for (var builder : map.values()) {
			builder.hidden = builder.data.hidden().orElse(null);
			builder.parent = builder.data.parent().map(map::get).orElse(null);
			builder.overrideType = builder.data.overrideType().orElse(null);

			if (builder.data.recipeFactory().isPresent()) {
				var fname = builder.data.recipeFactory().get();
				builder.recipeFactory = storage.recipeTypes.get(fname);

				if (builder.recipeFactory == null) {
					throw new NullPointerException("Recipe factory '" + fname + "' not found for recipe schema '" + builder.id + "'");
				}
			}

			if (builder.data.keys().isPresent()) {
				builder.keys = new ArrayList<>();

				for (var keyData : builder.data.keys().get()) {
					try {
						var type = keyData.type();
						var key = type.key(keyData.name(), keyData.role());

						if (keyData.defaultOptional()) {
							key.defaultOptional();
						} else if (keyData.optional().isPresent()) {
							var optionalJson = keyData.optional().get();

							try {
								//noinspection unchecked, rawtypes
								key.optional = (RecipeOptional) RecipeOptional.unit(key.codec.decode(ops, optionalJson).getOrThrow().getFirst());
							} catch (Exception ex) {
								throw new IllegalArgumentException("Failed to create optional value for key '" + key + "' of '" + builder.id + "' from " + optionalJson, ex);
							}
						}

						if (!keyData.alternativeNames().isEmpty()) {
							key.names.addAll(keyData.alternativeNames());
						}

						key.excluded = keyData.excluded();

						if (!keyData.functionNames().isEmpty()) {
							key.functionNames = keyData.functionNames();
						}

						key.alwaysWrite = keyData.alwaysWrite();
						builder.keys.add(key);
					} catch (Exception ex) {
						ConsoleJS.SERVER.error("Error parsing recipe schema '" + builder.id + "' key " + keyData.name(), ex);
					}
				}
			}

			if (builder.data.constructors().isPresent()) {
				builder.constructors = builder.data.constructors().get();
			}

			if (builder.data.unique().isPresent()) {
				builder.unique = builder.data.unique().get();
			}

			if (builder.data.functions().isPresent()) {
				builder.functions = builder.data.functions().get();
			}

			if (!builder.data.overrideKeys().isEmpty()) {
				builder.overrideKeys = builder.data.overrideKeys();
			}

			if (builder.data.postProcessors().isPresent()) {
				builder.postProcessors = builder.data.postProcessors().get();
			}
		}

		for (var builder : map.values()) {
			var schema = builder.getSchema(ops);
			// System.out.println("SCHEMA " + builder.id);

			for (var constructor : schema.constructors().values()) {
				for (var key : schema.keys) {
					if (key.optional != null && !constructor.keys.contains(key) && !constructor.overrides.containsKey(key)) {
						constructor.defaultValue(key, Cast.to(key.optional));
						// System.out.println("- V DEF " + key.toString());
					} else {
						// System.out.println("- X NOP " + key.toString());
					}
				}

				// System.out.println("^ " + constructor);
			}
		}

		for (var builder : map.values()) {
			var schema = builder.getSchema(ops);
			event.namespace(builder.id.getNamespace()).register(builder.id.getPath(), schema);
		}
	}
}
