// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using Microsoft.Diagnostics.DataContractReader.ExecutionManagerHelpers;

namespace Microsoft.Diagnostics.DataContractReader.Contracts;

internal partial class ExecutionManagerCore<T> : IExecutionManager
{
    private sealed class EEJitManager : JitManager
    {
        private readonly INibbleMap _nibbleMap;
        private readonly RuntimeFunctionLookup _runtimeFunctions;
        public EEJitManager(Target target, INibbleMap nibbleMap) : base(target)
        {
            _nibbleMap = nibbleMap;
            _runtimeFunctions = RuntimeFunctionLookup.Create(target);
        }

        public override bool GetMethodInfo(RangeSection rangeSection, TargetCodePointer jittedCodeAddress, [NotNullWhen(true)] out CodeBlock? info)
        {
            info = null;
            // EEJitManager::JitCodeToMethodInfo
            Debug.Assert(!rangeSection.IsRangeList);

            if (rangeSection.Data == null)
                throw new ArgumentException(nameof(rangeSection));

            TargetPointer codeStart = FindMethodCode(rangeSection, jittedCodeAddress);
            if (codeStart == TargetPointer.Null)
                return false;

            Debug.Assert(codeStart.Value <= jittedCodeAddress.Value);
            TargetNUInt relativeOffset = new TargetNUInt(jittedCodeAddress.Value - codeStart.Value);

            if (!GetRealCodeHeader(rangeSection, codeStart, out Data.RealCodeHeader? realCodeHeader))
                return false;

            info = new CodeBlock(codeStart, realCodeHeader.MethodDesc, relativeOffset, rangeSection.Data!.JitManager);
            return true;
        }

        public override void GetMethodRegionInfo(
            RangeSection rangeSection,
            TargetCodePointer jittedCodeAddress,
            out uint hotSize,
            out TargetPointer coldStart,
            out uint coldSize)
        {
            // cold regions are not supported for JITted code
            coldStart = TargetPointer.Null;
            coldSize = 0;

            IGCInfo gcInfo = Target.Contracts.GCInfo;
            GetGCInfo(rangeSection, jittedCodeAddress, out TargetPointer pGcInfo, out uint gcVersion);
            IGCInfoHandle gcInfoHandle = gcInfo.DecodePlatformSpecificGCInfo(pGcInfo, gcVersion);
            hotSize = gcInfo.GetCodeLength(gcInfoHandle);
            Debug.Assert(hotSize > 0);
        }

        public override TargetPointer GetUnwindInfo(RangeSection rangeSection, TargetCodePointer jittedCodeAddress)
        {
            Debug.Assert(!rangeSection.IsRangeList);
            if (rangeSection.Data == null)
                throw new ArgumentException(nameof(rangeSection));

            TargetPointer codeStart = FindMethodCode(rangeSection, jittedCodeAddress);
            if (codeStart == TargetPointer.Null)
                return TargetPointer.Null;
            Debug.Assert(codeStart.Value <= jittedCodeAddress.Value);

            if (!GetRealCodeHeader(rangeSection, codeStart, out Data.RealCodeHeader? realCodeHeader))
                return TargetPointer.Null;

            if (realCodeHeader.NumUnwindInfos == 0)
            {
                return TargetPointer.Null;
            }

            // Find the relative address that we are looking for
            TargetPointer addr = CodePointerUtils.AddressFromCodePointer(jittedCodeAddress, Target);
            TargetPointer imageBase = rangeSection.Data.RangeBegin;
            TargetPointer relativeAddr = addr - imageBase;

            if (!_runtimeFunctions.TryGetRuntimeFunctionIndexForAddress(realCodeHeader.UnwindInfos, realCodeHeader.NumUnwindInfos, relativeAddr, out uint index))
                return TargetPointer.Null;

            return _runtimeFunctions.GetRuntimeFunctionAddress(realCodeHeader.UnwindInfos, index);
        }

        public override TargetPointer GetDebugInfo(RangeSection rangeSection, TargetCodePointer jittedCodeAddress, out bool hasFlagByte)
        {
            hasFlagByte = false;
            Debug.Assert(!rangeSection.IsRangeList);
            if (rangeSection.Data == null)
                throw new ArgumentException(nameof(rangeSection));

            TargetPointer codeStart = FindMethodCode(rangeSection, jittedCodeAddress);
            if (codeStart == TargetPointer.Null)
                return TargetPointer.Null;
            Debug.Assert(codeStart.Value <= jittedCodeAddress.Value);

            if (!GetRealCodeHeader(rangeSection, codeStart, out Data.RealCodeHeader? realCodeHeader))
                return TargetPointer.Null;

            bool featureOnStackReplacement = Target.ReadGlobal<byte>(Constants.Globals.FeatureOnStackReplacement) != 0;
            Data.EEJitManager eeJitManager = Target.ProcessedData.GetOrAdd<Data.EEJitManager>(rangeSection.Data.JitManager);
            if (featureOnStackReplacement || eeJitManager.StoreRichDebugInfo)
                hasFlagByte = true;

            return realCodeHeader.DebugInfo;
        }

