using Microsoft.Xna.Framework;
using Microsoft.Xna.Framework.Graphics;

namespace Kav
{
    public class PointLightCollection
    {
        private readonly Vector3[] positions = new Vector3[4];
        private readonly Vector3[] colors = new Vector3[4];

        readonly EffectParameter lightPositionsParam;
        readonly EffectParameter lightColorsParam;

        public PointLightCollection(EffectParameter lightPositionsParam, EffectParameter lightColorsParam)
        {
            this.lightPositionsParam = lightPositionsParam;
            this.lightColorsParam = lightColorsParam;
        }

        public PointLight this[int i]
        {
            get
            {
                return new PointLight(
                    positions[i],
                    new Color(
                        colors[i].X,
                        colors[i].Y,
                        colors[i].Z,
                        1f
                    )
                );
            }
            set
            {
                positions[i] = value.Position;
                colors[i] = value.Color.ToVector3() * value.Intensity;
                lightPositionsParam.SetValue(positions);
                lightColorsParam.SetValue(colors);
            }
        }
    }

    public class PBREffect : Effect, TransformEffect, PointLightEffect
    {
        EffectParameter worldParam;
        EffectParameter worldViewProjectionParam;
        EffectParameter worldInverseTransposeParam;

        EffectParameter albedoTextureParam;
        EffectParameter normalTextureParam;
        EffectParameter emissionTextureParam;
        EffectParameter occlusionTextureParam;
        EffectParameter metallicRoughnessTextureParam;
        EffectParameter envDiffuseTextureParam;
        EffectParameter brdfLutTextureParam;
        EffectParameter envSpecularTextureParam;

        EffectParameter albedoParam;
        EffectParameter metallicParam;
        EffectParameter roughnessParam;
        EffectParameter aoParam;

        EffectParameter eyePositionParam;

        EffectParameter shaderIndexParam;

        Matrix world = Matrix.Identity;
        Matrix view = Matrix.Identity;
        Matrix projection = Matrix.Identity;
        PointLightCollection pointLightCollection;

        Vector3 albedo;
        float metallic;
        float roughness;
        float ao;

        bool albedoTextureEnabled = false;
        bool metallicRoughnessMapEnabled = false;
        bool normalMapEnabled = false;

        EffectDirtyFlags dirtyFlags = EffectDirtyFlags.All;

        // FIXME: lazily set properties for performance

        public Matrix World
        {
            get { return world; }
            set
            {
                world = value;
                dirtyFlags |= EffectDirtyFlags.World | EffectDirtyFlags.WorldViewProj;
            }
        }

        public Matrix View
        {
            get { return view; }
            set
            {
                view = value;
                dirtyFlags |= EffectDirtyFlags.WorldViewProj | EffectDirtyFlags.EyePosition;
            }
        }

        public Matrix Projection
        {
            get { return projection; }
            set
            {
                projection = value;
                dirtyFlags |= EffectDirtyFlags.WorldViewProj;
            }
        }

        public PointLightCollection PointLights
        {
            get { return pointLightCollection; }
            internal set { pointLightCollection = value; }
        }

        public Vector3 Albedo
        {
            get { return albedo; }
            set
            {
                albedo = value;
                albedoParam.SetValue(albedo);
            }
        }

        public float Metallic
        {
            get { return metallic; }
            set
            {
                metallic = value;
                metallicParam.SetValue(metallic);
            }
        }

        public float Roughness
        {
            get { return roughness; }
            set
            {
                roughness = value;
                roughnessParam.SetValue(roughness);
            }
        }

        public float AO
        {
            get { return ao; }
            set
            {
                ao = value;
                aoParam.SetValue(ao);
            }
        }

        public Texture2D AlbedoTexture
        {
            get { return albedoTextureParam.GetValueTexture2D(); }
            set
            {
                albedoTextureParam.SetValue(value);
                albedoTextureEnabled = value != null;
                dirtyFlags |= EffectDirtyFlags.ShaderIndex;
            }
        }

        public Texture2D NormalTexture
        {
            get { return normalTextureParam.GetValueTexture2D(); }
            set
            {
                normalTextureParam.SetValue(value);
                normalMapEnabled = value != null;
                dirtyFlags |= EffectDirtyFlags.ShaderIndex;
            }
        }

        public Texture2D EmissionTexture
        {
            get { return emissionTextureParam.GetValueTexture2D(); }
            set { emissionTextureParam.SetValue(value); }
        }

        public Texture2D OcclusionTexture
        {
            get { return occlusionTextureParam.GetValueTexture2D(); }
            set { occlusionTextureParam.SetValue(value); }
        }

        public Texture2D MetallicRoughnessTexture
        {
            get { return metallicRoughnessTextureParam.GetValueTexture2D(); }
            set
            {
                metallicRoughnessTextureParam.SetValue(value);
                metallicRoughnessMapEnabled = value != null;
                dirtyFlags |= EffectDirtyFlags.ShaderIndex;
            }
        }

        public TextureCube EnvDiffuseTexture
        {
            get { return envDiffuseTextureParam.GetValueTextureCube(); }
            set { envDiffuseTextureParam.SetValue(value); }
        }

