using UnityEngine;
using System.Collections.Generic;
using UnityEditor.Graphing;
using System;
using Unity.GraphAuthoring.Editor.ProviderSystem;

namespace UnityEditor.ShaderGraph.ProviderSystem
{
    // A Base Node serving as the entry point for Provider based function definitions.
    // This Model object can work with any sort of function provider, interpret the function definition,
    // and generate a valid model representation of that object. This abstracts the need for node definitions
    // to be aware of how the model functions.

    // NOTE: There exists an ExpressionProvider currently, so it's not a problem to have Models that are specializations,
    // but ideally the expression input type would be gracefully handled in a generic way and it wouldn't be necessary.
    [Serializable]
    [HasDependencies(typeof(MinimalProviderNode))]
    internal class ProviderNode : AbstractMaterialNode, IHasAssetDependencies, IGeneratesBodyCode, IGeneratesFunction
    {
        [Serializable]
        internal class MinimalProviderNode : IHasDependencies
        {
            [SerializeReference]
            IProvider<IShaderFunction> m_provider;

            public void GetSourceAssetDependencies(AssetCollection assetCollection)
            {
                if (m_provider == null || m_provider.AssetID == default)
                    return;

                assetCollection.AddAssetDependency(m_provider.AssetID,
                    AssetCollection.Flags.SourceDependency
                    | AssetCollection.Flags.ArtifactDependency
                    | AssetCollection.Flags.IncludeInExportPackage);
            }
        }


        // Reserves a control for the function selector. This is a bit of a hack.
        [GroupVariantControl]
        internal ProviderNode groupVariantSentinel => this;

        [SerializeReference]
        IProvider<IShaderFunction> m_provider;

        internal FunctionHeader Header { get; private set; }

        internal Dictionary<string, ParameterHeader> ParamHeaders { get; private set; }

        internal IProvider<IShaderFunction> Provider => m_provider;

        const int kReservedOutputSlot = 0;

        public override bool hasPreview => true;

        public override bool hasUserAuthoredCode => true;

        public override bool canSetPrecision => Header?.allowPrecision ?? false;

        public override string documentationURL => Header?.helpUrl ?? base.documentationURL;

        internal override bool ExposeToSearcher => false;

        internal virtual bool requiresGeneration => false;

        public ProviderNode()
        {
            name = "Provider Based Node";
            Header = new();
        }

        internal void InitializeFromProvider(IProvider<IShaderFunction> provider)
        {
            this.m_provider = (IProvider<IShaderFunction>)provider.Clone();
            UpdateModel();
        }

        public override void UpdateNodeAfterDeserialization()
        {
            base.UpdateNodeAfterDeserialization();
            Refresh();
        }

        public override void Concretize()
        {
            UpdateModel();
            base.Concretize();
        }

        public bool Reload(HashSet<string> changedAssetGuids)
        {
            // Scripted providers don't react to asset changes — their definition is in-memory.
            if (Provider.AssetID != default && changedAssetGuids.Contains(Provider.AssetID.ToString()))
            {
                Refresh();
                return true;
            }
            return false;
        }

        protected void Refresh()
        {
            Provider?.Reload();
            UpdateModel();
            ValidateNode();
            Dirty(ModificationScope.Graph);
        }

