using Ryujinx.Graphics.Shader.IntermediateRepresentation;
using System.Collections.Generic;

using static Ryujinx.Graphics.Shader.IntermediateRepresentation.OperandHelper;
using static Ryujinx.Graphics.Shader.Translation.GlobalMemory;

namespace Ryujinx.Graphics.Shader.Translation.Optimizations
{
    static class GlobalToStorage
    {
        private struct SearchResult
        {
            public static SearchResult NotFound => new SearchResult(-1, 0);
            public bool Found => SbCbSlot != -1;
            public int SbCbSlot { get; }
            public int SbCbOffset { get; }

            public SearchResult(int sbCbSlot, int sbCbOffset)
            {
                SbCbSlot = sbCbSlot;
                SbCbOffset = sbCbOffset;
            }
        }

        public static void RunPass(BasicBlock block, ShaderConfig config, ref int sbUseMask, ref int ubeUseMask)
        {
            int sbStart = GetStorageBaseCbOffset(config.Stage);
            int sbEnd = sbStart + StorageDescsSize;

            int ubeStart = UbeBaseOffset;
            int ubeEnd = UbeBaseOffset + UbeDescsSize;

            for (LinkedListNode<INode> node = block.Operations.First; node != null; node = node.Next)
            {
                for (int index = 0; index < node.Value.SourcesCount; index++)
                {
                    Operand src = node.Value.GetSource(index);

                    int storageIndex = GetStorageIndex(src, sbStart, sbEnd);

                    if (storageIndex >= 0)
                    {
                        sbUseMask |= 1 << storageIndex;
                    }

                    if (config.Stage == ShaderStage.Compute)
                    {
                        int constantIndex = GetStorageIndex(src, ubeStart, ubeEnd);

                        if (constantIndex >= 0)
                        {
                            ubeUseMask |= 1 << constantIndex;
                        }
                    }
                }

                if (!(node.Value is Operation operation))
                {
                    continue;
                }

                if (UsesGlobalMemory(operation.Inst, operation.StorageKind))
                {
                    Operand source = operation.GetSource(0);

                    var result = SearchForStorageBase(config, block, source);
                    if (!result.Found)
                    {
                        continue;
                    }

                    if (config.Stage == ShaderStage.Compute &&
                        operation.Inst == Instruction.LoadGlobal &&
                        result.SbCbSlot == DriverReservedCb &&
                        result.SbCbOffset >= UbeBaseOffset &&
                        result.SbCbOffset < UbeBaseOffset + UbeDescsSize)
                    {
                        // Here we effectively try to replace a LDG instruction with LDC.
                        // The hardware only supports a limited amount of constant buffers
                        // so NVN "emulates" more constant buffers using global memory access.
                        // Here we try to replace the global access back to a constant buffer
                        // load.
                        node = ReplaceLdgWithLdc(node, config, (result.SbCbOffset - UbeBaseOffset) / StorageDescSize);
                    }
                    else
                    {
                        // Storage buffers are implemented using global memory access.
                        // If we know from where the base address of the access is loaded,
                        // we can guess which storage buffer it is accessing.
                        // We can then replace the global memory access with a storage
                        // buffer access.
                        node = ReplaceGlobalWithStorage(block, node, config, config.GetSbSlot((byte)result.SbCbSlot, (ushort)result.SbCbOffset));
                    }
                }
            }

            config.SetAccessibleBufferMasks(sbUseMask, ubeUseMask);
        }

        private static LinkedListNode<INode> ReplaceGlobalWithStorage(BasicBlock block, LinkedListNode<INode> node, ShaderConfig config, int storageIndex)
        {
            Operation operation = (Operation)node.Value;

            bool isAtomic = operation.Inst.IsAtomic();
            bool isStg16Or8 = operation.Inst == Instruction.StoreGlobal16 || operation.Inst == Instruction.StoreGlobal8;
            bool isWrite = isAtomic || operation.Inst == Instruction.StoreGlobal || isStg16Or8;

            config.SetUsedStorageBuffer(storageIndex, isWrite);

            Operand[] sources = new Operand[operation.SourcesCount];

            sources[0] = Const(storageIndex);
            sources[1] = GetStorageOffset(block, node, config, storageIndex, operation.GetSource(0), isStg16Or8);

            for (int index = 2; index < operation.SourcesCount; index++)
            {
                sources[index] = operation.GetSource(index);
            }

            Operation storageOp;

            if (isAtomic)
            {
                storageOp = new Operation(operation.Inst, StorageKind.StorageBuffer, operation.Dest, sources);
            }
            else if (operation.Inst == Instruction.LoadGlobal)
            {
                storageOp = new Operation(Instruction.LoadStorage, operation.Dest, sources);
            }
            else
            {
                Instruction storeInst = operation.Inst switch
                {
                    Instruction.StoreGlobal16 => Instruction.StoreStorage16,
                    Instruction.StoreGlobal8 => Instruction.StoreStorage8,
                    _ => Instruction.StoreStorage
                };

                storageOp = new Operation(storeInst, null, sources);
            }

            for (int index = 0; index < operation.SourcesCount; index++)
            {
                operation.SetSource(index, null);
            }

            LinkedListNode<INode> oldNode = node;

            node = node.List.AddBefore(node, storageOp);

            node.List.Remove(oldNode);

            return node;
        }

