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

namespace Smuggler
{
    public struct PBRLight
    {
        public Vector3 position;
        public Vector3 color;

        public PBRLight(Vector3 position, Vector3 colour)
        {
            this.position = position;
            this.color = colour;
        }
    }

    public class PBRLightCollection
    {
        private readonly Vector3[] positions = new Vector3[4];
        private readonly Vector3[] colors = new Vector3[4];

        readonly EffectParameter lightPositionsParam;
        readonly EffectParameter lightColorsParam;

        public PBRLightCollection(EffectParameter lightPositionsParam, EffectParameter lightColorsParam)
        {
            this.lightPositionsParam = lightPositionsParam;
            this.lightColorsParam = lightColorsParam;
        }
        
        public PBRLight this[int i]
        {
            get { return new PBRLight(positions[i], colors[i]); }
            set 
            {
                positions[i] = value.position;
                colors[i] = value.color;
                lightPositionsParam.SetValue(positions);
                lightColorsParam.SetValue(colors);
            }
        }
    }

    public class PBREffect : Effect
    {
        EffectParameter worldParam;
        EffectParameter viewParam;
        EffectParameter projectionParam;
        EffectParameter worldViewProjectionParam;
        EffectParameter worldInverseTransposeParam;

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

        EffectParameter lightPositionsParam;
        EffectParameter lightColorsParam;

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

        EffectParameter eyePositionParam;

        Matrix world = Matrix.Identity;
        Matrix view = Matrix.Identity;
        Matrix projection = Matrix.Identity;
        PBRLightCollection pbrLightCollection;

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

        // FIXME: lazily set properties for performance

        public Matrix World
        {
            get { return world; }
            set 
            { 
                world = value;
                worldParam.SetValue(world);
                worldViewProjectionParam.SetValue(world * view * projection);
                worldInverseTransposeParam.SetValue(Matrix.Transpose(Matrix.Invert(world)));
            }
        }

        public Matrix View
        {
            get { return view; }
            set 
            { 
                view = value;
                viewParam.SetValue(view);
                worldViewProjectionParam.SetValue(world * view * projection);
                eyePositionParam.SetValue(Matrix.Invert(view).Translation);
            }
        }

        public Matrix Projection
        {
            get { return projection; }
            set 
            { 
                projection = value;
                projectionParam.SetValue(projection);
                worldViewProjectionParam.SetValue(world * view * projection);
            }
        }

        public PBRLightCollection Lights
        {
            get { return pbrLightCollection; }
            internal set { pbrLightCollection = 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 BaseColourTexture
        {
            get { return baseColorTextureParam.GetValueTexture2D(); }
            set { baseColorTextureParam.SetValue(value); }
        }

        public Texture2D NormalTexture
        {
            get { return normalTextureParam.GetValueTexture2D(); }
            set { normalTextureParam.SetValue(value); }
        }

        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); }
        }

        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(null);

            pbrLightCollection = new PBRLightCollection(
                Parameters["LightPositions"],
                Parameters["LightColors"]
            );
        }

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

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

            Lights = new PBRLightCollection(
                Parameters["LightPositions"],
                Parameters["LightColors"]
            );

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

            BaseColourTexture = cloneSource.BaseColourTexture;
            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);
        }

        // FIXME: do param applications here for performance
        protected override void OnApply()
        {
            base.OnApply();
        }

        void CacheEffectParameters(PBREffect cloneSource)
        {
            worldParam = Parameters["World"];
            viewParam = Parameters["View"];
            projectionParam = Parameters["Projection"];
            worldViewProjectionParam = Parameters["WorldViewProjection"];
            worldInverseTransposeParam = Parameters["WorldInverseTranspose"];

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

            lightPositionsParam = Parameters["LightPositions"];
            lightColorsParam = Parameters["LightColors"];

            albedoParam = Parameters["Albedo"];
            metallicParam = Parameters["Metallic"];
            roughnessParam = Parameters["Roughness"];
            aoParam = Parameters["AO"];

            eyePositionParam = Parameters["EyePosition"];
        }
    }
}