using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Text;

namespace ProfilerHelpers
{
    public class MethodStore : IDisposable
    {
        // JITed methods information (start address + size + signature)
        private readonly List<MethodInfo> _methods;
        
        // addresses from callstacks already matching (address -> full name)
        private readonly Dictionary<ulong, string> _cache;

        // for native methods, rely on dbghelp API
        // so the process handle is required
        private IntPtr _hProcess;
        private Process _process;
        private readonly int _pid;

        public MethodStore(int pid, bool loadModules = false)
        {
            // it may be possible to open the process
            // in that case, _hProcess = IntPtr.Zero
            _pid = pid;

            _methods = new List<MethodInfo>(1024);
            _cache = new Dictionary<ulong, string>();

            _hProcess = BindToProcess(pid, loadModules);
        }

        private IntPtr BindToProcess(int pid, bool loadModules)
        {
            try
            {
                _process = Process.GetProcessById(pid);

                if (!SymInitialize(_process.Handle, loadModules))
                    return IntPtr.Zero;

                return _process.Handle;
            }
            catch (Exception x)
            {
                Console.WriteLine($"Error while binding pid #{pid} to DbgHelp:");
                Console.WriteLine(x.Message);
                return IntPtr.Zero;
            }
        }

        private bool SymInitialize(IntPtr hProcess, bool loadModules = false)
        {
            // read https://docs.microsoft.com/en-us/windows/win32/api/dbghelp/nf-dbghelp-symsetoptions for more details
            // maybe SYMOPT_NO_PROMPTS and SYMOPT_FAIL_CRITICAL_ERRORS could be used
            NativeDbgHelp.SymSetOptions(
                NativeDbgHelp.SYMOPT_DEFERRED_LOADS |   // performance optimization
                NativeDbgHelp.SYMOPT_UNDNAME            // C++ names are not mangled
                );

            // https://docs.microsoft.com/en-us/windows/win32/api/dbghelp/nf-dbghelp-syminitialize
            // search path for symbols:
            //   - The current working directory of the application
            //   - The _NT_SYMBOL_PATH environment variable
            //   - The _NT_ALTERNATE_SYMBOL_PATH environment variable
            //
            // passing false as last parameter means that we will need to call SymLoadModule64 
            // each time a module is loaded in the process
            return NativeDbgHelp.SymInitialize(hProcess, null, loadModules);
        }

        public MethodInfo Add(ulong address, int size, string namespaceAndTypeName, string name, string signature)
        {
            var method = new MethodInfo(address, size, namespaceAndTypeName, name, signature);
            _methods.Add(method);
            return method;
        }

        public string GetFullName(ulong address)
        {
            if (_cache.TryGetValue(address, out var fullName))
                return fullName;
            
            // look for managed methods
            for (int i = 0; i < _methods.Count; i++)
            {
                var method = _methods[i];
                
                if ((address >= method.StartAddress) && (address < method.StartAddress + (ulong)method.Size))
                {
                    fullName = method.FullName;
                    _cache[address] = fullName;
                    return fullName;
                }
            }

            // look for native methods
            fullName = GetNativeMethodName(address);
            _cache[address] = fullName;

            return fullName;
        }

        private string GetNativeMethodName(ulong address)
        {
            var symbol = new NativeDbgHelp.SYMBOL_INFO();
            symbol.MaxNameLen = 1024;
            symbol.SizeOfStruct = (uint)Marshal.SizeOf(symbol) - 1024;   // char buffer is not counted
            // the ANSI version of SymFromAddr is called so each character is 1 byte long

            if (NativeDbgHelp.SymFromAddr(_hProcess, address, out var displacement, ref symbol))
            {
                var buffer = new StringBuilder(symbol.Name.Length);

                // remove weird "$##" at the end of some symbols
                var pos = symbol.Name.LastIndexOf("$##");
                if (pos == -1)
                    buffer.Append(symbol.Name);
                else
                    buffer.Append(symbol.Name, 0, pos);

                // add offset if any
                if (displacement != 0)
                    buffer.Append($"+0x{displacement}");

                return buffer.ToString();
            }

            // default value is just the address in HEX
#if DEBUG
            return ($"0x{address:x}  (SymFromAddr failed with 0x{Marshal.GetLastWin32Error():x})");
#else
            return $"0x{address:x}";
#endif

        }

        const int ERROR_SUCCESS = 0;
        public void AddModule(string filename, ulong baseOfDll, int sizeOfDll)
        {
            var baseAddress = NativeDbgHelp.SymLoadModule64(_hProcess, IntPtr.Zero, filename, null, baseOfDll, (uint)sizeOfDll);
            if (baseAddress == 0)
            {
                // should work if the same module is added more than once
                if (Marshal.GetLastWin32Error() == ERROR_SUCCESS) return;

                Console.WriteLine($"SymLoadModule64 failed for {filename}");
            }
        }

        public void Dispose()
        {
            if (_hProcess == IntPtr.Zero)
                return;
            _hProcess = IntPtr.Zero;

            _process.Dispose();
        }
    }

    internal static class NativeDbgHelp
    {
        // from C:\Program Files (x86)\Windows Kits\10\Debuggers\inc\dbghelp.h
        public const uint SYMOPT_UNDNAME = 0x00000002;
        public const uint SYMOPT_DEFERRED_LOADS = 0x00000004;

        [StructLayout(LayoutKind.Sequential)]
        public struct SYMBOL_INFO
        {
            public uint SizeOfStruct;
            public uint TypeIndex;      // Type Index of symbol
            private ulong Reserved1;
            private ulong Reserved2;
            public uint Index;
            public uint Size;
            public ulong ModBase;       // Base Address of module containing this symbol
            public uint Flags;
            public ulong Value;         // Value of symbol, ValuePresent should be 1
            public ulong Address;       // Address of symbol including base address of module
            public uint Register;       // register holding value or pointer to value
            public uint Scope;          // scope of the symbol
            public uint Tag;            // pdb classification
            public uint NameLen;        // Actual length of name
            public uint MaxNameLen;
            [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 1024)]
            public string Name;
        }

        [DllImport("dbghelp.dll", SetLastError = true)]
        public static extern bool SymInitialize(IntPtr hProcess, string userSearchPath, bool invadeProcess);

        [DllImport("dbghelp.dll", SetLastError = true)]
        public static extern uint SymSetOptions(uint symOptions);

        [DllImport("dbghelp.dll", SetLastError = true, CharSet = CharSet.Ansi)]
        public static extern ulong SymLoadModule64(IntPtr hProcess, IntPtr hFile, string imageName, string moduleName, ulong baseOfDll, uint sizeOfDll);

        // use ANSI version to ensure the right size of the structure 
        // read https://docs.microsoft.com/en-us/windows/win32/api/dbghelp/ns-dbghelp-symbol_info
        [DllImport("dbghelp.dll", SetLastError = true, CharSet = CharSet.Ansi)]
        public static extern bool SymFromAddr(IntPtr hProcess, ulong address, out ulong displacement, ref SYMBOL_INFO symbol);

        [DllImport("dbghelp.dll", SetLastError = true)]
        public static extern bool SymCleanup(IntPtr hProcess);
    }
}