        public override CodeKind GetCodeKind(RangeSection rangeSection, TargetCodePointer codeAddress)
        {
            TargetPointer startAddr = FindMethodCode(rangeSection, codeAddress); // validate that the code address is within the method's code range
            if (startAddr == TargetPointer.Null)
                return CodeKind.Unknown;
            return GetCodeHeaderStubKind(rangeSection, startAddr);
        }

        public override void GetGCInfo(RangeSection rangeSection, TargetCodePointer jittedCodeAddress, out TargetPointer gcInfo, out uint gcVersion)
        {
            gcInfo = TargetPointer.Null;
            gcVersion = 0;

            // EEJitManager::GetGCInfoToken
            Debug.Assert(!rangeSection.IsRangeList);

            if (rangeSection.Data == null)
                throw new ArgumentException(nameof(rangeSection));

            TargetPointer codeStart = FindMethodCode(rangeSection, jittedCodeAddress);
            if (codeStart == TargetPointer.Null)
                return;
            Debug.Assert(codeStart.Value <= jittedCodeAddress.Value);

            if (!GetRealCodeHeader(rangeSection, codeStart, out Data.RealCodeHeader? realCodeHeader))
                return;

            gcVersion = Target.ReadGlobal<uint>(Constants.Globals.GCInfoVersion);
            gcInfo = realCodeHeader.GCInfo;
        }

        private TargetPointer FindMethodCode(RangeSection rangeSection, TargetCodePointer codeAddress)
        {
            // EEJitManager::FindMethodCode
            Debug.Assert(rangeSection.Data != null);

            if (!rangeSection.IsCodeHeap)
                throw new InvalidOperationException("RangeSection is not a code heap");

            TargetPointer heapListAddress = rangeSection.Data.HeapList;
            Data.CodeHeapListNode heapListNode = Target.ProcessedData.GetOrAdd<Data.CodeHeapListNode>(heapListAddress);
            return _nibbleMap.FindMethodCode(heapListNode, codeAddress);
        }

        private TargetPointer GetCodeHeaderAddress(RangeSection rangeSection, TargetPointer codeStart)
        {
            // EEJitManager::JitCodeToMethodInfo
            Debug.Assert(!rangeSection.IsRangeList);

            if (rangeSection.Data == null)
                throw new ArgumentException(nameof(rangeSection));

            // See EEJitManager::GetCodeHeaderFromStartAddress in vm/codeman.h
            int codeHeaderOffset = Target.PointerSize;
            TargetPointer codeHeaderIndirect = new TargetPointer(codeStart - (ulong)codeHeaderOffset);
            return Target.ReadPointer(codeHeaderIndirect);
        }

        private bool GetRealCodeHeader(RangeSection rangeSection, TargetPointer codeStart, [NotNullWhen(true)] out Data.RealCodeHeader? realCodeHeader)
        {
            realCodeHeader = null;
            TargetPointer codeHeaderAddress = GetCodeHeaderAddress(rangeSection, codeStart);
            if (RangeSection.IsStubCodeBlock(Target, codeHeaderAddress))
            {
                return false;
            }
            realCodeHeader = Target.ProcessedData.GetOrAdd<Data.RealCodeHeader>(codeHeaderAddress);
            return true;
        }

        private CodeKind GetCodeHeaderStubKind(RangeSection rangeSection, TargetPointer codeStart)
        {
            TargetPointer codeHeaderAddress = GetCodeHeaderAddress(rangeSection, codeStart);
            if (RangeSection.IsStubCodeBlock(Target, codeHeaderAddress))
            {
                return GetStubKind((StubKind)codeHeaderAddress.Value);
            }
            return CodeKind.Jitted;
        }

        public override void GetExceptionClauses(RangeSection rangeSection, CodeBlockHandle codeInfoHandle, out TargetPointer startAddr, out TargetPointer endAddr)
        {
            startAddr = TargetPointer.Null;
            endAddr = TargetPointer.Null;

            if (rangeSection.Data == null)
                throw new ArgumentException(nameof(rangeSection));

            Data.RealCodeHeader? realCodeHeader;
            TargetPointer codeStart = FindMethodCode(rangeSection, new TargetCodePointer(codeInfoHandle.Address));
            if (codeStart == TargetPointer.Null)
                return;
            if (!GetRealCodeHeader(rangeSection, codeStart, out realCodeHeader) || realCodeHeader == null)
                return;

            if (realCodeHeader.EHInfo == TargetPointer.Null)
                return;

            Data.EEILException ehInfo = Target.ProcessedData.GetOrAdd<Data.EEILException>(realCodeHeader.EHInfo);
            TargetNUInt numEHInfos = Target.ReadNUInt(ehInfo.Address - (ulong)Target.PointerSize);
            startAddr = ehInfo.Clauses;
            endAddr = startAddr + numEHInfos.Value * Target.GetTypeInfo(DataType.EEExceptionClause).Size!.Value;
        }
    }
}