package pokefenn.totemic.client.rendering.model;

import java.util.List;
import java.util.Map;
import java.util.function.Function;

import javax.annotation.Nullable;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;

import net.minecraft.block.state.IBlockState;
import net.minecraft.client.renderer.block.model.BakedQuad;
import net.minecraft.client.renderer.block.model.IBakedModel;
import net.minecraft.client.renderer.texture.TextureAtlasSprite;
import net.minecraft.client.renderer.vertex.VertexFormat;
import net.minecraft.client.resources.IResourceManager;
import net.minecraft.util.EnumFacing;
import net.minecraft.util.ResourceLocation;
import net.minecraftforge.client.model.BakedModelWrapper;
import net.minecraftforge.client.model.ICustomModelLoader;
import net.minecraftforge.client.model.IModel;
import net.minecraftforge.client.model.ModelLoaderRegistry;
import net.minecraftforge.common.model.IModelState;
import net.minecraftforge.common.model.TRSRTransformation;
import net.minecraftforge.common.property.IExtendedBlockState;
import pokefenn.totemic.api.TotemicRegistries;
import pokefenn.totemic.api.totem.TotemEffect;
import pokefenn.totemic.block.totem.BlockTotemPole;
import pokefenn.totemic.init.ModContent;

public class ModelTotemPole implements IModel
{
    private final IModel blankModel;
    private final Map<TotemEffect, IModel> totemModels;

    public ModelTotemPole(IModel blankModel, Map<TotemEffect, IModel> totemModels)
    {
        this.blankModel = blankModel;
        this.totemModels = totemModels;
    }

    public ModelTotemPole()
    {
        this.blankModel = ModelLoaderRegistry.getModelOrMissing(new ResourceLocation("totemic", "block/totem_pole_blank"));

        ImmutableMap.Builder<TotemEffect, IModel> builder = ImmutableMap.builder();
        for(TotemEffect totem : TotemicRegistries.totemEffects())
        {
            ResourceLocation name = totem.getRegistryName();
            builder.put(totem, ModelLoaderRegistry.getModelOrLogError(new ResourceLocation(name.getResourceDomain(), "block/totem_pole_" + name.getResourcePath()),
                    "Could not load Totem Pole model for " + name));
        }
        this.totemModels = builder.build();
    }

    @Override
    public IBakedModel bake(IModelState state, VertexFormat format, Function<ResourceLocation, TextureAtlasSprite> bakedTextureGetter)
    {
        Map<TotemEffect, IBakedModel> bakedTotemModels = Maps.transformEntries(totemModels, (totem, model) -> {
            if(totem == ModContent.buffaloTotem) //FIXME: Band-aid to fix wrongly rotated Buffalo model
                return model.bake(new TRSRTransformation(EnumFacing.EAST), format, bakedTextureGetter);
            else
                return model.bake(state, format, bakedTextureGetter);
        });
        return new BakedTotemPole(blankModel.bake(state, format, bakedTextureGetter), ImmutableMap.copyOf(bakedTotemModels));
    }

    @Override
    public IModel retexture(ImmutableMap<String, String> textures)
    {
        return new ModelTotemPole(blankModel.retexture(textures),
                ImmutableMap.copyOf(Maps.transformValues(totemModels, model -> model.retexture(textures))));
    }

    public static final class BakedTotemPole extends BakedModelWrapper<IBakedModel>
    {
        private final ImmutableMap<TotemEffect, IBakedModel> bakedTotemModels;

        public BakedTotemPole(IBakedModel bakedBlankModel, ImmutableMap<TotemEffect, IBakedModel> bakedTotemModels)
        {
            super(bakedBlankModel);
            this.bakedTotemModels = bakedTotemModels;
        }

        @Override
        public List<BakedQuad> getQuads(@Nullable IBlockState state, @Nullable EnumFacing side, long rand)
        {
            if(state instanceof IExtendedBlockState)
            {
                TotemEffect effect = ((IExtendedBlockState) state).getValue(BlockTotemPole.TOTEM);
                if(effect != null)
                    return bakedTotemModels.get(effect).getQuads(state, side, rand);
            }
            return originalModel.getQuads(state, side, rand);
        }
    }

    public enum Loader implements ICustomModelLoader
    {
        INSTANCE;

        @Override
        public void onResourceManagerReload(IResourceManager resourceManager)
        { }

        @Override
        public boolean accepts(ResourceLocation modelLocation)
        {
            return modelLocation.getResourceDomain().equals("totemic") && modelLocation.getResourcePath().equals("models/block/totem_pole");
        }

        @Override
        public IModel loadModel(ResourceLocation modelLocation) throws Exception
        {
            return new ModelTotemPole();
        }
    }
}

@@ -64,34 +64,9 @@
 til.
-List;%0Aimport java.util.Map
+*
 ;%0Aim
@@ -1825,17 +1825,16 @@
 ct totem
- 
 : Totemi
@@ -1861,16 +1861,16 @@
 ects())%0A
+
         
@@ -2167,32 +2167,32 @@
 me));%0A        %7D%0A
-
         this.tot
@@ -2218,32 +2218,654 @@
 build();%0A    %7D%0A%0A
+    @Override%0A    public Collection%3CResourceLocation%3E getDependencies()%0A    %7B%0A        Set%3CResourceLocation%3E deps = new HashSet%3C%3E();%0A        deps.addAll(blankModel.getDependencies());%0A        for(IModel model: totemModels.values())%0A            deps.addAll(model.getDependencies());%0A        return deps;%0A    %7D%0A%0A    @Override%0A    public Collection%3CResourceLocation%3E getTextures()%0A    %7B%0A        Set%3CResourceLocation%3E textures = new HashSet%3C%3E();%0A        textures.addAll(blankModel.getTextures());%0A        for(IModel model: totemModels.values())%0A            textures.addAll(model.getTextures());%0A        return textures;%0A    %7D%0A%0A
     @Override%0A  

Make sure totem pole textures are being loaded