#define LOG_CALLS

using System;
using System.Runtime.InteropServices;

namespace VisualStudioMock
{
  delegate int GetInterfaceDelegate(IntPtr @this, [MarshalAs(UnmanagedType.LPStruct)] Guid iid, out IntPtr ppv);
  delegate int AddRefDelegate(IntPtr @this);
  delegate int ReleaseDelegate(IntPtr @this);

  sealed class ComProxy
  {
    private const int S_OK = 0;
    private const int E_NOINTERFACE = -2147467262;

    private static readonly Guid IID_IUnknown = new Guid("00000000-0000-0000-c000-000000000046");
    private static readonly Guid IID_IManagedObject = new Guid("C3FCC19E-A970-11D2-8B5A-00A0C9B7C9C4");

    [DllImport("kernel32.dll", SetLastError = true)]
    private static extern IntPtr VirtualAlloc(IntPtr lpAddress, UIntPtr dwSize, AllocationType lAllocationType, MemoryProtection flProtect);

    [DllImport("KERNEL32.dll", ExactSpelling = true, SetLastError = true)]
    private static extern bool VirtualProtect(IntPtr lpAddress, UIntPtr dwSize, PAGE_PROTECTION_FLAGS flNewProtect, out PAGE_PROTECTION_FLAGS lpflOldProtect);

    [DllImport("KERNEL32.dll", ExactSpelling = true, SetLastError = true)]
    public static extern bool FlushInstructionCache(IntPtr hProcess, [Optional] IntPtr lpBaseAddress, IntPtr dwSize);

    private static IntPtr _execBlockBase;
    private static int _execBlockNext;
    private static int _execBlockSize;

    private readonly object _realInstance;
    private readonly IntPtr _realInstanceN;
    private readonly IntPtr _proxyInstanceN;
    private readonly GetInterfaceDelegate _getInterface;
    private readonly IntPtr _getInterfaceN;
    private readonly AddRefDelegate _addRef;
    private readonly ReleaseDelegate _release;

    internal ComProxy(object instance)
    {
      this._realInstance = instance;
      this._realInstanceN = Marshal.GetIUnknownForObject(instance);
      this._proxyInstanceN = Marshal.AllocHGlobal(IntPtr.Size * 4);

      this._getInterfaceN = Marshal.GetFunctionPointerForDelegate((Delegate)(this._getInterface = new GetInterfaceDelegate(GetInterface)));

      IntPtr vtbl = _proxyInstanceN + IntPtr.Size;
      Marshal.WriteIntPtr(_proxyInstanceN, vtbl);
      Marshal.WriteIntPtr(vtbl, this._getInterfaceN = Marshal.GetFunctionPointerForDelegate((Delegate)(this._getInterface = new GetInterfaceDelegate(GetInterface))));
      Marshal.WriteIntPtr(vtbl, IntPtr.Size * 1, Marshal.GetFunctionPointerForDelegate((Delegate)(this._addRef = new AddRefDelegate(AddRef))));
      Marshal.WriteIntPtr(vtbl, IntPtr.Size * 2, Marshal.GetFunctionPointerForDelegate((Delegate)(this._release = new ReleaseDelegate(Release))));
    }

    internal object GetIUnknown()
    {
      return Marshal.GetObjectForIUnknown(_proxyInstanceN);
    }