        private static Operand GetStorageOffset(
            BasicBlock block,
            LinkedListNode<INode> node,
            ShaderConfig config,
            int storageIndex,
            Operand addrLow,
            bool isStg16Or8)
        {
            (int sbCbSlot, int sbCbOffset) = config.GetSbCbInfo(storageIndex);

            bool storageAligned = !(config.GpuAccessor.QueryHasUnalignedStorageBuffer() || config.GpuAccessor.QueryHostStorageBufferOffsetAlignment() > Constants.StorageAlignment);

            (Operand byteOffset, int constantOffset) = storageAligned ?
                GetStorageOffset(block, Utils.FindLastOperation(addrLow, block), sbCbSlot, sbCbOffset) :
                (null, 0);

            if (byteOffset != null)
            {
                ReplaceAddressAlignment(node.List, addrLow, byteOffset, constantOffset);
            }

            if (byteOffset == null)
            {
                Operand baseAddrLow = Cbuf(sbCbSlot, sbCbOffset);
                Operand baseAddrTrunc = Local();

                Operand alignMask = Const(-config.GpuAccessor.QueryHostStorageBufferOffsetAlignment());

                Operation andOp = new Operation(Instruction.BitwiseAnd, baseAddrTrunc, baseAddrLow, alignMask);

                node.List.AddBefore(node, andOp);

                Operand offset = Local();
                Operation subOp = new Operation(Instruction.Subtract, offset, addrLow, baseAddrTrunc);

                node.List.AddBefore(node, subOp);

                byteOffset = offset;
            }
            else if (constantOffset != 0)
            {
                Operand offset = Local();
                Operation addOp = new Operation(Instruction.Add, offset, byteOffset, Const(constantOffset));

                node.List.AddBefore(node, addOp);

                byteOffset = offset;
            }

            if (isStg16Or8)
            {
                return byteOffset;
            }

            Operand wordOffset = Local();
            Operation shrOp = new Operation(Instruction.ShiftRightU32, wordOffset, byteOffset, Const(2));

            node.List.AddBefore(node, shrOp);

            return wordOffset;
        }

        private static bool IsCbOffset(Operand operand, int slot, int offset)
        {
            return operand.Type == OperandType.ConstantBuffer && operand.GetCbufSlot() == slot && operand.GetCbufOffset() == offset;
        }

        private static void ReplaceAddressAlignment(LinkedList<INode> list, Operand address, Operand byteOffset, int constantOffset)
        {
            // When we emit 16/8-bit LDG, we add extra code to determine the address alignment.
            // Eliminate the storage buffer base address from this too, leaving only the byte offset.

            foreach (INode useNode in address.UseOps)
            {
                if (useNode is Operation op && op.Inst == Instruction.BitwiseAnd)
                {
                    Operand src1 = op.GetSource(0);
                    Operand src2 = op.GetSource(1);

                    int addressIndex = -1;

                    if (src1 == address && src2.Type == OperandType.Constant && src2.Value == 3)
                    {
                        addressIndex = 0;
                    }
                    else if (src2 == address && src1.Type == OperandType.Constant && src1.Value == 3)
                    {
                        addressIndex = 1;
                    }

                    if (addressIndex != -1)
                    {
                        LinkedListNode<INode> node = list.Find(op);

                        // Add offset calculation before the use. Needs to be on the same block.
                        if (node != null)
                        {
                            Operand offset = Local();
                            Operation addOp = new Operation(Instruction.Add, offset, byteOffset, Const(constantOffset));
                            list.AddBefore(node, addOp);

                            op.SetSource(addressIndex, offset);
                        }
                    }
                }
            }
        }

        private static (Operand, int) GetStorageOffset(BasicBlock block, Operand address, int cbSlot, int baseAddressCbOffset)
        {
            if (IsCbOffset(address, cbSlot, baseAddressCbOffset))
            {
                // Direct offset: zero.
                return (Const(0), 0);
            }

            (address, int constantOffset) = GetStorageConstantOffset(block, address);

            address = Utils.FindLastOperation(address, block);

            if (IsCbOffset(address, cbSlot, baseAddressCbOffset))
            {
                // Only constant offset
                return (Const(0), constantOffset);
            }

            if (!(address.AsgOp is Operation offsetAdd) || offsetAdd.Inst != Instruction.Add)
            {
                return (null, 0);
            }

            Operand src1 = offsetAdd.GetSource(0);
            Operand src2 = Utils.FindLastOperation(offsetAdd.GetSource(1), block);

            if (IsCbOffset(src2, cbSlot, baseAddressCbOffset))
            {
                return (src1, constantOffset);
            }
            else if (IsCbOffset(src1, cbSlot, baseAddressCbOffset))
            {
                return (src2, constantOffset);
            }

            return (null, 0);
        }