        internal void UpdateModel()
        {
            if (Provider == null || Provider.Definition == null || !Provider.IsValid)
                return;

            Header.Process(Provider.Definition, Provider);

            var header = Header;
            this.name = header.displayName;
            this.synonyms = header.searchTerms;
            var parameters = Provider.Definition.Parameters;

            List<MaterialSlot> previousSlots = new();
            GetSlots(previousSlots);

            Dictionary<string, (int inputId, int outputId)> oldSlotMap = new();
            HashSet<int> usedSlotIds = new();

            // build a map of parameter names to old slot ids (tuple is to support inout params).
            foreach (var oldSlot in previousSlots)
            {
                if (!oldSlotMap.TryGetValue(oldSlot.shaderOutputName, out var idTuple))
                    idTuple = (-1, -1);

                if (oldSlot.isInputSlot) idTuple.inputId = oldSlot.id;
                if (oldSlot.isOutputSlot) idTuple.outputId = oldSlot.id;

                oldSlotMap[oldSlot.shaderOutputName] = idTuple;
            }

            ParamHeaders = new();
            List<ParameterHeader> paramOrder = new();
            List<int> desiredSlotOrder = new();

            // return type is a special case, because it has no parameter but still needs a slot
            // we reserve the 0 slotId for this and always make sure it's added first/at the top.
            if (header.hasReturnValueType)
            {
                usedSlotIds.Add(kReservedOutputSlot);
                desiredSlotOrder.Add(kReservedOutputSlot);
                AddSlotFromParameter(header.returnHeader, kReservedOutputSlot, SlotType.Output);
            }

            // build the header data for our parameters and mark which slot ids are being reused.
            foreach(var param in parameters)
            {
                var paramHeader = new ParameterHeader(param, Provider);
                ParamHeaders.TryAdd(param.Name, paramHeader);
                paramOrder.Add(paramHeader);
                if (oldSlotMap.TryGetValue(param.Name, out var idTuple))
                {
                    if (idTuple.inputId > -1 && param.IsInput)
                        usedSlotIds.Add(idTuple.inputId);
                    if (idTuple.outputId > -1 && param.IsOutput)
                        usedSlotIds.Add(idTuple.outputId);
                }
            }

            // walk through our header data and build the actual slots.
            int nextSlot = kReservedOutputSlot;
            foreach (var paramHeader in paramOrder)
            {
                if (!oldSlotMap.TryGetValue(paramHeader.referenceName, out var idTuple))
                    idTuple = (-1, -1);

                void DoSlot(int slotId, SlotType dir)
                {
                    if (slotId == -1)
                    {
                        while (usedSlotIds.Contains(++nextSlot));
                        usedSlotIds.Add(nextSlot);
                        slotId = nextSlot;
                    }
                    AddSlotFromParameter(paramHeader, slotId, dir);
                    desiredSlotOrder.Add(slotId);
                }

                if (paramHeader.isInput)
                    DoSlot(idTuple.inputId, SlotType.Input);
                if (paramHeader.isOutput)
                    DoSlot(idTuple.outputId, SlotType.Output);
            }
            RemoveSlotsNameNotMatching(usedSlotIds, true);
            SetSlotOrder(desiredSlotOrder);
        }

        static Rendering.ShaderCompilerMessageSeverity ToCompilerSeverity(MessageType severity) => severity switch
        {
            MessageType.Error => Rendering.ShaderCompilerMessageSeverity.Error,
            _                  => Rendering.ShaderCompilerMessageSeverity.Warning,
        };

        void AddSlotFromParameter(ParameterHeader header, int slotId, SlotType dir)
        {
            var slot = HeaderUtils.MakeSlotFromParameter(header, slotId, dir);
            if (slot != null)
                // Slots that use label control have no indication in the Model that they use them.
                // However, the only _existing_ case where they'd meaningfully change are SpaceMaterialSlots,
                // which also don't handle default values well either. In this case, we go nuclear and completely
                // rebuild the slot.
                AddSlot(slot, slot is not SpaceMaterialSlot, slot is not SpaceMaterialSlot);
        }

