using System.Collections.Generic;
using Microsoft.Xna.Framework;
using Microsoft.Xna.Framework.Graphics;

namespace Kav
{
    public class Renderer
    {
        private const int MAX_SHADOW_CASCADES = 4;
        private int ShadowMapSize { get; }

        private GraphicsDevice GraphicsDevice { get; }
        private int RenderDimensionsX { get; }
        private int RenderDimensionsY { get; }

        private VertexBuffer FullscreenTriangle { get; }
        private int NumShadowCascades { get; }

        private RenderTarget2D ColorRenderTarget { get; }
        private RenderTarget2D DirectionalRenderTarget { get; }
        private RenderTarget2D[] ShadowRenderTargets { get; }

        private DeferredPBREffect DeferredPBREffect { get; }
        private DeferredPBR_AmbientLightEffect DeferredAmbientLightEffect { get; }
        private DeferredPBR_PointLightEffect DeferredPointLightEffect { get; }
        private DeferredPBR_DirectionalLightEffect DeferredDirectionalLightEffect { get; }
        private SimpleDepthEffect SimpleDepthEffect { get; }
        private Effect ToneMapEffect { get; }

        private RenderTarget2D gPosition { get; }
        private RenderTarget2D gNormal { get; }
        private RenderTarget2D gAlbedo { get; }
        private RenderTarget2D gMetallicRoughness { get; }

        private RenderTargetBinding[] GBuffer { get; }

        private SpriteBatch SpriteBatch { get; }

        public Renderer(
            GraphicsDevice graphicsDevice,
            int renderDimensionsX,
            int renderDimensionsY,
            int numShadowCascades,
            int shadowMapSize
        ) {
            GraphicsDevice = graphicsDevice;
            RenderDimensionsX = renderDimensionsX;
            RenderDimensionsY = renderDimensionsY;

            ShadowMapSize = shadowMapSize;

            NumShadowCascades = (int)MathHelper.Clamp(numShadowCascades, 1, MAX_SHADOW_CASCADES);
            ShadowRenderTargets = new RenderTarget2D[numShadowCascades];

            for (var i = 0; i < numShadowCascades; i++)
            {
                ShadowRenderTargets[i] = new RenderTarget2D(
                    GraphicsDevice,
                    ShadowMapSize,
                    ShadowMapSize,
                    false,
                    SurfaceFormat.Single,
                    DepthFormat.Depth24
                );
            }

            ColorRenderTarget = new RenderTarget2D(
                graphicsDevice,
                renderDimensionsX,
                renderDimensionsY,
                false,
                SurfaceFormat.Color,
                DepthFormat.None,
                0,
                RenderTargetUsage.PreserveContents
            );

            DirectionalRenderTarget = new RenderTarget2D(
                graphicsDevice,
                renderDimensionsX,
                renderDimensionsY,
                false,
                SurfaceFormat.Color,
                DepthFormat.None,
                0,
                RenderTargetUsage.PreserveContents
            );

            gPosition = new RenderTarget2D(
                GraphicsDevice,
                renderDimensionsX,
                renderDimensionsY,
                false,
                SurfaceFormat.Vector4,
                DepthFormat.Depth24
            );

            gNormal = new RenderTarget2D(
                GraphicsDevice,
                renderDimensionsX,
                renderDimensionsY,
                false,
                SurfaceFormat.Vector4,
                DepthFormat.None
            );

            gAlbedo = new RenderTarget2D(
                GraphicsDevice,
                renderDimensionsX,
                renderDimensionsY,
                false,
                SurfaceFormat.Color,
                DepthFormat.None
            );

            gMetallicRoughness = new RenderTarget2D(
                GraphicsDevice,
                renderDimensionsX,
                renderDimensionsY,
                false,
                SurfaceFormat.HalfVector2,
                DepthFormat.None
            );

            GBuffer = new RenderTargetBinding[4] {
                new RenderTargetBinding(gPosition),
                new RenderTargetBinding(gNormal),
                new RenderTargetBinding(gAlbedo),
                new RenderTargetBinding(gMetallicRoughness)
            };

            SimpleDepthEffect = new SimpleDepthEffect(GraphicsDevice);
            DeferredPBREffect = new DeferredPBREffect(GraphicsDevice);
            DeferredAmbientLightEffect = new DeferredPBR_AmbientLightEffect(GraphicsDevice);
            DeferredPointLightEffect = new DeferredPBR_PointLightEffect(GraphicsDevice);
            DeferredDirectionalLightEffect = new DeferredPBR_DirectionalLightEffect(GraphicsDevice);
            DeferredDirectionalLightEffect.ShadowMapSize = ShadowMapSize;
            ToneMapEffect = new Effect(graphicsDevice, Resources.ToneMapEffect);

            FullscreenTriangle = new VertexBuffer(GraphicsDevice, typeof(VertexPositionTexture), 3, BufferUsage.WriteOnly);
            FullscreenTriangle.SetData(new VertexPositionTexture[3] {
                new VertexPositionTexture(new Vector3(-1, -3, 0), new Vector2(0, 2)),
                new VertexPositionTexture(new Vector3(-1, 1, 0), new Vector2(0, 0)),
                new VertexPositionTexture(new Vector3(3, 1, 0), new Vector2(2, 0))
            });

            SpriteBatch = new SpriteBatch(graphicsDevice);
        }

