import UIKit

extension UIControl.State: Hashable { }

protocol StateValueContainable {
    var animatesUpdate: Bool { get }
    func update<Object: AnyObject>(object: Object, for state: UIButton.State, using keyPath: PartialKeyPath<Object>)
}

struct AnyStateValueContainer: StateValueContainable {
    
    let base: Any
    
    init<Container>(_ wrapped: Container) where Container: StateValueContainable {
        base = wrapped
    }
    
    var animatesUpdate: Bool {
        (base as? StateValueContainable)?.animatesUpdate == true
    }
    
    func update<Object: AnyObject>(object: Object, for state: UIButton.State, using keyPath: PartialKeyPath<Object>) {
        guard let container = base as? StateValueContainable else {
            return
        }
        container.update(object: object, for: state, using: keyPath)
    }
}

/// Handles storing and retrieving values linked to button states
/// This allows for adding specific values to a UIButton linked
/// to any (combo of) button state, like fonts and colors
final class StateValueContainer<Value: Equatable>: StateValueContainable {
    
    private(set) var animatesUpdate: Bool = false
    
    private var storage = [UIButton.State: Value]()
    private var updater: ((UIButton.State, StateValueContainer<Value>) -> Void)?
    
    subscript(state: UIButton.State) -> Value? {
        get {
            return value(for: state)
        }
        set {
            set(newValue, for: state)
        }
    }
    
    /// Sets the value for the given state, removes it if `nil`
    func set(_ value: Value?, for state: UIButton.State) {
        guard let value = value else {
            storage.removeValue(forKey: state)
            return
        }
        storage[state] = value
    }
    
    /// Returns the value if available
    func value(for state: UIButton.State) -> Value? {
        storage[state]
    }
    
    /// Whether any values have been stored
    var isEmpty: Bool {
        return storage.isEmpty
    }
    
    /// To closely approximate `UIButton`'s default behavior of using the normal state as a fallback state
    /// and only use a default (or nil) value when there is no value set for the normal state, use this method
    /// when you need to update a specific value in your view
    ///
    /// - This method will return nil when there's no values set, as to not override any view properties
    /// set outside of the set/get methods.
    /// - Will either use the value for the state given or the value set for the `.normal` state
    /// - Only returns `defaultValue` when the requested state is `.normal`, otherwise the view's
    /// properties should not be changed
    ///
    /// - Parameters:
    ///   - state: Control state to request the value for
    ///   - defaultValue: Value to use when no value available for the given state or `.normal` state
    /// - Returns: Tuple with the used state and the value to set for the value
    func processedValue(for state: UIButton.State, defaultValue: Value? = nil) -> (usedState: UIButton.State, value: Value?)? {
        guard !isEmpty else {
            return nil
        }
        
        var result = (usedState: state, value: nil as Value?)
        
        if let value = value(for: state) {
            result.value = value
        } else if state == .normal || value(for: .normal) != nil {
            result.usedState = .normal
            result.value = value(for: .normal) ?? defaultValue
        }
        return result
    }
    
    /// To closely approximate `UIButton`'s default behavior of using the normal state as a fallback state
    /// and only use a default (or nil) value when there is no value set for the normal state, use this method
    /// when you need to update a specific value in your view
    ///
    /// - This method will return nil when there's no values set, as to not override any view properties
    /// set outside of the set/get methods.
    /// - Will either use the value for the state given or the value set for the `.normal` state
    /// - Only returns `defaultValue` when the requested state is `.normal`, otherwise the view's
    /// properties should not be changed
    ///
    /// - Note: `defaultValue` is not optional in this method to make sure to always get a value in the
    /// `.normal` state
    ///
    /// - Parameters:
    ///   - state: Control state to request the value for
    ///   - defaultValue: Value to use when no value available for the given state or `.normal` state
    /// - Returns: Tuple with the used state and the value to set for the value
    func processedValue(for state: UIButton.State, defaultValue: Value) -> (usedState: UIButton.State, value: Value)? {
        guard let result = processedValue(for: state, defaultValue: Optional(defaultValue)) as? (UIButton.State, Value) else {
            return nil
        }
        return result
    }
    
    func update<Object: AnyObject>(object: Object, for state: UIButton.State, using keyPath: PartialKeyPath<Object>) {
        if let updater = updater {
            updater(state, self)
            return
        }
        
        if let keyPath = keyPath as? ReferenceWritableKeyPath<Object, Value> {
            if let value = processedValue(for: state)?.value {
                object[keyPath: keyPath] = value
            }
        } else if let keyPath = keyPath as? ReferenceWritableKeyPath<Object, Optional<Value>> {
            if let result = processedValue(for: state) {
                object[keyPath: keyPath] = result.value
            }
        }
        
    }
}

extension StateValueContainer {
    @discardableResult
    func onUpdate(_ handler: @escaping (UIButton.State, StateValueContainer<Value>) -> Void) -> Self {
        updater = handler
        return self
    }
    
    @discardableResult
    func animatesUpdate(_ animates: Bool = true) -> Self {
        animatesUpdate = true
        return self
    }
}