    private int GetInterface(IntPtr @this, [MarshalAs(UnmanagedType.LPStruct)] Guid iid, out IntPtr ppv)
    {
      if (iid == IID_IUnknown)
      {
        ppv = _proxyInstanceN;
#if LOG_CALLS
        Ext.LogCall(_realInstance.GetType().Name, nameof(GetInterface), "IUnknown OK");
#endif
        return S_OK;
      }

      if (iid == IID_IManagedObject)
      {
        ppv = IntPtr.Zero;
#if LOG_CALLS
        Ext.LogCall(_realInstance.GetType().Name, nameof(GetInterface), "IManagedObject Fail");
#endif
        return E_NOINTERFACE;
      }

      int result = Marshal.QueryInterface(_realInstanceN, ref iid, out ppv);
      if (result == S_OK)
      {
        // Patch the vtable with our own thunks
        IntPtr origVtbl = Marshal.ReadIntPtr(ppv);

        Type iType = null;
        foreach (Type i in _realInstance.GetType().GetInterfaces())
          if (i.GUID == iid)
          {
            iType = i;
            break;
          }

        if (iType != null)
        {
          // Generate proper thunks in machine code
          int slotCount = Marshal.GetEndComSlot(iType) + 1;
          int slotsSizeBytes = IntPtr.Size * slotCount;
          IntPtr newPpv = Marshal.AllocHGlobal(slotsSizeBytes + IntPtr.Size);
          IntPtr newVtbl = newPpv + IntPtr.Size;
          Marshal.WriteIntPtr(newPpv, newVtbl);
          Marshal.WriteIntPtr(newVtbl, _getInterfaceN);
          for (int i = 1; i < slotCount; i++)
            Marshal.WriteIntPtr(newVtbl, i * IntPtr.Size, CreateThunk(ppv, Marshal.ReadIntPtr(origVtbl, i * IntPtr.Size)));
          FlushThunkCode();
          ppv = newPpv;
        }

#if LOG_CALLS
        Ext.LogCall(_realInstance.GetType().Name, nameof(GetInterface), (iType == null ? iid.ToString() : iType.Name) + " OK");
#endif
      }
      else
      {
#if LOG_CALLS
        Ext.LogCall(_realInstance.GetType().Name, nameof(GetInterface), iid.ToString() + " Fail");
#endif
      }
      return result;
    }

    private int AddRef(IntPtr @this)
    {
      return Marshal.AddRef(_realInstanceN);
    }

    private int Release(IntPtr @this)
    {
      return Marshal.Release(_realInstanceN);
    }

    private static IntPtr CreateThunk(IntPtr @this, IntPtr jmpTarget)
    {
      int allocSize = Environment.Is64BitProcess ? 22 : 13;

      if (_execBlockSize - _execBlockNext < allocSize)
      {
        // Flush the now complete page
        if (_execBlockBase != IntPtr.Zero)
          FlushThunkCode();

        // Prepare the memory for the machine code
        _execBlockSize = Environment.SystemPageSize;
        _execBlockBase = VirtualAlloc(IntPtr.Zero, new UIntPtr((uint)_execBlockSize), AllocationType.Commit | AllocationType.Reserve, MemoryProtection.ReadWrite);
        if (!VirtualProtect(_execBlockBase, new UIntPtr((uint)_execBlockSize), PAGE_PROTECTION_FLAGS.PAGE_EXECUTE_READWRITE, out _))
          throw new InvalidOperationException("Failed to make thunk data executable");

        _execBlockNext = 0;

#if LOG_CALLS
        Ext.Log("Allocated thunk page at 0x" + _execBlockBase.ToString("X"));
#endif
      }

#if LOG_CALLS
      Ext.LogCall(nameof(ComProxy), nameof(CreateThunk), "This = 0x" + @this.ToString("X") + " Target = 0x" + jmpTarget.ToString("X"));
#endif

      IntPtr result = _execBlockBase + _execBlockNext;
      if (Environment.Is64BitProcess)
      {
        // movabs rcx, this
        Marshal.WriteInt16(_execBlockBase, _execBlockNext, unchecked((short)(0xB948)));
        Marshal.WriteIntPtr(_execBlockBase, _execBlockNext + 2, @this);
        _execBlockNext += 2 + IntPtr.Size;

        // movabs rax, jmpTarget
        // jmp rax
        Marshal.WriteInt16(_execBlockBase, _execBlockNext, unchecked((short)(0xB848)));
        Marshal.WriteIntPtr(_execBlockBase, _execBlockNext + 2, jmpTarget);
        Marshal.WriteInt16(_execBlockBase, _execBlockNext + 10, unchecked((short)(0xE0FF)));
        _execBlockNext += 2 + IntPtr.Size + 2;
      }
      else
      {
        // NOT tested!
        // pop eax
        // push this
        Marshal.WriteByte(_execBlockBase, _execBlockNext, 0x58);
        Marshal.WriteByte(_execBlockBase, _execBlockNext + 1, 0x68);
        Marshal.WriteIntPtr(_execBlockBase, _execBlockNext + 2, @this);
        _execBlockNext += 1 + 1 + IntPtr.Size;

        // mov eax target_address
        // jmp eax
        Marshal.WriteByte(_execBlockBase, _execBlockNext, 0xB8);
        Marshal.WriteIntPtr(_execBlockBase, _execBlockNext + 1, jmpTarget);
        Marshal.WriteInt16(_execBlockBase, _execBlockNext + 5, unchecked((short)(0xE0FF)));
        _execBlockNext += 1 + IntPtr.Size + 2;
      }
      return result;
    }

