package net.minecraft.world.item.equipment.trim;

import com.google.common.collect.Maps;
import com.mojang.serialization.Codec;
import com.mojang.serialization.MapCodec;
import com.mojang.serialization.codecs.RecordCodecBuilder;
import io.netty.buffer.ByteBuf;
import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
import java.util.Map;
import net.minecraft.network.codec.ByteBufCodecs;
import net.minecraft.network.codec.StreamCodec;
import net.minecraft.resources.ResourceKey;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.util.ExtraCodecs;
import net.minecraft.world.item.equipment.EquipmentAsset;
import net.minecraft.world.item.equipment.EquipmentAssets;

public record MaterialAssetGroup(MaterialAssetGroup.AssetInfo base, Map<ResourceKey<EquipmentAsset>, MaterialAssetGroup.AssetInfo> overrides) {
    public static final String SEPARATOR = "_";
    public static final MapCodec<MaterialAssetGroup> MAP_CODEC = RecordCodecBuilder.mapCodec(
        instance -> instance.group(
                MaterialAssetGroup.AssetInfo.CODEC.fieldOf("asset_name").forGetter(MaterialAssetGroup::base),
                Codec.unboundedMap(ResourceKey.codec(EquipmentAssets.ROOT_ID), MaterialAssetGroup.AssetInfo.CODEC)
                    .optionalFieldOf("override_armor_assets", Map.of())
                    .forGetter(MaterialAssetGroup::overrides)
            )
            .apply(instance, MaterialAssetGroup::new)
    );
    public static final StreamCodec<ByteBuf, MaterialAssetGroup> STREAM_CODEC = StreamCodec.composite(
        MaterialAssetGroup.AssetInfo.STREAM_CODEC,
        MaterialAssetGroup::base,
        ByteBufCodecs.map(Object2ObjectOpenHashMap::new, ResourceKey.streamCodec(EquipmentAssets.ROOT_ID), MaterialAssetGroup.AssetInfo.STREAM_CODEC),
        MaterialAssetGroup::overrides,
        MaterialAssetGroup::new
    );
    public static final MaterialAssetGroup QUARTZ = create("quartz");
    public static final MaterialAssetGroup IRON = create("iron", Map.of(EquipmentAssets.IRON, "iron_darker"));
    public static final MaterialAssetGroup NETHERITE = create("netherite", Map.of(EquipmentAssets.NETHERITE, "netherite_darker"));
    public static final MaterialAssetGroup REDSTONE = create("redstone");
    public static final MaterialAssetGroup COPPER = create("copper");
    public static final MaterialAssetGroup GOLD = create("gold", Map.of(EquipmentAssets.GOLD, "gold_darker"));
    public static final MaterialAssetGroup EMERALD = create("emerald");
    public static final MaterialAssetGroup DIAMOND = create("diamond", Map.of(EquipmentAssets.DIAMOND, "diamond_darker"));
    public static final MaterialAssetGroup LAPIS = create("lapis");
    public static final MaterialAssetGroup AMETHYST = create("amethyst");
    public static final MaterialAssetGroup RESIN = create("resin");

    public static MaterialAssetGroup create(String name) {
        return new MaterialAssetGroup(new MaterialAssetGroup.AssetInfo(name), Map.of());
    }

    public static MaterialAssetGroup create(String name, Map<ResourceKey<EquipmentAsset>, String> overrides) {
        return new MaterialAssetGroup(new MaterialAssetGroup.AssetInfo(name), Map.copyOf(Maps.transformValues(overrides, MaterialAssetGroup.AssetInfo::new)));
    }

    public MaterialAssetGroup.AssetInfo assetId(ResourceKey<EquipmentAsset> assetKey) {
        return this.overrides.getOrDefault(assetKey, this.base);
    }

    public record AssetInfo(String suffix) {
        public static final Codec<MaterialAssetGroup.AssetInfo> CODEC = ExtraCodecs.RESOURCE_PATH_CODEC
            .xmap(MaterialAssetGroup.AssetInfo::new, MaterialAssetGroup.AssetInfo::suffix);
        public static final StreamCodec<ByteBuf, MaterialAssetGroup.AssetInfo> STREAM_CODEC = ByteBufCodecs.STRING_UTF8
            .map(MaterialAssetGroup.AssetInfo::new, MaterialAssetGroup.AssetInfo::suffix);

        public AssetInfo(String suffix) {
            if (!ResourceLocation.isValidPath(suffix)) {
                throw new IllegalArgumentException("Invalid string to use as a resource path element: " + suffix);
            } else {
                this.suffix = suffix;
            }
        }
    }
}
