// SPDX-FileCopyrightText: 2023 The CC: Tweaked Developers
//
// SPDX-License-Identifier: MPL-2.0

package dan200.computercraft.shared.util;

import com.mojang.serialization.DataResult;
import com.mojang.serialization.MapCodec;
import net.minecraft.core.HolderLookup;
import net.minecraft.core.Registry;
import net.minecraft.core.registries.BuiltInRegistries;
import net.minecraft.network.RegistryFriendlyByteBuf;
import net.minecraft.network.codec.StreamCodec;
import net.minecraft.resources.Identifier;
import net.minecraft.resources.RegistryOps;
import net.minecraft.resources.ResourceKey;
import net.minecraft.util.ExtraCodecs;
import org.jetbrains.annotations.ApiStatus;

/**
 * Additioanl functions for working with {@linkplain Registry registries}.
 */
@ApiStatus.Internal
public final class RegistryHelper {
    private RegistryHelper() {
    }

    /**
     * Find a registry from a {@link ResourceKey}, throwing if it does not exist.
     *
     * @param id  The id of the registry.
     * @param <T> The contents of the registry
     * @return The associated registry.
     */
    @SuppressWarnings("unchecked")
    public static <T> Registry<T> getRegistry(ResourceKey<Registry<T>> id) {
        var registry = (Registry<T>) BuiltInRegistries.REGISTRY.getValue(id.identifier());
        if (registry == null) throw new IllegalArgumentException("Unknown registry " + id);
        return registry;
    }

    /**
     * Get the key of a registry entry, throwing if it is not registered.
     *
     * @param registry The registry to look up in.
     * @param object   The object to look up.
     * @param <T>      The type of this registry.
     * @return The ID of this object
     * @see Registry#getResourceKey(Object)
     */
    public static <T> Identifier getKeyOrThrow(Registry<T> registry, T object) {
        var key = registry.getResourceKey(object);
        if (key.isEmpty()) throw new IllegalArgumentException(object + " was not registered in " + registry.key());
        return key.get().identifier();
    }

    /**
     * Create a {@link MapCodec} that extracts a {@link HolderLookup.RegistryLookup} from the deserialisation context.
     * <p>
     * This is similar to {@link RegistryOps#retrieveGetter(ResourceKey)}, but returning a more general type.
     *
     * @param registryKey The registry to retrieve.
     * @param <T>         The type of object the registry contains.
     * @return A codec that retrieves the given registry, or errors if not available.
     * @see #retrieveRegistryStreamCodec(ResourceKey)
     */
    public static <T> MapCodec<HolderLookup.RegistryLookup<T>> retrieveRegistryCodec(ResourceKey<? extends Registry<? extends T>> registryKey) {
        return ExtraCodecs.retrieveContext(ops -> {
            if (!(ops instanceof RegistryOps<?> registryOps)) return DataResult.error(() -> "Not a registryKey ops");

            var holderGetter = registryOps.getter(registryKey);
            if (holderGetter.isEmpty()) return DataResult.error(() -> "Unknown registry key" + registryKey);

            if (!(holderGetter.get() instanceof HolderLookup.RegistryLookup<T> registryLookup)) {
                return DataResult.error(() -> "Registry " + registryKey + " is not a RegistryLookup");
            }

            return DataResult.success(registryLookup, registryLookup.registryLifecycle());
        });
    }

    /**
     * Create a {@link StreamCodec} that extracts a {@link Registry} from the deserialisation context.
     *
     * @param registryKey The registry to retrieve.
     * @param <T>         The type of object the registry contains.
     * @return A codec that retrieves the given registry, or errors if not available.
     * @see #retrieveRegistryCodec(ResourceKey)
     */
    public static <T> StreamCodec<RegistryFriendlyByteBuf, Registry<T>> retrieveRegistryStreamCodec(ResourceKey<? extends Registry<? extends T>> registryKey) {
        return StreamCodec.of(
            (_, _) -> {
            },
            x -> x.registryAccess().lookupOrThrow(registryKey)
        );
    }
}
