session-ios/SessionUIKit/Style Guide/ThemeManager.swift

// Copyright © 2022 Rangeproof Pty Ltd. All rights reserved.

import UIKit
import GRDB
import SessionUtilitiesKit

// MARK: - Preferences

public extension Setting.EnumKey {
    /// Controls what theme should be used
    static let theme: Setting.EnumKey = "selectedTheme"
    
    /// Controls what primary color should be used for the theme
    static let themePrimaryColor: Setting.EnumKey = "selectedThemePrimaryColor"
}

public extension Setting.BoolKey {
    /// A flag indicating whether the app should match system day/night settings
    static let themeMatchSystemDayNightCycle: Setting.BoolKey = "themeMatchSystemDayNightCycle"
}

// MARK: - ThemeManager

public enum ThemeManager {
    private static var hasSetInitialSystemTrait: Bool = false
    
    /// **Note:** Using `weakToStrongObjects` means that the value types will continue to be maintained until the map table resizes
    /// itself (ie. until a new UI element is registered to the table)
    ///
    /// Unfortunately if we don't do this the `ThemeApplier` is immediately deallocated and we can't use it to update the theme
    private static var uiRegistry: NSMapTable<AnyObject, ThemeApplier> = NSMapTable.weakToStrongObjects()
    
    public static var currentTheme: Theme = {
        Storage.shared[.theme].defaulting(to: Theme.classicDark)
    }() {
        didSet {
            // Only update if it was changed
            guard oldValue != currentTheme else { return }
            
            Storage.shared.writeAsync { db in
                db[.theme] = currentTheme
            }
            
            // Only trigger the UI update if the primary colour wasn't changed (otherwise we'd be doing
            // an extra UI update
            if let defaultPrimaryColor: Theme.PrimaryColor = Theme.PrimaryColor(color: currentTheme.colors[.defaultPrimary]) {
                guard primaryColor == defaultPrimaryColor else {
                    ThemeManager.primaryColor = defaultPrimaryColor
                    return
                }
            }
            
            updateAllUI()
        }
    }
    
    public static var primaryColor: Theme.PrimaryColor = {
        Storage.shared[.themePrimaryColor].defaulting(to: Theme.PrimaryColor.green)
    }() {
        didSet {
            // Only update if it was changed
            guard oldValue != primaryColor else { return }
            
            Storage.shared.writeAsync { db in
                db[.themePrimaryColor] = primaryColor
            }
            
            updateAllUI()
        }
    }
    
    public static var matchSystemNightModeSetting: Bool = {
        Storage.shared[.themeMatchSystemDayNightCycle]
    }() {
        didSet {
            // Only update if it was changed
            guard oldValue != matchSystemNightModeSetting else { return }
            
            Storage.shared.writeAsync { db in
                db[.themeMatchSystemDayNightCycle] = matchSystemNightModeSetting
            }
            
            // Note: We have to trigger this directly or the 'TraitObservingWindow' won't actually
            // trigger the trait change if the app launched with this setting switched off
            applyWindowStyling()
        }
    }
    
    // When this gets set we need to update the UI to ensure the global appearance stuff is set
    // correctly on launch
    public static weak var mainWindow: UIWindow? {
        didSet { updateAllUI() }
    }
    
    // MARK: - Functions
    
    public static func traitCollectionDidChange(_ previousTraitCollection: UITraitCollection?) {
        let currentUserInterfaceStyle: UIUserInterfaceStyle = UITraitCollection.current.userInterfaceStyle
        
        // Only trigger updates if the style changed and the device is set to match the system style
        guard
            currentUserInterfaceStyle != ThemeManager.currentTheme.interfaceStyle,
            ThemeManager.matchSystemNightModeSetting
        else { return }
        
        // Swap to the appropriate light/dark mode
        switch (currentUserInterfaceStyle, ThemeManager.currentTheme) {
            case (.light, .classicDark): ThemeManager.currentTheme = .classicLight
            case (.light, .oceanDark): ThemeManager.currentTheme = .oceanLight
            case (.dark, .classicLight): ThemeManager.currentTheme = .classicDark
            case (.dark, .oceanLight): ThemeManager.currentTheme = .oceanDark
            default: break
        }
    }
    
    public static func applySavedTheme() {
        ThemeManager.primaryColor = Storage.shared[.themePrimaryColor].defaulting(to: Theme.PrimaryColor.green)
        ThemeManager.currentTheme = Storage.shared[.theme].defaulting(to: Theme.classicDark)
    }
    