    private static void FlushThunkCode()
    {
      FlushInstructionCache(new IntPtr(-1), _execBlockBase, new IntPtr(_execBlockSize));
    }

        [Flags]
    enum AllocationType
    {
      Commit = 0x1000,
      Reserve = 0x2000,
      Decommit = 0x4000,
      Release = 0x8000,
      Reset = 0x80000,
      Physical = 0x400000,
      TopDown = 0x100000,
      WriteWatch = 0x200000,
      LargePages = 0x20000000
    }

    [Flags]
    enum MemoryProtection
    {
      Execute = 0x10,
      ExecuteRead = 0x20,
      ExecuteReadWrite = 0x40,
      ExecuteWriteCopy = 0x80,
      NoAccess = 0x01,
      ReadOnly = 0x02,
      ReadWrite = 0x04,
      WriteCopy = 0x08,
      GuardModifierflag = 0x100,
      NoCacheModifierflag = 0x200,
      WriteCombineModifierflag = 0x400
    }

    [Flags]
    enum PAGE_PROTECTION_FLAGS : uint
    {
      PAGE_NOACCESS = 0x00000001,
      PAGE_READONLY = 0x00000002,
      PAGE_READWRITE = 0x00000004,
      PAGE_WRITECOPY = 0x00000008,
      PAGE_EXECUTE = 0x00000010,
      PAGE_EXECUTE_READ = 0x00000020,
      PAGE_EXECUTE_READWRITE = 0x00000040,
      PAGE_EXECUTE_WRITECOPY = 0x00000080,
      PAGE_GUARD = 0x00000100,
      PAGE_NOCACHE = 0x00000200,
      PAGE_WRITECOMBINE = 0x00000400,
      PAGE_GRAPHICS_NOACCESS = 0x00000800,
      PAGE_GRAPHICS_READONLY = 0x00001000,
      PAGE_GRAPHICS_READWRITE = 0x00002000,
      PAGE_GRAPHICS_EXECUTE = 0x00004000,
      PAGE_GRAPHICS_EXECUTE_READ = 0x00008000,
      PAGE_GRAPHICS_EXECUTE_READWRITE = 0x00010000,
      PAGE_GRAPHICS_COHERENT = 0x00020000,
      PAGE_GRAPHICS_NOCACHE = 0x00040000,
      PAGE_ENCLAVE_THREAD_CONTROL = 0x80000000,
      PAGE_REVERT_TO_FILE_MAP = 0x80000000,
      PAGE_TARGETS_NO_UPDATE = 0x40000000,
      PAGE_TARGETS_INVALID = 0x40000000,
      PAGE_ENCLAVE_UNVALIDATED = 0x20000000,
      PAGE_ENCLAVE_MASK = 0x10000000,
      PAGE_ENCLAVE_DECOMMIT = 0x10000000,
      PAGE_ENCLAVE_SS_FIRST = 0x10000001,
      PAGE_ENCLAVE_SS_REST = 0x10000002,
      SEC_PARTITION_OWNER_HANDLE = 0x00040000,
      SEC_64K_PAGES = 0x00080000,
      SEC_FILE = 0x00800000,
      SEC_IMAGE = 0x01000000,
      SEC_PROTECTED_IMAGE = 0x02000000,
      SEC_RESERVE = 0x04000000,
      SEC_COMMIT = 0x08000000,
      SEC_NOCACHE = 0x10000000,
      SEC_WRITECOMBINE = 0x40000000,
      SEC_LARGE_PAGES = 0x80000000,
      SEC_IMAGE_NO_EXECUTE = 0x11000000,
    }
  }
}