package pokefenn.totemic.client.rendering.model;

import static pokefenn.totemic.Totemic.logger;

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.block.model.ModelResourceLocation;
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.property.IExtendedBlockState;
import pokefenn.totemic.api.TotemicRegistries;
import pokefenn.totemic.api.totem.TotemEffect;
import pokefenn.totemic.block.totem.BlockTotemPole;

public class ModelTotemPole implements IModel
{
    public static final ModelTotemPole MODEL = new ModelTotemPole();

    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()
    {
        ImmutableMap<String, String> defaultTextures = ImmutableMap.of("wood", "totemic:blocks/stripped_cedar_log",  "particle", "totemic:blocks/stripped_cedar_log");

        this.blankModel = ModelLoaderRegistry.getModelOrMissing(new ResourceLocation("totemic", "block/totem_pole_blank")).retexture(defaultTextures);

        ImmutableMap.Builder<TotemEffect, IModel> builder = ImmutableMap.builder();
        for(TotemEffect totem : TotemicRegistries.totemEffects())
        {
            ResourceLocation name = totem.getRegistryName();
            try
            {
                builder.put(totem, ModelLoaderRegistry.getModel(new ResourceLocation(name.getResourceDomain(), "block/totem_pole_" + name.getResourcePath())).retexture(defaultTextures));
            }
            catch(Exception e)
            {
                logger.error("Could not load Totem Pole model for " + name, e);
                builder.put(totem, blankModel);
            }
        }
        this.totemModels = builder.build();
    }

    @Override
    public IBakedModel bake(IModelState state, VertexFormat format, Function<ResourceLocation, TextureAtlasSprite> bakedTextureGetter)
    {
        return new BakedTotemPole(blankModel.bake(state, format, bakedTextureGetter),
                ImmutableMap.copyOf(Maps.transformValues(totemModels, model -> model.bake(state, format, bakedTextureGetter))));
    }

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

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

        public BakedTotemPole(IBakedModel bakedBlankModel, Map<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("totem_pole");
        }

        @Override
        public IModel loadModel(ResourceLocation modelLocation) throws Exception
        {
            //TODO: Hardcoded until I figure out how to get variant definitions from the blockstate JSON.
            String variant = ((ModelResourceLocation) modelLocation).getVariant();
            if(variant.equals("wood=cedar"))
                return MODEL;
            if(variant.startsWith("wood="))
            {
                String texture = "totemic:blocks/stripped_" + variant.substring(5) + "_log";
                return MODEL.retexture(ImmutableMap.of("wood", texture,  "particle", texture));
            }
            return MODEL;
        }
    }
}

@@ -47,56 +47,8 @@
 l;%0A%0A
-import static pokefenn.totemic.Totemic.logger;%0A%0A
 impo
@@ -66,16 +66,16 @@
 l.List;%0A
+
 import j
@@ -1286,78 +1286,8 @@
 l%0A%7B%0A
-    public static final ModelTotemPole MODEL = new ModelTotemPole();%0A%0A
     
@@ -2129,42 +2129,8 @@
 ();%0A
-            try%0A            %7B%0A    
     
@@ -2184,16 +2184,26 @@
 getModel
+OrLogError
 (new Res
@@ -2291,83 +2291,9 @@
 h())
-).retexture(defaultTextures));%0A            %7D%0A            catch(Exception e)
+,
 %0A   
@@ -2291,36 +2291,32 @@
 h()),%0A          
-  %7B%0A
                 
@@ -2313,27 +2313,8 @@
     
-      logger.error(
 %22Cou
@@ -2358,75 +2358,38 @@
 name
-, e);%0A                builder.put(totem, blankModel);%0A            %7D
+).retexture(defaultTextures));
 %0A   
@@ -4050,16 +4050,62 @@
 TANCE;%0A%0A
+        private ModelTotemPole MODEL = null;%0A%0A
         
@@ -4201,16 +4201,50 @@
        %7B
+%0A            MODEL = null;%0A       
  %7D%0A%0A    
@@ -4568,16 +4568,16 @@
 ception%0A
-
         
@@ -4570,32 +4570,109 @@
 ption%0A        %7B%0A
+            if(MODEL == null)%0A                MODEL = new ModelTotemPole();%0A%0A
             //TO

Allow realoding of totem pole models