From b5e74291d646cc418bf427212f4c16e3f103ac21 Mon Sep 17 00:00:00 2001
From: Toni Sevener <tsevener@wikimedia.org>
Date: Tue, 9 Jan 2024 16:19:08 -0600
Subject: [PATCH] Bug fixes after merge

---
 .../Input Views/WKEditorInputView.swift       | 45 ++++++++++++++++---
 ...EditorFormatterHeading+ButtonActions.swift |  2 +-
 .../WKSourceEditorViewController.swift        |  4 +-
 3 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/Components/Sources/Components/Components/Editors/Common Views/Input Views/WKEditorInputView.swift b/Components/Sources/Components/Components/Editors/Common Views/Input Views/WKEditorInputView.swift
index 9e6fead4fc2..3b14344aa86 100644
--- a/Components/Sources/Components/Components/Editors/Common Views/Input Views/WKEditorInputView.swift	
+++ b/Components/Sources/Components/Components/Editors/Common Views/Input Views/WKEditorInputView.swift	
@@ -6,6 +6,7 @@ protocol WKEditorInputViewDelegate: AnyObject {
     func didTapBold(isSelected: Bool)
     func didTapItalics(isSelected: Bool)
     func didTapTemplate(isSelected: Bool)
+    func didTapHeading(type: WKEditorInputView.HeadingButtonType)
     func didTapStrikethrough(isSelected: Bool)
 }
 
@@ -221,6 +222,18 @@ class WKEditorInputView: WKComponentView {
         ])
         
         updateColors()
+        
+        NotificationCenter.default.addObserver(self, selector: #selector(updateButtonSelectionState(_:)), name: Notification.WKSourceEditorSelectionState, object: nil)
+    }
+    
+    // MARK: - Notifications
+        
+    @objc private func updateButtonSelectionState(_ notification: NSNotification) {
+        guard let selectionState = notification.userInfo?[Notification.WKSourceEditorSelectionStateKey] as? WKSourceEditorSelectionState else {
+            return
+        }
+
+        configure(selectionState: selectionState)
     }
     
     // MARK: - Overrides
@@ -283,20 +296,40 @@ class WKEditorInputView: WKComponentView {
             
             switch type {
             case .paragraph:
-                paragraphButton.isSelected.toggle()
+                paragraphButton.isSelected = true
             case .heading:
-                headerButton.isSelected.toggle()
+                headerButton.isSelected = true
             case .subheading1:
-                subheader1Button.isSelected.toggle()
+                subheader1Button.isSelected = true
             case .subheading2:
-                subheader2Button.isSelected.toggle()
+                subheader2Button.isSelected = true
             case .subheading3:
-                subheader3Button.isSelected.toggle()
+                subheader3Button.isSelected = true
             case .subheading4:
-                subheader4Button.isSelected.toggle()
+                subheader4Button.isSelected = true
             }
+            
+            delegate?.didTapHeading(type: type)
         })
         
         return UIButton(configuration: configuration, primaryAction: action)
     }
+    
+    func configure(selectionState: WKSourceEditorSelectionState) {
+        headingButtons.forEach { $0.isSelected = false }
+
+        if selectionState.isHeading {
+            headerButton.isSelected = true
+        } else if selectionState.isSubheading1 {
+            subheader1Button.isSelected = true
+        } else if selectionState.isSubheading2 {
+            subheader2Button.isSelected = true
+        } else if selectionState.isSubheading3 {
+            subheader3Button.isSelected = true
+        } else if selectionState.isSubheading4 {
+            subheader4Button.isSelected = true
+        } else {
+            paragraphButton.isSelected = true
+        }
+    }
 }
diff --git a/Components/Sources/Components/Components/Editors/Source Editor/Formatter Extensions/WKSourceEditorFormatterHeading+ButtonActions.swift b/Components/Sources/Components/Components/Editors/Source Editor/Formatter Extensions/WKSourceEditorFormatterHeading+ButtonActions.swift
index 123e10f1903..869c6502bcd 100644
--- a/Components/Sources/Components/Components/Editors/Source Editor/Formatter Extensions/WKSourceEditorFormatterHeading+ButtonActions.swift	
+++ b/Components/Sources/Components/Components/Editors/Source Editor/Formatter Extensions/WKSourceEditorFormatterHeading+ButtonActions.swift	
@@ -2,7 +2,7 @@ import Foundation
 import ComponentsObjC
 
 extension WKSourceEditorFormatterHeading {
-    func toggleHeadingFormatting(selectedHeading: WKEditorHeaderSelectViewModel.Configuration, currentSelectionState: WKSourceEditorSelectionState, textView: UITextView) {
+    func toggleHeadingFormatting(selectedHeading: WKEditorInputView.HeadingButtonType, currentSelectionState: WKSourceEditorSelectionState, textView: UITextView) {
         
         var currentStateIsParagraph = false
         if currentSelectionState.isHeading {
diff --git a/Components/Sources/Components/Components/Editors/Source Editor/WKSourceEditorViewController.swift b/Components/Sources/Components/Components/Editors/Source Editor/WKSourceEditorViewController.swift
index bdea5b34ccd..bb43e24f3b9 100644
--- a/Components/Sources/Components/Components/Editors/Source Editor/WKSourceEditorViewController.swift	
+++ b/Components/Sources/Components/Components/Editors/Source Editor/WKSourceEditorViewController.swift	
@@ -291,8 +291,8 @@ extension WKSourceEditorViewController: WKEditorToolbarHighlightViewDelegate {
 // MARK: - WKEditorInputViewDelegate
 
 extension WKSourceEditorViewController: WKEditorInputViewDelegate {
-    func didTapHeading(selectedHeading: WKEditorHeaderSelectViewModel.Configuration) {
-        textFrameworkMediator.headingFormatter?.toggleHeadingFormatting(selectedHeading: selectedHeading, currentSelectionState: selectionState(), textView: textView)
+    func didTapHeading(type: WKEditorInputView.HeadingButtonType) {
+        textFrameworkMediator.headingFormatter?.toggleHeadingFormatting(selectedHeading: type, currentSelectionState: selectionState(), textView: textView)
     }
     
     func didTapBold(isSelected: Bool) {