package com.mineplex.studio;

import com.google.common.base.MoreObjects;
import com.google.common.base.Objects;
import com.google.common.collect.ImmutableList;
import com.mojang.datafixers.util.Pair;
import com.mojang.serialization.Codec;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import net.minecraft.core.component.DataComponentPatch;
import net.minecraft.core.component.DataComponentType;
import net.minecraft.core.component.DataComponents;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.NbtOps;
import net.minecraft.nbt.Tag;
import net.minecraft.network.RegistryFriendlyByteBuf;
import net.minecraft.network.chat.Component;
import net.minecraft.network.chat.ComponentSerialization;
import net.minecraft.network.codec.StreamCodec;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.item.component.CustomData;
import org.jspecify.annotations.NonNull;

public class MineplexI18NUtil {
    private static final String STORED_NAME = "i18n.item_data";

    private static final ImmutableList<TypeTransformer<?>> COMPONENT_TYPES = ImmutableList.of(
            new TypeTransformer<>("minecraft:item_name", DataComponents.ITEM_NAME),
            new TypeTransformer<>("minecraft:custom_name", DataComponents.CUSTOM_NAME),
            new TypeTransformer<>("minecraft:lore", DataComponents.LORE));

    private MineplexI18NUtil() {}

    @NonNull public static DataComponentPatch translate(@NonNull final ItemStack stack) {
        final DataComponentPatch patch = stack.getComponentsPatch();
        if (patch.isEmpty()) {
            return DataComponentPatch.EMPTY;
        }

        // Store the untranslated data in the custom data to restore it later during receive.
        final Map<String, Tag> transformedTags = HashMap.newHashMap(COMPONENT_TYPES.size());
        for (final TypeTransformer<?> typeTransformer : COMPONENT_TYPES) {
            typeTransformer.encode(patch).ifPresent(tag -> transformedTags.put(typeTransformer.getKey(), tag));
        }
        if (transformedTags.isEmpty()) {
            return patch;
        }

        // We need to rebuild the patch with the new custom data.
        // Reusing the same patch is not an option since we only want this data to exist on the wire or the client.
        final DataComponentPatch.Builder newBuilder = DataComponentPatch.builder();
        newBuilder.copy(patch);

        final CustomData currentCustomData = stack.get(DataComponents.CUSTOM_DATA);
        final CompoundTag tag =
                currentCustomData == null ? new CompoundTag(HashMap.newHashMap(1)) : currentCustomData.copyTag();
        tag.put(STORED_NAME, new CompoundTag(transformedTags));

        newBuilder.set(DataComponents.CUSTOM_DATA, CustomData.of(tag));

        return newBuilder.build();
    }

    @NonNull public static List<Component> translateComponents(
            @NonNull final List<Component> components, @NonNull final Locale locale) {
        // The DONT_RENDER_TRANSLATABLES logic needs to be in sync with the part from ItemStrack
        final boolean prev = ComponentSerialization.DONT_RENDER_TRANSLATABLES.get();
        try {
            ComponentSerialization.DONT_RENDER_TRANSLATABLES.set(true);
            final List<Component> lines = new ArrayList<>(components.size());
            for (final Component line : components) {
                final net.kyori.adventure.text.Component advLine =
                        io.papermc.paper.adventure.PaperAdventure.asAdventure(line);
                net.kyori.adventure.text.Component i18nLine =
                        net.kyori.adventure.translation.GlobalTranslator.render(advLine, locale);
                Component vanillaLine = io.papermc.paper.adventure.PaperAdventure.asVanilla(i18nLine);
                lines.add(vanillaLine);
            }
            return lines;
        } finally {
            ComponentSerialization.DONT_RENDER_TRANSLATABLES.set(prev);
        }
    }

    public static void removeTranslations(@NonNull final ItemStack stack) {
        final CustomData customData = stack.get(DataComponents.CUSTOM_DATA);
        if (customData == null || customData.isEmpty()) {
            return;
        }

        // GetCompound would create a new CompoundTag if the tag is not present.
        if (!customData.contains(STORED_NAME)) {
            return;
        }

        final CompoundTag savedTag = customData.getUnsafe().getCompound(STORED_NAME);
        // We never want the data to be available on the server itself.
        customData.getUnsafe().remove(STORED_NAME);

        for (final TypeTransformer<?> typeTransformer : COMPONENT_TYPES) {
            typeTransformer.decode(stack, savedTag);
        }
    }

    public static class LoreCodec implements StreamCodec<RegistryFriendlyByteBuf, List<Component>> {
        @NonNull private final StreamCodec<RegistryFriendlyByteBuf, List<Component>> currentCodec;

        public LoreCodec(@NonNull StreamCodec<RegistryFriendlyByteBuf, List<Component>> currentCodec) {
            this.currentCodec = currentCodec;
        }

        @Override
        public List<Component> decode(final RegistryFriendlyByteBuf buf) {
            return currentCodec.decode(buf);
        }

        @Override
        public void encode(final RegistryFriendlyByteBuf buf, final List<Component> value) {
            final List<Component> components =
                    buf.adventure$locale == null ? value : translateComponents(value, buf.adventure$locale);
            currentCodec.encode(buf, components);
        }
    }

    private static final class TypeTransformer<T> {
        @NonNull private final String key;

        @NonNull private final DataComponentType<T> componentType;

        public TypeTransformer(@NonNull String key, @NonNull DataComponentType<T> componentType) {
            this.key = key;
            this.componentType = componentType;
        }

        @NonNull public Codec<T> getCodec() {
            return componentType.codecOrThrow();
        }

        public void decode(@NonNull final ItemStack itemStack, final CompoundTag compoundTag) {
            final Tag encoded = compoundTag.get(getKey());
            if (encoded == null) {
                return;
            }

            getCodec()
                    .decode(NbtOps.INSTANCE, encoded)
                    .result()
                    .map(Pair::getFirst)
                    .ifPresent(decoded -> itemStack.set(getComponentType(), decoded));
        }

        @NonNull public Optional<@NonNull Tag> encode(@NonNull final DataComponentPatch dataComponentPatch) {
            //noinspection unchecked
            final Optional<T> value = (Optional<T>) dataComponentPatch.get(componentType);
            //noinspection OptionalAssignedToNull
            if (value == null || value.isEmpty()) {
                return Optional.empty();
            }

            return getCodec().encodeStart(NbtOps.INSTANCE, value.get()).result();
        }

        public @NonNull String getKey() {
            return this.key;
        }

        public @NonNull DataComponentType<T> getComponentType() {
            return this.componentType;
        }

        @Override
        public String toString() {
            return MoreObjects.toStringHelper(this)
                    .add("key", key)
                    .add("componentType", componentType)
                    .toString();
        }

        @Override
        public boolean equals(Object o) {
            if (!(o instanceof TypeTransformer<?> that)) return false;
            return Objects.equal(getKey(), that.getKey()) && Objects.equal(getComponentType(), that.getComponentType());
        }

        @Override
        public int hashCode() {
            return Objects.hashCode(getKey(), getComponentType());
        }
    }
}
