Adding a click event listener on WKWebView not working

I am working on an app that is using a WKWebView and I am trying to add an event listener when the user clicks on a button on the web page. The button as no id so I am working with the class.

I am able to set it up properly on chrome using the inspect console:

For my Swift development I have tried to add a function inside a WKWebView extension so it can be reused like that:

extension WKWebView {
    /// Adds a event listener that will be call on WKScriptMessageHandler - didReceiveMessage
    /// - Parameters:
    ///   - elementID: The name of the element
    ///   - callbackID: The ID for the callback
    ///   - elementType: The type of element to get
    ///   - completion: Callback triggered went script has been appended to WKWebView
    //NOT WORKING
    func addEventListener(id: String, callbackID: String, elementType: ElementType, handler: WKScriptMessageHandler, completion: ((Error?)->Void)?) {
        let scriptString: String
        
        switch elementType {
        case .id:
            scriptString = """
                document.getElementById('\(id)').addEventListener('click', function(){
                    webkit.messageHandlers.refreshWebPage.postMessage("status":"ok");
                });
                """
        case .class:
            scriptString = """
                document.getElementsByClassName('\(id)')[0].addEventListener('click', function(){
                    webkit.messageHandlers.refreshWebPage.postMessage("status":"ok");
                });
                """
        }
        
        let script = WKUserScript(source: scriptString, injectionTime: .atDocumentStart, forMainFrameOnly: false)
        configuration.userContentController.addUserScript(script)
        configuration.userContentController.add(handler, name: "postMessage")
    }
} 

Then, on the UIViewController, that contains the WKWebView:

//MARK: WKNavigationDelegate implementation
extension ArtpadWebViewController: WKNavigationDelegate {
    func webView(_ webView: WKWebView, didFinish navigation: WKNavigation!) {
        webView.addEventListener(id: "top_left_menu", callbackID: backCallbackID, elementType: .class, handler: self, completion: nil)
    }
}

//MARK: WKScriptMessageHandler implementation
extension ArtpadWebViewController: WKScriptMessageHandler {
    func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) {
        print("HO!")
    }
}

However, the WKScriptMessageHandler function is not called... I have been trying to different solutions I found but non of them has worked. What am I missing?

Thank you

For some reason, Everything started working after adding the code after evaluating that javascript. I am still unable to know why.

Here is the code for the full extension so far:

extension WKWebView {
    /// Type of HTML element to get from DOM
    enum ElementType {
        /// ID Element
        case id
        /// Class element
        case `class`
    }
    
    /// List of errors for WKWebView injection
    enum InjectionError: Error {
        /// The Listener is already added
        case listenerAlreadyAdded
    }
    
    /// Cahnges the CSS Visibiltiy
    /// - Parameters:
    ///   - elementID: The name of the element
    ///   - isVisible: Wether or not is visible
    ///   - elementType: The type of element to get
    ///   - completion: Callback triggered went script has been appended to WKWebView
    func changeCSSVisibility(elementID: String, isVisible: Bool, elementType: ElementType, completion: ((Error?)->Void)?) {
        let script: String
        
        switch elementType {
        case .id:
            script = "document.getElementById('\(elementID)').style.visibility = \(isVisible ? "'visible'":"'hidden'");"
        case .class:
        script =
            """
            [].forEach.call(document.querySelectorAll('.\(elementID)'), function (el) {
              el.style.visibility = \(isVisible ? "'visible'":"'hidden'");
            });
            """
        }
        
        evaluateJavaScript(script, completionHandler: { (result, error) in
            completion?(error)
        })
    }
    
    /// Adds a event listener that will be call on WKScriptMessageHandler - didReceiveMessage
    /// - Parameters:
    ///   - elementID: The name of the element
    ///   - callbackID: The ID for the callback
    ///   - elementType: The type of element to get
    ///   - completion: Callback triggered went script has been appended to WKWebView
    func addEventListener(elementID: String, callbackID: String, elementType: ElementType, handler: WKScriptMessageHandler, completion: ((Error?)->Void)?) {
        let element: String
        
        switch elementType {
        case .id:
            element = "document.getElementById('\(elementID)')"
        case .class:
            element = "document.getElementsByClassName('\(elementID)')[0]"
        }
        
        let scriptString = """
            function callback () {
                console.log('\(callbackID) clicked!')
                window.webkit.messageHandlers.\(callbackID).postMessage({
                    message: 'WKWebView-onClickListener-\(callbackID)'
                });
            }

            \(element).addEventListener('click', callback);
            """
        
        if configuration.userContentController.userScripts.first(where: { $0.source == scriptString }) == nil {
            evaluateJavaScript(scriptString) { [weak self] (result, error) -> Void in
                guard let self = self else { return }
                
                if let error = error {
                    completion?(error)
                } else {
                    self.configuration.userContentController.removeScriptMessageHandler(forName: callbackID)
                    self.configuration.userContentController.add(handler, name: callbackID)
                    self.configuration.userContentController.addUserScript(WKUserScript(source: scriptString, injectionTime: .atDocumentEnd, forMainFrameOnly: false))
                }
            }
        } else {
            completion?(InjectionError.listenerAlreadyAdded)
        }
    }
}