        public Texture2D BRDFLutTexture
        {
            get { return brdfLutTextureParam.GetValueTexture2D(); }
            set { brdfLutTextureParam.SetValue(value); }
        }

        public TextureCube EnvSpecularTexture
        {
            get { return envSpecularTextureParam.GetValueTextureCube(); }
            set { envSpecularTextureParam.SetValue(value); }
        }

        public PBREffect(GraphicsDevice graphicsDevice) : base(graphicsDevice, Resources.PBREffect)
        {
            CacheEffectParameters();

            pointLightCollection = new PointLightCollection(
                Parameters["LightPositions"],
                Parameters["LightColors"]
            );
        }

        protected PBREffect(PBREffect cloneSource) : base(cloneSource)
        {
            CacheEffectParameters();

            World = cloneSource.World;
            View = cloneSource.View;
            Projection = cloneSource.Projection;

            PointLights = new PointLightCollection(
                Parameters["LightPositions"],
                Parameters["LightColors"]
            );

            for (int i = 0; i < 4; i++)
            {
                PointLights[i] = cloneSource.PointLights[i];
            }

            AlbedoTexture = cloneSource.AlbedoTexture;
            NormalTexture = cloneSource.NormalTexture;
            EmissionTexture = cloneSource.EmissionTexture;
            OcclusionTexture = cloneSource.OcclusionTexture;
            MetallicRoughnessTexture = cloneSource.MetallicRoughnessTexture;
            EnvDiffuseTexture = cloneSource.EnvDiffuseTexture;
            BRDFLutTexture = cloneSource.BRDFLutTexture;
            EnvSpecularTexture = cloneSource.EnvSpecularTexture;

            Albedo = cloneSource.Albedo;
            Metallic = cloneSource.Metallic;
            Roughness = cloneSource.Roughness;
            AO = cloneSource.AO;
        }

        public override Effect Clone()
        {
            return new PBREffect(this);
        }

        protected override void OnApply()
        {
            if ((dirtyFlags & EffectDirtyFlags.World) != 0)
            {
                worldParam.SetValue(world);

                Matrix.Invert(ref world, out Matrix worldInverse);
                Matrix.Transpose(ref worldInverse, out Matrix worldInverseTranspose);
                worldInverseTransposeParam.SetValue(worldInverseTranspose);

                dirtyFlags &= ~EffectDirtyFlags.World;
            }

            if ((dirtyFlags & EffectDirtyFlags.WorldViewProj) != 0)
            {
                Matrix.Multiply(ref world, ref view, out Matrix worldView);
                Matrix.Multiply(ref worldView, ref projection, out Matrix worldViewProj);
                worldViewProjectionParam.SetValue(worldViewProj);

                dirtyFlags &= ~EffectDirtyFlags.WorldViewProj;
            }

            if ((dirtyFlags & EffectDirtyFlags.EyePosition) != 0)
            {
                Matrix.Invert(ref view, out Matrix inverseView);
                eyePositionParam.SetValue(inverseView.Translation);

                dirtyFlags &= ~EffectDirtyFlags.EyePosition;
            }

            if ((dirtyFlags & EffectDirtyFlags.ShaderIndex) != 0)
            {
                int shaderIndex = 0;

                if (albedoTextureEnabled && metallicRoughnessMapEnabled && normalMapEnabled)
                {
                    shaderIndex = 7;
                }
                else if (metallicRoughnessMapEnabled && normalMapEnabled)
                {
                    shaderIndex = 6;
                }
                else if (albedoTextureEnabled && normalMapEnabled)
                {
                    shaderIndex = 5;
                }
                else if (albedoTextureEnabled && metallicRoughnessMapEnabled)
                {
                    shaderIndex = 4;
                }
                else if (normalMapEnabled)
                {
                    shaderIndex = 3;
                }
                else if (metallicRoughnessMapEnabled)
                {
                    shaderIndex = 2;
                }
                else if (albedoTextureEnabled)
                {
                    shaderIndex = 1;
                }

                shaderIndexParam.SetValue(shaderIndex);

                dirtyFlags &= ~EffectDirtyFlags.ShaderIndex;
            }
        }

        void CacheEffectParameters()
        {
            worldParam = Parameters["World"];
            worldViewProjectionParam = Parameters["WorldViewProjection"];
            worldInverseTransposeParam = Parameters["WorldInverseTranspose"];

            albedoTextureParam = Parameters["AlbedoTexture"];
            normalTextureParam = Parameters["NormalTexture"];
            emissionTextureParam = Parameters["EmissionTexture"];
            occlusionTextureParam = Parameters["OcclusionTexture"];
            metallicRoughnessTextureParam = Parameters["MetallicRoughnessTexture"];
            envDiffuseTextureParam = Parameters["EnvDiffuseTexture"];
            brdfLutTextureParam = Parameters["BrdfLutTexture"];
            envSpecularTextureParam = Parameters["EnvSpecularTexture"];

            albedoParam = Parameters["AlbedoValue"];
            metallicParam = Parameters["MetallicValue"];
            roughnessParam = Parameters["RoughnessValue"];
            aoParam = Parameters["AO"];

            eyePositionParam = Parameters["EyePosition"];

            shaderIndexParam = Parameters["ShaderIndex"];
        }
    }
}