    public static func applyNavigationStyling() {
        let textPrimary: UIColor = (ThemeManager.currentTheme.colors[.textPrimary] ?? .white)
        
        // Set the `mainWindow.tintColor` for system screens to use the right colour for text
        ThemeManager.mainWindow?.tintColor = textPrimary
        ThemeManager.mainWindow?.rootViewController?.setNeedsStatusBarAppearanceUpdate()
        
        // Update the nav bars to use the right colours
        UINavigationBar.appearance().barTintColor = ThemeManager.currentTheme.colors[.backgroundPrimary]
        UINavigationBar.appearance().isTranslucent = false
        UINavigationBar.appearance().tintColor = textPrimary
        UINavigationBar.appearance().shadowImage = ThemeManager.currentTheme.colors[.backgroundPrimary]?.toImage()
        UINavigationBar.appearance().titleTextAttributes = [
            NSAttributedString.Key.foregroundColor: textPrimary
        ]
        UINavigationBar.appearance().largeTitleTextAttributes = [
            NSAttributedString.Key.foregroundColor: textPrimary
        ]
        
        // Update the bar button item appearance
        UIBarButtonItem.appearance().tintColor = textPrimary

        // Update toolbars to use the right colours
        UIToolbar.appearance().barTintColor = ThemeManager.currentTheme.colors[.backgroundPrimary]
        UIToolbar.appearance().isTranslucent = false
        UIToolbar.appearance().tintColor = textPrimary
        
        // Note: Looks like there were changes to the appearance behaviour in iOS 15, unfortunately
        // this breaks parts of the old 'UINavigationBar.appearance()' logic so we need to do everything
        // again using the new API...
        if #available(iOS 15.0, *) {
            let appearance = UINavigationBarAppearance()
            appearance.configureWithOpaqueBackground()
            appearance.backgroundColor = ThemeManager.currentTheme.colors[.backgroundPrimary]
            appearance.shadowImage = ThemeManager.currentTheme.colors[.backgroundPrimary]?.toImage()
            appearance.titleTextAttributes = [
                NSAttributedString.Key.foregroundColor: textPrimary
            ]
            appearance.largeTitleTextAttributes = [
                NSAttributedString.Key.foregroundColor: textPrimary
            ]
            
            // Apply the button item appearance as well
            let barButtonItemAppearance = UIBarButtonItemAppearance(style: .plain)
            barButtonItemAppearance.normal.titleTextAttributes = [ .foregroundColor: textPrimary ]
            barButtonItemAppearance.disabled.titleTextAttributes = [ .foregroundColor: textPrimary ]
            barButtonItemAppearance.highlighted.titleTextAttributes = [ .foregroundColor: textPrimary ]
            barButtonItemAppearance.focused.titleTextAttributes = [ .foregroundColor: textPrimary ]
            appearance.buttonAppearance = barButtonItemAppearance
            appearance.backButtonAppearance = barButtonItemAppearance
            appearance.doneButtonAppearance = barButtonItemAppearance
            
            UINavigationBar.appearance().standardAppearance = appearance
            UINavigationBar.appearance().scrollEdgeAppearance = appearance
        }
        
        // Note: 'UINavigationBar.appearance' only affects newly created nav bars so we need
        // to force-update the current navigation bar (unfortunately the only way to do this
        // is to remove the nav controller from the view hierarchy and then re-add it)
        let currentNavController: UINavigationController? = {
            var targetViewController: UIViewController? = ThemeManager.mainWindow?.rootViewController
            
            while targetViewController?.presentedViewController != nil {
                targetViewController = targetViewController?.presentedViewController
            }
            
            return (
                (targetViewController as? UINavigationController) ??
                targetViewController?.navigationController
            )
        }()
        