        public void DeferredRender(
            PerspectiveCamera camera,
            IEnumerable<(Model, Matrix)> modelTransforms,
            IEnumerable<PointLight> pointLights,
            DirectionalLight directionalLight
        ) {
            // g-buffer pass

            GraphicsDevice.SetRenderTargets(GBuffer);
            GraphicsDevice.Clear(ClearOptions.Target | ClearOptions.DepthBuffer, Color.Black, 1f, 0);
            GraphicsDevice.DepthStencilState = DepthStencilState.Default;
            GraphicsDevice.BlendState = BlendState.Opaque;

            foreach (var (model, transform) in modelTransforms)
            {
                foreach (var modelMesh in model.Meshes)
                {
                    foreach (var meshPart in modelMesh.MeshParts)
                    {
                        GraphicsDevice.SetVertexBuffer(meshPart.VertexBuffer);
                        GraphicsDevice.Indices = meshPart.IndexBuffer;

                        if (meshPart.Effect is TransformEffect transformEffect)
                        {
                            transformEffect.World = transform;
                            transformEffect.View = camera.View;
                            transformEffect.Projection = camera.Projection;
                        }

                        foreach (var pass in meshPart.Effect.CurrentTechnique.Passes)
                        {
                            pass.Apply();

                            GraphicsDevice.DrawIndexedPrimitives(
                                PrimitiveType.TriangleList,
                                0,
                                0,
                                meshPart.VertexBuffer.VertexCount,
                                0,
                                meshPart.Triangles.Length
                            );
                        }
                    }
                }
            }

            GraphicsDevice.SetRenderTarget(ColorRenderTarget);
            GraphicsDevice.Clear(Color.Black);
            GraphicsDevice.BlendState = BlendState.Additive;
            GraphicsDevice.DepthStencilState = DepthStencilState.None;

            DeferredAmbientLightEffect.GPosition = gPosition;
            DeferredAmbientLightEffect.GAlbedo = gAlbedo;

            foreach (var pass in DeferredAmbientLightEffect.CurrentTechnique.Passes)
            {
                pass.Apply();
                GraphicsDevice.SetVertexBuffer(FullscreenTriangle);
                GraphicsDevice.DrawPrimitives(PrimitiveType.TriangleList, 0, 1);
            }

            DeferredPointLightEffect.EyePosition = camera.Position;

            foreach (var pointLight in pointLights)
            {
                PointLightRender(pointLight);
            }

            DirectionalLightRender(camera, modelTransforms, directionalLight);
            // return;
            // GraphicsDevice.SetRenderTarget(null);
            // SpriteBatch.Begin(SpriteSortMode.Deferred, BlendState.NonPremultiplied);
            // SpriteBatch.Draw(DirectionalRenderTarget, Vector2.Zero, Color.White);
            // SpriteBatch.End();

            GraphicsDevice.SetRenderTarget(null);
            GraphicsDevice.Clear(Color.Black);
            SpriteBatch.Begin(SpriteSortMode.Deferred, BlendState.Opaque, null, null, null, ToneMapEffect);
            SpriteBatch.Draw(ColorRenderTarget, Vector2.Zero, Color.White);
            SpriteBatch.End();
        }

