using System;
using System.Collections.Generic;
using System.Linq;
using System.Reactive.Linq;
using DynamicData;
using Intromat.Graphics;
using Intromat.Model.Compiler;
using Intromat.Nodes.Code;
using Intromat.PersistentModel;
using Intromat.ViewModels;
using Intromat.ViewModels.Editors;
using Intromat.Views;
using NodeNetwork.ViewModels;
using ReactiveUI;
using SharpDX;
using SharpDX.D3DCompiler;
using SharpDX.Direct3D;
using SharpDX.Direct3D11;
using Splat;
using Buffer = SharpDX.Direct3D11.Buffer;

namespace Intromat.Nodes.Textures
{
    public class ShaderTextureNode : DxTextureNodeBase
    {
        private readonly List<Buffer> _constantBuffers = new();
        private ComputeShader? _cs;
        private byte[]? _byteCode;
        private TextureValue[]? _customOutputs;
        private readonly List<(NodeInputViewModel, InputDesc)> _customInputs = new();
        private readonly EndpointGroup _group;
        private readonly Dictionary<string, object> _inputValues = new();

        public record InputDesc(int ConstantBufferIndex, int Offset, Type VariableType, string Name);

        static ShaderTextureNode()
        {
            Locator.CurrentMutable.Register(() => new CodeGenNodeView(), typeof(IViewFor<ShaderTextureNode>));
        }

        public ShaderTextureNode()
        {
            Name = "Shader Texture";
            _group = new EndpointGroup("Shader Texture");
            Inputs.Add(Source = new CodeGenInputViewModel<ITypedExpression<string>>(EPortType.String)
            {
                Name = "Source",
                Group = _group,
                Editor = SourceEditor
            });

            var dxHost = Locator.Current.GetService<DxHost>()!;
            Source.ValueChanged
                .ObserveOn(dxHost.RenderDispatcher)
                .Subscribe(_ =>
                {
                    _resourcesCreated = false;
                    UpdateFrame(dxHost.Device, dxHost.Device.ImmediateContext);
                });
        }

        public override NodeModelBase CreateModel()
        {
            return new ShaderTextureModel();
        }

        public override void SaveModel(NodeModelBase model)
        {
            base.SaveModel(model);
            var shaderTexture = (ShaderTextureModel)model;
            shaderTexture.Source = SourceEditor.CreateModel();
            shaderTexture.InputValues = _customInputs
                .Where(i => i.Item1.Editor != null)
                .Select(i => new LiteralModelEntry()
                {
                    Key = i.Item1.Name,
                    Literal = CreateLiteralModel((dynamic)i.Item1.Editor)
                })
                .ToArray();
        }

        private TModel CreateLiteralModel<T, TValue, TModel>(ExpressionEditorBaseViewModel<T, TValue, TModel> editor) where TValue : LiteralValueBase<T>, new() where TModel : LiteralModelBase<T>, new()
        {
            return editor.CreateModel();
        }

        public override void LoadModel(NodeModelBase model)
        {
            base.LoadModel(model);
            var shaderTexture = (ShaderTextureModel)model;
            SourceEditor.LoadModel(shaderTexture.Source);
            foreach (var inputValue in shaderTexture.InputValues)
            {
                _inputValues[inputValue.Key] = inputValue.Literal;
            }
            
            var dxHost = Locator.Current.GetService<DxHost>()!;
            _resourcesCreated = true;
            CreateDeviceResources(dxHost.Device);
        }