        if
            let navController: UINavigationController = currentNavController,
            let superview: UIView = navController.view.superview,
            !navController.isNavigationBarHidden
        {
            navController.view.removeFromSuperview()
            superview.addSubview(navController.view)
            navController.topViewController?.setNeedsStatusBarAppearanceUpdate()
        }
    }
    
    public static func applyWindowStyling() {
        mainWindow?.overrideUserInterfaceStyle = {
            guard !ThemeManager.matchSystemNightModeSetting else { return .unspecified }
            
            switch ThemeManager.currentTheme.interfaceStyle {
                case .light: return .light
                case .dark, .unspecified: return .dark
                @unknown default: return .dark
            }
        }()
        mainWindow?.backgroundColor = ThemeManager.currentTheme.colors[.backgroundPrimary]
    }
    
    public static func onThemeChange(observer: AnyObject, callback: @escaping (Theme, Theme.PrimaryColor) -> ()) {
        ThemeManager.uiRegistry.setObject(
            ThemeApplier(
                existingApplier: nil,
                info: []
            ) { theme in callback(theme, ThemeManager.primaryColor) },
            forKey: observer
        )
    }
    
    private static func updateAllUI() {
        guard Thread.isMainThread else {
            DispatchQueue.main.async {
                updateAllUI()
            }
            return
        }
        
        ThemeManager.uiRegistry.objectEnumerator()?.forEach { applier in
            (applier as? ThemeApplier)?.apply(theme: currentTheme)
        }
        
        applyNavigationStyling()
        applyWindowStyling()
        
        if !hasSetInitialSystemTrait {
            traitCollectionDidChange(nil)
            hasSetInitialSystemTrait = true
        }
    }
    
    internal static func set<T: AnyObject>(
        _ view: T,
        keyPath: ReferenceWritableKeyPath<T, UIColor?>,
        to value: ThemeValue?,
        for state: UIControl.State = .normal
    ) {
        ThemeManager.uiRegistry.setObject(
            ThemeApplier(
                existingApplier: ThemeManager.get(for: view),
                info: [ keyPath ]
            ) { [weak view] theme in
                guard let value: ThemeValue = value else {
                    view?[keyPath: keyPath] = nil
                    return
                }

                view?[keyPath: keyPath] = ThemeManager.resolvedColor(theme.colors[value])
            },
            forKey: view
        )
    }
    
    internal static func set<T: AnyObject>(
        _ view: T,
        keyPath: ReferenceWritableKeyPath<T, CGColor?>,
        to value: ThemeValue?,
        for state: UIControl.State = .normal
    ) {
        ThemeManager.uiRegistry.setObject(
            ThemeApplier(
                existingApplier: ThemeManager.get(for: view),
                info: [ keyPath ]
            ) { [weak view] theme in
                guard let value: ThemeValue = value else {
                    view?[keyPath: keyPath] = nil
                    return
                }
                
                view?[keyPath: keyPath] = ThemeManager.resolvedColor(theme.colors[value])?.cgColor
            },
            forKey: view
        )
    }
    
    internal static func set<T: AnyObject>(
        _ view: T,
        to applier: ThemeApplier,
        for state: UIControl.State = .normal
    ) {
        ThemeManager.uiRegistry.setObject(applier, forKey: view)
    }
    
    /// Using a `UIColor(dynamicProvider:)` unfortunately doesn't seem to work properly for some controls (eg. UISwitch) so
    /// since we are already explicitly updating all UI when changing colours & states we just force-resolve the primary colour to avoid
    /// running into these glitches
    internal static func resolvedColor(_ color: UIColor?) -> UIColor? {
        return color?.resolvedColor(with: UITraitCollection())
    }
    
    internal static func get(for view: AnyObject) -> ThemeApplier? {
        return ThemeManager.uiRegistry.object(forKey: view)
    }
}

// MARK: - ThemeApplier

internal class ThemeApplier {
    enum InfoKey: String {
        case keyPath
        case controlState
    }
    
    private let applyTheme: (Theme) -> ()
    private let info: [AnyHashable]
    private var otherAppliers: [ThemeApplier]?
    
    init(
        existingApplier: ThemeApplier?,
        info: [AnyHashable],
        applyTheme: @escaping (Theme) -> ()
    ) {
        self.applyTheme = applyTheme
        self.info = info
        
        // Store any existing "appliers" (removing their 'otherApplier' references to prevent
        // loops and excluding any which match the current "info" as they should be replaced
        // by this applier)
        self.otherAppliers = [existingApplier]
            .appending(contentsOf: existingApplier?.otherAppliers)
            .compactMap { $0?.clearingOtherAppliers() }
            .filter { $0.info != info }
        
        // Automatically apply the theme immediately
        self.apply(theme: ThemeManager.currentTheme)
    }
    
    // MARK: - Functions
    
    private func clearingOtherAppliers() -> ThemeApplier {
        self.otherAppliers = nil
        
        return self
    }
    
    fileprivate func apply(theme: Theme) {
        self.applyTheme(theme)
        
        // If there are otherAppliers stored against this one then trigger those as well
        self.otherAppliers?.forEach { applier in
            applier.applyTheme(theme)
        }
    }
}

// MARK: - Convenience Extensions

extension Array {
    fileprivate func appending(contentsOf other: [Element]?) -> [Element] {
        guard let other: [Element] = other else { return self }
        
        var updatedArray: [Element] = self
        updatedArray.append(contentsOf: other)
        return updatedArray
    }
}