        private void PointLightRender(PointLight pointLight) 
        {
            DeferredPointLightEffect.GPosition = gPosition;
            DeferredPointLightEffect.GAlbedo = gAlbedo;
            DeferredPointLightEffect.GNormal = gNormal;
            DeferredPointLightEffect.GMetallicRoughness = gMetallicRoughness;    

            DeferredPointLightEffect.PointLightPosition = pointLight.Position;
            DeferredPointLightEffect.PointLightColor = 
                pointLight.Color.ToVector3() * pointLight.Intensity;

            foreach (var pass in DeferredPointLightEffect.CurrentTechnique.Passes)
            {
                pass.Apply();
                GraphicsDevice.SetVertexBuffer(FullscreenTriangle);
                GraphicsDevice.DrawPrimitives(PrimitiveType.TriangleList, 0, 1);
            }
        }

        private void DirectionalLightRender(
            PerspectiveCamera camera,
            IEnumerable<(Model, Matrix)> modelTransforms,
            DirectionalLight directionalLight
        ) {
            // render the individual shadow cascades
            var previousFarPlane = camera.NearPlane;
            for (var i = 0; i < NumShadowCascades; i++)
            {
                var farPlane = camera.FarPlane / (MathHelper.Max((NumShadowCascades - i - 1) * 2f, 1f));

                // divide the view frustum 
                var shadowCamera = new PerspectiveCamera(
                    camera.Position,
                    camera.Forward,
                    camera.Up,
                    camera.FieldOfView,
                    camera.AspectRatio,
                    previousFarPlane,
                    farPlane
                );
                
                // TODO: This is tightly coupled to the effect and it sucks
                RenderShadowMap(shadowCamera, modelTransforms, directionalLight, i);

                previousFarPlane = farPlane;
            }

            DeferredDirectionalLightEffect.GPosition = gPosition;
            DeferredDirectionalLightEffect.GAlbedo = gAlbedo;
            DeferredDirectionalLightEffect.GNormal = gNormal;
            DeferredDirectionalLightEffect.GMetallicRoughness = gMetallicRoughness;

            DeferredDirectionalLightEffect.ShadowMapOne = ShadowRenderTargets[0];
            if (NumShadowCascades > 1)
            {
                DeferredDirectionalLightEffect.ShadowMapTwo = ShadowRenderTargets[1];
            }
            if (NumShadowCascades > 2)
            {
                DeferredDirectionalLightEffect.ShadowMapThree = ShadowRenderTargets[2];
            }
            if (NumShadowCascades > 3)
            {
                DeferredDirectionalLightEffect.ShadowMapFour = ShadowRenderTargets[3];
            }
            
            DeferredDirectionalLightEffect.DirectionalLightDirection = directionalLight.Direction;
            DeferredDirectionalLightEffect.DirectionalLightColor = 
                directionalLight.Color.ToVector3() * directionalLight.Intensity;
            
            DeferredDirectionalLightEffect.ViewMatrix = camera.View;
            DeferredDirectionalLightEffect.EyePosition = Matrix.Invert(camera.View).Translation;

            GraphicsDevice.SetRenderTarget(ColorRenderTarget);
            GraphicsDevice.BlendState = BlendState.Additive;

            foreach (EffectPass pass in DeferredDirectionalLightEffect.CurrentTechnique.Passes)
            {
                pass.Apply();
                GraphicsDevice.SetVertexBuffer(FullscreenTriangle);
                GraphicsDevice.DrawPrimitives(PrimitiveType.TriangleList, 0, 1);
            }
        }

