 extension ChatViewController: UICollectionViewDelegateFlowLayout {
     
     func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, sizeForItemAt indexPath: IndexPath) -> CGSize {
-        otherMessageCellPrototype.messageLabel.text = messages[indexPath.row]
-        otherMessageCellPrototype.setNeedsLayout()
-        otherMessageCellPrototype.layoutIfNeeded()
-        let width = collectionView.frame.width
-        let height = otherMessageCellPrototype.containerView.frame.maxY + LayoutDimension().spacing * 0.5
-        return CGSize(width: width, height: height)
+        let item = messages[indexPath.row]
+        
+        if item.isMe {
+            myMessageCellPrototype.messageLabel.text = item.message
+            myMessageCellPrototype.setNeedsLayout()
+            myMessageCellPrototype.layoutIfNeeded()
+            let width = collectionView.frame.width
+            let height = myMessageCellPrototype.messageLabel.frame.maxY + LayoutDimension().spacing * 0.5
+            return CGSize(width: width, height: height)
+            
+        } else {
+            otherMessageCellPrototype.messageLabel.text = item.message
+            otherMessageCellPrototype.setNeedsLayout()
+            otherMessageCellPrototype.layoutIfNeeded()
+            let width = collectionView.frame.width
+            let height = otherMessageCellPrototype.containerView.frame.maxY + LayoutDimension().spacing * 0.5
+            return CGSize(width: width, height: height)
+        }
     }
     
     func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, minimumLineSpacingForSectionAt section: Int) -> CGFloat {

    }
    
    func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
        let cell = ChatOtherMessageCell.dequeue(from: collectionView, at: indexPath)
        cell.messageLabel.text = messages[indexPath.row]
        cell.setNeedsLayout()
        cell.layoutIfNeeded()
        return cell
    }
}

        chatView.collectionView.dataSource = self
        
        ChatOtherMessageCell.register(in: chatView.collectionView)
        ChatMyMessageCell.register(in: chatView.collectionView)
        
        otherMessageCellPrototype.frame.size.width = size.width
        myMessageCellPrototype.frame.size.width = size.width
        
        view = chatView
    }