        protected override void CreateDeviceResources(Device device)
        {
            foreach (var cb in _constantBuffers)
                cb.Dispose();

            _constantBuffers.Clear();

            var flags = ShaderFlags.None;
#if DEBUG
            flags |= ShaderFlags.Debug;
#endif
            var errors = new List<string>();
            var shader = Source.Value.Evaluate();
            if (string.IsNullOrEmpty(shader))
            {
                errors.Add("Shader is empty");
            }
            else
            {
                try
                {
                    var compilationResult = ShaderBytecode.Compile(shader, "main", "cs_5_0", flags);
                    if (compilationResult.HasErrors)
                    {
                        errors.Add(compilationResult.Message);
                    }
                    else
                    {
                        _byteCode = compilationResult.Bytecode.Data;
                        var reflection = new ShaderReflection(_byteCode);
                        var newInputs = new List<InputDesc>();
                        var newOutputs = new List<InputDesc>();
                        for (int i = 0; i < reflection.Description.BoundResources; ++i)
                        {
                            var resBindingDesc = reflection.GetResourceBindingDescription(i);
                            switch (resBindingDesc.Type)
                            {
                                case ShaderInputType.ConstantBuffer:
                                {
                                    var cb = reflection.GetConstantBuffer(resBindingDesc.Name);
                                    var cbIndex = resBindingDesc.BindPoint;
                                    var cbSize = cb.Description.Size;
                                    var bufferDesc = new BufferDescription(cbSize, ResourceUsage.Dynamic, BindFlags.ConstantBuffer, CpuAccessFlags.Write, ResourceOptionFlags.None, 0);
                                    _constantBuffers.Add(new Buffer(device, bufferDesc));
                                    for (int j = 0; j < cb.Description.VariableCount; ++j)
                                    {
                                        var variable = cb.GetVariable(j);
                                        var varTypeDesc = variable.GetVariableType()
                                            .Description;
                                        var type = varTypeDesc.Type switch
                                        {
                                            ShaderVariableType.Float => typeof(float),
                                            ShaderVariableType.Int   => typeof(int),
                                            ShaderVariableType.Bool  => typeof(bool),
                                            _                        => null
                                        };
                                        if (type == null)
                                        {
                                            errors.Add($"Constant buffer {resBindingDesc.Name} at slot {cbIndex} has a variable {variable.Description.Name} with unsupported type {varTypeDesc.Type}");
                                            continue;
                                        }

                                        for (int k = 0; k < varTypeDesc.ColumnCount; ++k)
                                        {
                                            var suffix = "" + 'X' + k;
                                            var inputName = varTypeDesc.ColumnCount == 1 ? variable.Description.Name : $"{variable.Description.Name}.{suffix}";
                                            newInputs.Add(new(cbIndex, variable.Description.StartOffset + k * variable.Description.Size, type, inputName));
                                        }
                                    }

                                    break;
                                }
                                case ShaderInputType.Texture:
                                {
                                    var texName = resBindingDesc.Name;
                                    newInputs.Add(new(-1, resBindingDesc.BindPoint, typeof(TextureValue), texName));
                                    break;
                                }
                                case ShaderInputType.Sampler:
                                {
                                    var texName = resBindingDesc.Name;
                                    newInputs.Add(new(-1, resBindingDesc.BindPoint, typeof(SamplerLiteralValue), texName));
                                    break;
                                }
                                case ShaderInputType.UnorderedAccessViewRWTyped:
                                {
                                    var texName = resBindingDesc.Name;
                                    if (resBindingDesc.Dimension != ShaderResourceViewDimension.Texture2D)
                                    {
                                        errors.Add($"UAV {texName} at slot {resBindingDesc.BindPoint} has unsupported dimension {resBindingDesc.Dimension}");
                                        break;
                                    }

                                    newOutputs.Add(new(-1, resBindingDesc.BindPoint, typeof(TextureValue), texName));
                                    break;
                                }
                                default:
                                {
                                    errors.Add($"Resource {resBindingDesc.Name} at slot {resBindingDesc.BindPoint} has unsupported type {resBindingDesc.Type}");
                                    break;
                                }
                            }
                        }

                        if (newOutputs.Count == 0)
                        {
                            errors.Add("The shader doesn't contain any UAV suitable for the texture output");
                        }
                        
                        Inputs.Edit(inputs =>
                        {
                            var maxSamplerStateIndex = 0;
                            for (var i = _customInputs.Count - 1; i >= 0; --i)
                            {
                                var existingInput = _customInputs[i];
                                if (newInputs.All(inputDesc => inputDesc.Name != existingInput.Item1.Name))
                                {
                                    inputs.Remove(existingInput.Item1);
                                    _customInputs.RemoveAt(i);
                                }
                            }

                            foreach (var newInputDesc in newInputs)
                            {
                                if (_customInputs.Any(i => i.Item2 == newInputDesc))
                                    continue;

                                NodeInputViewModel newInput;
                                if (newInputDesc.VariableType == typeof(float))
                                {
                                    newInput = new CodeGenInputViewModel<ITypedExpression<float>>(EPortType.Float)
                                    {
                                        Name = newInputDesc.Name,
                                        Group = _group,
                                        Editor = new FloatExpressionEditorViewModel { MaxValue = 1 }
                                    };
                                    if (_inputValues.TryGetValue(newInputDesc.Name, out var inputValue) && inputValue is FloatLiteralModel model)
                                        ((FloatExpressionEditorViewModel)newInput.Editor).LoadModel(model);
                                }
                                else if (newInputDesc.VariableType == typeof(int))
                                {
                                    newInput = new CodeGenInputViewModel<ITypedExpression<int>>(EPortType.Integer)
                                    {
                                        Name = newInputDesc.Name,
                                        Group = _group,
                                        Editor = new IntegerExpressionEditorViewModel { MaxValue = 100 }
                                    };
                                    if (_inputValues.TryGetValue(newInputDesc.Name, out var inputValue) && inputValue is IntLiteralModel model)
                                        ((IntegerExpressionEditorViewModel)newInput.Editor).LoadModel(model);
                                }
                                else if (newInputDesc.VariableType == typeof(bool))
                                {
                                    newInput = new CodeGenInputViewModel<ITypedExpression<bool>>(EPortType.Boolean)
                                    {
                                        Name = newInputDesc.Name,
                                        Group = _group,
                                        Editor = new BooleanExpressionEditorViewModel()
                                    };
                                    if (_inputValues.TryGetValue(newInputDesc.Name, out var inputValue) && inputValue is BooleanLiteralModel model)
                                        ((BooleanExpressionEditorViewModel)newInput.Editor).LoadModel(model);
                                }
                                else if (newInputDesc.VariableType == typeof(TextureValue))
                                {
                                    newInput = new CodeGenInputViewModel<TextureValue>(EPortType.Texture) { Name = newInputDesc.Name };
                                    if (newInputDesc.Offset == 0)
                                    {
                                        var inputObservable = ((CodeGenInputViewModel<TextureValue>)newInput).ValueChanged;
                                        WidthEditor.InputValueProvider = inputObservable.Select(input => input?.Width ?? 0);
                                        HeightEditor.InputValueProvider = inputObservable.Select(input => input?.Height ?? 0);
                                    }
                                }
                                else if (newInputDesc.VariableType == typeof(SamplerLiteralValue))
                                {
                                    maxSamplerStateIndex = Math.Max(maxSamplerStateIndex, newInputDesc.Offset);
                                    newInput = new CodeGenInputViewModel<ITypedExpression<SamplerDesc>>(EPortType.None)
                                    {
                                        Name = newInputDesc.Name,
                                        Group = _group,
                                        Editor = new SamplerEditorViewModel { CustomValue = new SamplerDesc() }
                                    };
                                    if (_inputValues.TryGetValue(newInputDesc.Name, out var inputValue) && inputValue is SamplerLiteralModel model)
                                        ((SamplerEditorViewModel)newInput.Editor).LoadModel(model);
                                }
                                else
                                {
                                    continue;
                                }

                                _customInputs.Add((newInput, newInputDesc));
                                inputs.Add(newInput);
                            }
                        });

                        Outputs.Edit(outputs =>
                        {
                            if (outputs.Count > newOutputs.Count)
                            {
                                if (newOutputs.Count == 0)
                                {
                                    if (outputs.Count > 1)
                                        outputs.RemoveRange(1, outputs.Count - 1);
                                    outputs[0]
                                        .Name = "Output";
                                }
                                else
                                {
                                    outputs.RemoveRange(newOutputs.Count, outputs.Count - newOutputs.Count);
                                }
                            }

                            var outputIndex = 0;
                            if (outputIndex < newOutputs.Count)
                                for (; outputIndex < outputs.Count; ++outputIndex)
                                    outputs[outputIndex]
                                        .Name = newOutputs[outputIndex]
                                        .Name;

                            for (; outputIndex < newOutputs.Count; ++outputIndex)
                            {
                                var newOutputDesc = newOutputs[outputIndex];
                                NodeOutputViewModel newOutput;
                                if (newOutputDesc.VariableType == typeof(TextureValue))
                                    newOutput = new CodeGenOutputViewModel<TextureValue>(EPortType.Texture) { Name = newOutputDesc.Name };
                                else
                                    continue;

                                outputs.Add(newOutput);
                            }

                            _customOutputs = new TextureValue[outputs.Count];
                        });
                    }
                }
                catch (CompilationException e)
                {
                    errors.Add(e.Message);
                }
            }
            
            _cs?.Dispose();
            _cs = errors.Count == 0 ? new ComputeShader(device, _byteCode) : null;

            Errors.Edit(e =>
            {
                e.Clear();
                e.AddRange(errors);
            });
        }