        private void RenderShadowMap(
            PerspectiveCamera camera, 
            IEnumerable<(Model, Matrix)> modelTransforms, 
            DirectionalLight directionalLight,
            int shadowCascadeIndex
        ) {
            GraphicsDevice.SetRenderTarget(ShadowRenderTargets[shadowCascadeIndex]);
            GraphicsDevice.Clear(Color.White);
            GraphicsDevice.DepthStencilState = DepthStencilState.Default;
            GraphicsDevice.BlendState = BlendState.Opaque;
            
            var cameraBoundingFrustum = new BoundingFrustum(camera.View * camera.Projection);
            Vector3[] frustumCorners = cameraBoundingFrustum.GetCorners();

            Vector3 frustumCenter = Vector3.Zero;
            for (var i = 0; i < frustumCorners.Length; i++)
            {
                frustumCenter += frustumCorners[i];
            }
            frustumCenter /= 8f;

            var lightView = Matrix.CreateLookAt(frustumCenter + directionalLight.Direction, frustumCenter, Vector3.Backward);

            for (var i = 0; i < frustumCorners.Length; i++)
            {
                frustumCorners[i] = Vector3.Transform(frustumCorners[i], lightView);
            }

            BoundingBox lightBox = BoundingBox.CreateFromPoints(frustumCorners);
            
            SimpleDepthEffect.View = lightView;
            SimpleDepthEffect.Projection = Matrix.CreateOrthographicOffCenter(
                lightBox.Min.X,
                lightBox.Max.X,
                lightBox.Min.Y,
                lightBox.Max.Y,
                -lightBox.Max.Z - 10f, // TODO: near clip plane needs scene AABB info to get rid of this magic value
                -lightBox.Min.Z
            );

            var lightSpaceMatrix = SimpleDepthEffect.View * SimpleDepthEffect.Projection;

            if (shadowCascadeIndex == 0)
            {
                DeferredDirectionalLightEffect.LightSpaceMatrixOne = lightSpaceMatrix;
            }
            else if (shadowCascadeIndex == 1)
            {
                DeferredDirectionalLightEffect.LightSpaceMatrixTwo = lightSpaceMatrix;
            }
            else if (shadowCascadeIndex == 2)
            {
                DeferredDirectionalLightEffect.LightSpaceMatrixThree = lightSpaceMatrix;
            }
            else if (shadowCascadeIndex == 3)
            {
                DeferredDirectionalLightEffect.LightSpaceMatrixFour = lightSpaceMatrix;
            }

            DeferredDirectionalLightEffect.CascadeFarPlanes[shadowCascadeIndex] = camera.FarPlane;

            foreach (var (model, transform) in modelTransforms)
            {
                foreach (var modelMesh in model.Meshes)
                {
                    foreach (var meshPart in modelMesh.MeshParts)
                    {
                        GraphicsDevice.SetVertexBuffer(meshPart.VertexBuffer);
                        GraphicsDevice.Indices = meshPart.IndexBuffer;

                        SimpleDepthEffect.Model = transform;

                        foreach (var pass in SimpleDepthEffect.CurrentTechnique.Passes)
                        {
                            pass.Apply();

                            GraphicsDevice.DrawIndexedPrimitives(
                                PrimitiveType.TriangleList,
                                0,
                                0,
                                meshPart.VertexBuffer.VertexCount,
                                0,
                                meshPart.Triangles.Length
                            );
                        }
                    }
                }
            }
        }

        public void Render(
            PerspectiveCamera camera, 
            IEnumerable<(Model, Matrix)> modelTransforms, 
            IEnumerable<PointLight> pointLights,
            IEnumerable<DirectionalLight> directionalLights
        ) {
            Render(camera.View, camera.Projection, modelTransforms, pointLights, directionalLights);
        }

        private void Render(
            Matrix view,
            Matrix projection,
            IEnumerable<(Model, Matrix)> modelTransforms,
            IEnumerable<PointLight> pointLights,
            IEnumerable<DirectionalLight> directionalLights
        ) {
            foreach (var (model, transform) in modelTransforms)
            {
                foreach (var modelMesh in model.Meshes)
                {
                    foreach (var meshPart in modelMesh.MeshParts)
                    {
                        GraphicsDevice.SetVertexBuffer(meshPart.VertexBuffer);
                        GraphicsDevice.Indices = meshPart.IndexBuffer;

                        if (meshPart.Effect is TransformEffect transformEffect)
                        {
                            transformEffect.World = transform;
                            transformEffect.View = view;
                            transformEffect.Projection = projection;
                        }

                        if (meshPart.Effect is PointLightEffect pointLightEffect)
                        {
                            int i = 0;
                            foreach (var pointLight in pointLights)
                            {
                                if (i > pointLightEffect.MaxPointLights) { break; }
                                pointLightEffect.PointLights[i] = pointLight;
                                i++;
                            }
                        }

                        foreach (var pass in meshPart.Effect.CurrentTechnique.Passes)
                        {
                            pass.Apply();

                            GraphicsDevice.DrawIndexedPrimitives(
                                PrimitiveType.TriangleList,
                                0,
                                0,
                                meshPart.VertexBuffer.VertexCount,
                                0,
                                meshPart.Triangles.Length
                            );
                        }
                    }
                }
            }
        }
    }
}