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


    var chatView: ChatView!
    
    lazy var otherMessageCellPrototype = ChatOtherMessageCell()
    lazy var myMessageCellPrototype = ChatMyMessageCell()
    lazy var messages: [ChatDisplayData] = generateRandomChatDisplayItems()
    
    override func loadView() {

        chatView.collectionView.delegate = self
        chatView.collectionView.dataSource = self
        
        ChatMessageCell.register(in: chatView.collectionView)
        
        messageCellPrototype.frame.size.width = size.width
        
        view = chatView
    }