        protected override unsafe void UpdateFrame(Device device, DeviceContext context)
        {
            if (!_resourcesCreated)
            {
                _resourcesCreated = true;
                var dxHost = Locator.Current.GetService<DxHost>()!;
                dxHost.MainDispatcher!.Invoke(() => CreateDeviceResources(device));
            }

            EnsureBuffers(device);
            if (_customOutputs == null || _cs == null)
                return;

            _customOutputs[0] = _output;
            for (int i = 1; i < Outputs.Count; ++i)
            {
                EnsureTexture(device, ref _customOutputs![i]);
            }

            var cbData = _constantBuffers.Select(cb =>
            {
                context.MapSubresource(cb, 0, MapMode.WriteDiscard, MapFlags.None, out var stream);
                return stream.DataPointer;
            }).ToArray();

            int maxSrvIndex = 0;
            foreach (var input in _customInputs)
            {
                var desc = input.Item2;
                if (desc.ConstantBufferIndex >= 0)
                {
                    var dataPtr = cbData[desc.ConstantBufferIndex] + desc.Offset;
                    if (desc.VariableType == typeof(float))
                    {
                        _inputValues[input.Item1.Name] = ((FloatExpressionEditorViewModel)input.Item1.Editor).CreateModel();
                        *(float*)dataPtr = ((CodeGenInputViewModel<ITypedExpression<float>>)input.Item1).Value.Evaluate();
                    }
                    else if (desc.VariableType == typeof(int))
                    {
                        _inputValues[input.Item1.Name] = ((IntegerExpressionEditorViewModel)input.Item1.Editor).CreateModel();
                        *(int*)dataPtr = ((CodeGenInputViewModel<ITypedExpression<int>>)input.Item1).Value.Evaluate();
                    }
                    else if (desc.VariableType == typeof(bool))
                    {
                        _inputValues[input.Item1.Name] = ((BooleanExpressionEditorViewModel)input.Item1.Editor).CreateModel();
                        *(bool*)dataPtr = ((CodeGenInputViewModel<ITypedExpression<bool>>)input.Item1).Value.Evaluate();
                    }
                }
                else
                {
                    if (desc.VariableType == typeof(TextureValue))
                    {
                        maxSrvIndex = Math.Max(maxSrvIndex, desc.Offset);
                        context.ComputeShader.SetShaderResource(desc.Offset, ((CodeGenInputViewModel<TextureValue>)input.Item1).Value?.ShaderResourceView);
                    }
                    else if (desc.VariableType == typeof(SamplerLiteralValue))
                    {
                        _inputValues[input.Item1.Name] = ((SamplerEditorViewModel)input.Item1.Editor).CreateModel();
                        var samplerValue = (SamplerLiteralValue)((CodeGenInputViewModel<ITypedExpression<SamplerDesc>>)input.Item1).Value;
                        EnsureSamplerState(device, ref samplerValue);
                        context.ComputeShader.SetSampler(desc.Offset, samplerValue.SamplerState);
                    }
                }
            }

            for (var i = 0; i < _customOutputs!.Length; i++)
            {
                var output = _customOutputs![i];
                context.ComputeShader.SetUnorderedAccessView(i, output.UnorderedAccessView);
            }

            for (int i = 0; i < _constantBuffers.Count; ++i)
            {
                context.UnmapSubresource(_constantBuffers[i], 0);
                context.ComputeShader.SetConstantBuffer(i, _constantBuffers[i]);
            }

            context.ComputeShader.SetShader(_cs, null, 0);
            context.Dispatch(_texDesc.Width / 16, _texDesc.Height / 16, 1);

            for (var i = 0; i < _customOutputs!.Length; i++)
            {
                context.ComputeShader.SetUnorderedAccessView(i, null);
            }

            for (var i = 0; i <= maxSrvIndex; i++)
            {
                context.ComputeShader.SetShaderResource(i, null);
            }
        }

