class Button: UIButton {
    
    private class Subview: UIView {
        // Never allow isUserInteractionEnabled in any button subview
        override var isUserInteractionEnabled: Bool {
            get { false }
            set { super.isUserInteractionEnabled = false }
        }
    }
    
    // Creating custom classes for each subview for inspection purposes
    private class BackgroundView: Subview { }
    private class ShadowView: Subview { }
    
    var animatesStateChanges: Bool = false
    
    private(set) lazy var backgroundView: UIView = BackgroundView(frame: bounds)
    private(set) lazy var shadowView: UIView = ShadowView(frame: bounds)
    
    private var originalTintColor: UIColor!
    
    private var stateValueContainers: [PartialKeyPath<Button> : AnyStateValueContainer] = [:]
    
    override init(frame: CGRect) {
        super.init(frame: frame)
        originalTintColor = tintColor
        insertSubview(backgroundView, at: 0)
        insertSubview(shadowView, belowSubview: backgroundView)
    }
    
    required init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
}

// MARK: - Updating apperance
extension Button {
    
    override var isHighlighted: Bool { didSet {
        updateAppearance(animated: animatesStateChanges)
    }}
    
    override var isEnabled: Bool { didSet {
        updateAppearance()
    }}
    
    override var isSelected: Bool { didSet {
        updateAppearance(animated: animatesStateChanges)
    }}
    
    func updateAppearance(animated: Bool = false) {
        
        let updates = {
            // All self-animating value containers or all updates when non animated
            self.stateValueContainers
                .filter { !$1.animatesUpdate || !animated }
                .forEach { keyPath, container in
                    container.update(object: self, for: self.state, using: keyPath)
                }
        }
        if animated {
            UIView.animate(
                withDuration: 0.25, delay: 0,
                options: [.beginFromCurrentState, .allowUserInteraction],
                timingFunction: .quintOut,
                animations: updates)
            
            // Let all self-animating value containers animate
            self.stateValueContainers
                .filter { $1.animatesUpdate }
                .forEach { keyPath, container in
                    container.update(object: self, for: self.state, using: keyPath)
                }
        } else {
            UIView.performWithoutAnimation {
                // Disable all self-animating value container animations
                updates()
            }
        }
        
    }
}

// MARK: - View hierarchy consistency
extension Button {
    
    override func layoutSubviews() {
        super.layoutSubviews()
        
        // Make sure shadowView and backgroundView are always in the back
        [shadowView, backgroundView]
            .enumerated()
            .forEach { index, view in
                if view.superview?.subviews.firstIndex(of: view) != index {
                    view.removeFromSuperview()
                    view.frame = bounds
                    insertSubview(view, at: index)
                } else {
                    view.frame = bounds
                }
            }
    }
}

// MARK: - StateValueContainer getting and setting
extension Button {
    
    func stateValueContainer<Value>(for keyPath: KeyPath<Button, Value>) -> StateValueContainer<Value> {
        if let key = stateValueContainers.keys.first(where: { $0 == keyPath }),
           let container = stateValueContainers[key]?.base as? StateValueContainer<Value> {
            return container
        }
        let container = StateValueContainer<Value>()
        stateValueContainers[keyPath as PartialKeyPath] = AnyStateValueContainer(container)
        return container
    }

    @discardableResult
    func setValue<Value: Equatable>(_ value: Optional<Value>, for keyPath: KeyPath<Button, Value>, state: State) -> StateValueContainer<Value> {
        let container = stateValueContainer(for: keyPath)
        container.set(value, for: state)
        container.update(object: self, for: self.state, using: keyPath)
        return container
    }
    
    func value<Value: Equatable>(for keyPath: KeyPath<Button, Value>, state: State) -> Value? {
        stateValueContainer(for: keyPath).value(for: state)
    }
    
    func value<Value: Equatable>(for keyPath: KeyPath<Button, Optional<Value>>, state: UIButton.State) -> Value? {
        stateValueContainer(for: keyPath).value(for: state) as? Value
    }
    
    func processedValue<Value: Equatable>(for keyPath: KeyPath<Button, Value>, state: State, defaultValue: Value) -> (usedState: State, value: Value)? {
        stateValueContainer(for: keyPath).processedValue(for: state, defaultValue: defaultValue)
    }
    
    func processedValue<Value: Equatable>(for keyPath: KeyPath<Button, Optional<Value>>, state: State, defaultValue: Value? = nil) -> (usedState: State, value: Value?)? {
        stateValueContainer(for: keyPath).processedValue(for: state, defaultValue: defaultValue)
    }
}

// MARK: - Background color
extension Button {
    
    func setBackgroundColor(_ color: UIColor?, for state: State) {
        setValue(color, for: \.backgroundView.backgroundColor, state: state)
    }
    
    func backgroundColor(for state: State) -> UIColor? {
        value(for: \.backgroundView.backgroundColor, state: state)
    }
}

// MARK: - Image color
extension Button {
    
    override func setImage(_ image: UIImage?, for state: UIControl.State) {
        if image?.renderingMode != .alwaysTemplate, imageColor(for: state) != nil {
            super.setImage(image?.withRenderingMode(.alwaysTemplate), for: state)
            return
        }
        super.setImage(image, for: state)
    }
    
    func setImageColor(_ color: UIColor?, for state: State) {
        if let image = image(for: state) {
            setImage(image.withRenderingMode(.alwaysTemplate), for: state)
        }
        if color != nil {
            adjustsImageWhenHighlighted = false
            adjustsImageWhenDisabled = false
        }
        setValue(color, for: \.imageView?.tintColor, state: state)
            .onUpdate { [unowned self] state, container in
                self.adjustsImageWhenHighlighted = !container.isEmpty
                self.adjustsImageWhenDisabled = !container.isEmpty
                
                if let result = container.processedValue(for: state, defaultValue: self.originalTintColor) {
                    imageView?.tintColor = result.value
                }
            }
    }
    
    func imageColor(for state: State) -> UIColor? {
        value(for: \.imageView?.tintColor, state: state)
    }
}

// MARK: - Border color
extension Button {
    
    func setBorderColor(_ color: UIColor?, for state: State) {
        setValue(color?.cgColor, for: \.backgroundView.layer.borderColor, state: state)
    }
    
    func borderColor(for state: State) -> UIColor? {
        value(for: \.backgroundView.layer.borderColor, state: state).map { UIColor(cgColor: $0) }
    }
}