        private static (Operand, int) GetStorageConstantOffset(BasicBlock block, Operand address)
        {
            if (!(address.AsgOp is Operation offsetAdd) || offsetAdd.Inst != Instruction.Add)
            {
                return (address, 0);
            }

            Operand src1 = offsetAdd.GetSource(0);
            Operand src2 = offsetAdd.GetSource(1);

            if (src2.Type != OperandType.Constant)
            {
                return (address, 0);
            }

            return (src1, src2.Value);
        }

        private static LinkedListNode<INode> ReplaceLdgWithLdc(LinkedListNode<INode> node, ShaderConfig config, int storageIndex)
        {
            Operation operation = (Operation)node.Value;

            Operand GetCbufOffset()
            {
                Operand addrLow = operation.GetSource(0);

                Operand baseAddrLow = Cbuf(0, UbeBaseOffset + storageIndex * StorageDescSize);

                Operand baseAddrTrunc = Local();

                Operand alignMask = Const(-config.GpuAccessor.QueryHostStorageBufferOffsetAlignment());

                Operation andOp = new Operation(Instruction.BitwiseAnd, baseAddrTrunc, baseAddrLow, alignMask);

                node.List.AddBefore(node, andOp);

                Operand byteOffset = Local();
                Operand wordOffset = Local();

                Operation subOp = new Operation(Instruction.Subtract, byteOffset, addrLow, baseAddrTrunc);
                Operation shrOp = new Operation(Instruction.ShiftRightU32, wordOffset, byteOffset, Const(2));

                node.List.AddBefore(node, subOp);
                node.List.AddBefore(node, shrOp);

                return wordOffset;
            }

            Operand[] sources = new Operand[operation.SourcesCount];

            int cbSlot = UbeFirstCbuf + storageIndex;

            sources[0] = Const(cbSlot);
            sources[1] = GetCbufOffset();

            config.SetUsedConstantBuffer(cbSlot);

            for (int index = 2; index < operation.SourcesCount; index++)
            {
                sources[index] = operation.GetSource(index);
            }

            Operation ldcOp = new Operation(Instruction.LoadConstant, operation.Dest, sources);

            for (int index = 0; index < operation.SourcesCount; index++)
            {
                operation.SetSource(index, null);
            }

            LinkedListNode<INode> oldNode = node;

            node = node.List.AddBefore(node, ldcOp);

            node.List.Remove(oldNode);

            return node;
        }

        private static SearchResult SearchForStorageBase(ShaderConfig config, BasicBlock block, Operand globalAddress)
        {
            globalAddress = Utils.FindLastOperation(globalAddress, block);

            if (globalAddress.Type == OperandType.ConstantBuffer)
            {
                return GetStorageIndex(config, globalAddress);
            }

            Operation operation = globalAddress.AsgOp as Operation;

            if (operation == null || operation.Inst != Instruction.Add)
            {
                return SearchResult.NotFound;
            }

            Operand src1 = operation.GetSource(0);
            Operand src2 = operation.GetSource(1);

            if ((src1.Type == OperandType.LocalVariable && src2.Type == OperandType.Constant) ||
                (src2.Type == OperandType.LocalVariable && src1.Type == OperandType.Constant))
            {
                Operand baseAddr;

                if (src1.Type == OperandType.LocalVariable)
                {
                    baseAddr = Utils.FindLastOperation(src1, block);
                }
                else
                {
                    baseAddr = Utils.FindLastOperation(src2, block);
                }

                var result = GetStorageIndex(config, baseAddr);
                if (result.Found)
                {
                    return result;
                }

                operation = baseAddr.AsgOp as Operation;

                if (operation == null || operation.Inst != Instruction.Add)
                {
                    return SearchResult.NotFound;
                }
            }

            var selectedResult = SearchResult.NotFound;

            for (int index = 0; index < operation.SourcesCount; index++)
            {
                Operand source = operation.GetSource(index);

                var result = GetStorageIndex(config, source);

                // If we already have a result, we give preference to the ones from
                // the driver reserved constant buffer, as those are the ones that
                // contains the base address.
                if (result.Found && (!selectedResult.Found || result.SbCbSlot == GlobalMemory.DriverReservedCb))
                {
                    selectedResult = result;
                }
            }

            return selectedResult;
        }

        private static SearchResult GetStorageIndex(ShaderConfig config, Operand operand)
        {
            if (operand.Type == OperandType.ConstantBuffer)
            {
                int slot = operand.GetCbufSlot();
                int offset = operand.GetCbufOffset();

                if ((offset & 3) == 0)
                {
                    return new SearchResult(slot, offset);
                }
            }

            return SearchResult.NotFound;
        }

        private static int GetStorageIndex(Operand operand, int sbStart, int sbEnd)
        {
            if (operand.Type == OperandType.ConstantBuffer)
            {
                int slot = operand.GetCbufSlot();
                int offset = operand.GetCbufOffset();

                if (slot == 0 && offset >= sbStart && offset < sbEnd)
                {
                    int storageIndex = (offset - sbStart) / StorageDescSize;

                    return storageIndex;
                }
            }

            return -1;
        }
    }
}