        private void EnsureSamplerState(Device device, ref SamplerLiteralValue samplerLiteralValue)
        {
            if (samplerLiteralValue.SamplerState == null || samplerLiteralValue.SamplerState.Description.AddressU != samplerLiteralValue.Value.Address || samplerLiteralValue.SamplerState.Description.Filter != samplerLiteralValue.Value.Filter)
            {
                var desc = SamplerStateDescription.Default();
                desc.AddressU = desc.AddressV = desc.AddressW = samplerLiteralValue.Value.Address;
                desc.Filter = samplerLiteralValue.Value.Filter;
                samplerLiteralValue.SamplerState = new SamplerState(device, desc);
            }
        }

        public ISourceList<string> Errors { get; } = new SourceList<string>();
        public CodeGenInputViewModel<ITypedExpression<string>> Source { get; }

        public StringExpressionEditorViewModel SourceEditor { get; } = new()
        {
            CustomValue = @"
cbuffer cb
{
    float f;
}
Texture2D t0;
SamplerState ss;
RWTexture2D<float4> o;
[numthreads(16,16,1)]
void main(uint3 i : SV_DispatchThreadID)
{
    float2 s;
    o.GetDimensions(s.x, s.y);
    float2 uv = i.xy/s;
    float4 c = t0.SampleLevel(ss, uv, 0);
    o[i.xy] = float4(c.xyz, 1);
}"
        };
    }
}