        public void GenerateNodeCode(ShaderStringBuilder sb, GenerationMode generationMode)
        {
            if (this.Provider == null || this.Provider.Definition == null)
                return;

            // Apply legacy precision/dynamic only before code generation/usage.
            HeaderUtils.TryApplyLegacy(Provider.Definition, this.Header, this.ParamHeaders, PrecisionUtil.Token, lastKnownDynamicVectorLength, out var func);

            using (UnityEngine.Pool.ListPool<MaterialSlot>.Get(out var slots))
            {
                GetSlots<MaterialSlot>(slots);
                Dictionary<string, (MaterialSlot input, MaterialSlot output)> paramSlotMap = new();
                MaterialSlot returnSlot = null;

                // for various reasons, slots won't match one-to-one with parameters;
                // we'll need to make sure slots are matched to parameters in the order expected
                // while also accounting for inout duplicity.
                foreach (var slot in slots)
                {
                    if (slot.id == kReservedOutputSlot)
                        returnSlot = slot;
                    else
                    {
                        paramSlotMap.TryGetValue(slot.shaderOutputName, out var tup);
                        if (slot.isInputSlot) tup.input = slot;
                        if (slot.isOutputSlot) tup.output = slot;
                        paramSlotMap[slot.shaderOutputName] = tup;
                    }
                }

                ShaderStringBuilder call = new(); // for the callsite.
                ShaderStringBuilder args = new(); // for building up the argument list.

                // The return result is an output slot, but will instead get initialized by the function call.
                if (returnSlot != null)
                {
                    call.Append($"{HeaderUtils.ToShaderType(returnSlot)} {GetVariableNameForSlot(returnSlot.id)} = ");
                }

                bool first = true;
                foreach (var param in func.Parameters)
                {
                    var inputSlot = paramSlotMap[param.Name].input;
                    var outputSlot = paramSlotMap[param.Name].output;

                    var typeString = outputSlot != null ? HeaderUtils.ToShaderType(outputSlot) : HeaderUtils.ToShaderType(inputSlot);

                    var variableName = outputSlot != null ? GetVariableNameForSlot(outputSlot.id) : null;
                    var valueString = inputSlot != null ? GetSlotValue(inputSlot.id, generationMode) : null;

                    if (inputSlot != null && inputSlot.bareResource) // Texture/Sampler slots are normally 'Unity' wrapper structures
                        valueString += inputSlot is SamplerStateMaterialSlot ? ".samplerstate" : ".tex";

                    var argument = valueString; // assume it's an input, in which case the arg will be the upstream value/connection.

                    // apply linkage override
                    if (ParamHeaders.TryGetValue(param.Name, out var ph) && ph.isLinkage)
                    {
                        var targetSlots = paramSlotMap[ph.linkTarget];

                        // this concept is awkward for inout slots, but that could be improved in the future.
                        // ie. add support for int and set it to 0 = none, 1 = in, 2 = out, 3 = both.
                        if (targetSlots.input != null && targetSlots.input.isConnected
                        || targetSlots.output != null && targetSlots.output.isConnected)
                        {
                            argument = "true";
                        }
                        else argument = "false";
                    }

                    if (outputSlot != null)
                    {
                        // inout and out define a variable to be used as the argument.
                        // but only inout gets initialized by the value.
                        argument = variableName;
                        sb.AddLine($"{typeString} {variableName}{(inputSlot != null ? $" = {valueString}" : "")};");
                    }

                    if (!first)
                        args.Append(", ");
                    first = false;

                    args.Append(argument);
                }

                call.Append(ShaderObjectUtils.GenerateCall(func, args.ToString()));
                sb.AddLine(call.ToString());
            }
        }

        public void GenerateNodeFunction(FunctionRegistry registry, GenerationMode generationMode)
        {
            if (Provider == null || !Provider.IsValid || Provider.Definition == null)
                return;

            // Asset-backed providers contribute their source .hlsl as an include; scripted
            // providers carry their source inline via the generated function body below.
            if (Provider.AssetID != default)
            {
                var includePath = AssetDatabase.GUIDToAssetPath(Provider.AssetID);
                registry.RequiresIncludePath(includePath, false);
            }

            if (HeaderUtils.TryApplyLegacy(Provider.Definition, Header, ParamHeaders, PrecisionUtil.Token, lastKnownDynamicVectorLength, out var func) || requiresGeneration)
            {
                string code = ShaderObjectUtils.GenerateCode(func, false, false, true);
                registry.ProvideFunction(func.Name, s => s.AppendLine(code));
            }
        }

        public override void ValidateNode()
        {
            base.ValidateNode();
            if (Provider == null)
            {
                owner?.AddSetupError(this.objectId, "Node is in an invalid state and cannot recover.");
                return;
            }
            else if (Provider.AssetID != default && (Provider.Definition == null || !Provider.Definition.IsValid))
            {
                var path = AssetDatabase.GUIDToAssetPath(Provider.AssetID);
                owner?.AddSetupError(this.objectId, $"Data for '{Provider.ProviderKey}' expected in '{path}'.");
                return;
            }
            else
            {
                foreach (var (msg, severity) in Header.Messages)
                    owner?.AddValidationError(this.objectId, msg, ToCompilerSeverity(severity));

                foreach (var param in ParamHeaders.Values)
                    foreach (var (msg, severity) in param.Messages)
                        owner?.AddValidationError(this.objectId, msg, ToCompilerSeverity(severity));
            }
        }
    }
}