@preconcurrency import ScreenTime
import WebKit

private nonisolated(unsafe) var screenTimeConfigurationObserverKVOContext: UInt8 = 0

/// Observer instance to mimic the WKScreenTimeConfigurationObserver in WebKit.
@MainActor
final class MyObserver: NSObject {

    private var observer: STScreenTimeConfigurationObserver?

    // The same queue that it utilized by WebKit
    private let queue = DispatchQueue.global(qos: .default)

    func startObserving() {

        guard observer == nil else {
            return
        }

        let observer = STScreenTimeConfigurationObserver(
            updateQueue: queue // DANGER: A concurrent queue is passed in instead of serial.
        )
        self.observer = observer

        // The class is STScreenTimeConfigurationObserver
        print("Observer BEFORE KVO: \(String(cString: object_getClassName(observer)))")

        observer.addObserver(
            self,
            forKeyPath: "configuration.enforcesChildRestrictions",
            options: [],
            context: &screenTimeConfigurationObserverKVOContext
        )

        // The class is NSKVONotifying_STScreenTimeConfigurationObserver
        print("Observer After KVO: \(String(cString: object_getClassName(observer)))")

        // That means there is an "sudo" override of setConfiguration.
        // The actual code in the stack trace is _NSSetObjectValueAndNotify.
        // This is logically the same.
        // - (void) setConfiguration: STScreenTimeConfiguration* {
        //    [self willChangeValueForKey: @"configuration"] // not thread safe
        //    [super setConfiguration: configuration] // thread safe
        //    [self didChangeValueForKey: @"configuration"] // not thread safe
        // }


        // If the configuration is set to a non nil property before hand, crash doesn't occurs.
        // call me before observer.startObserving() is called.
        // This works b/c STScreenTimeConfigurationObserver does not create a dynamic KVO subclass.
        // It's properties are "safe" from crash but there are still bad protections.

        // UNCOMMENT ME TO FIX:
        observer.set(configuration: .create(enforcesChildRestrictions: false))

        observer.startObserving() // call not need to reproduce.

        simulateMultipleXPCConnectionUpdates()
    }

    // Simulate multiple calls from the underlying XPC Service at the same time b/c a the concurrent queue
    private func simulateMultipleXPCConnectionUpdates() {
        guard let observer else { return }
        for _ in 0..<1_000 {
            queue.asyncAfter(deadline: .now()) {
                let enforcesChildRestrictions = arc4random() % 2 == 0
                observer.update(
                    configuration: .create(
                        enforcesChildRestrictions: enforcesChildRestrictions
                    )
                )
                // The method observer.update(configuration:) doesn't have the crash
            }
        }
    }

    func stopObserving() {
        observer?.stopObserving()
        observer?.removeObserver(
            self,
            forKeyPath: "configuration.enforcesChildRestrictions",
            context: &screenTimeConfigurationObserverKVOContext
        )
        observer = nil
    }

    nonisolated override func observeValue(
        forKeyPath keyPath: String?,
        of object: Any?,
        change: [NSKeyValueChangeKey : Any]?,
        context: UnsafeMutableRawPointer?
    ) {
        if context == &screenTimeConfigurationObserverKVOContext {
            let enforcesChildRestrictions = (object as? STScreenTimeConfigurationObserver)?.configuration?.enforcesChildRestrictions
            let valueString = enforcesChildRestrictions?.description ?? "nil"
            print("observeValue(forKeyPath: \(keyPath ?? "nil"), value: \(valueString))")
        } else {
            super.observeValue(forKeyPath: keyPath, of: object, change: change, context: context)
        }
    }
}

extension STScreenTimeConfigurationObserver {

    /// Calls private `[STScreenTimeConfigurationObserver: _updateWithConfiguration]` method
    func update(configuration: STScreenTimeConfiguration) {
        let selector = NSSelectorFromString("_updateWithConfiguration:")
        guard self.responds(to: selector) else {
            fatalError("instances doesn't responds to selector: \(selector)")
        }
        self.perform(selector, with: configuration)
    }

    /// Calls private `[STScreenTimeConfigurationObserver: setConfiguration:]` method
    func set(configuration: STScreenTimeConfiguration) {
        let selector = NSSelectorFromString("setConfiguration:")
        guard self.responds(to: selector) else {
            fatalError("instances doesn't responds to selector: \(selector)")
        }
        self.perform(selector, with: configuration)
    }
}

extension STScreenTimeConfiguration {
    /// Creates an STScreenTimeConfiguration instance using the private initializer
    /// - Parameter enforcesChildRestrictions: Whether to enforce child restrictions
    /// - Returns: A configured STScreenTimeConfiguration instance
    static func create(enforcesChildRestrictions: Bool) -> STScreenTimeConfiguration {
        let allocSelector = NSSelectorFromString("alloc")

        // Allocate instance using perform
        guard let uninitializedConfig = STScreenTimeConfiguration.perform(allocSelector)?.takeUnretainedValue() else {
            fatalError("Failed to allocate STScreenTimeConfiguration instance")
        }

        let selector = NSSelectorFromString("initWithEnforcesChildRestrictions:")

        // Check if selector exists
        guard uninitializedConfig.responds(to: selector) else {
            fatalError("Private initializer 'initWithEnforcesChildRestrictions:' not found on STScreenTimeConfiguration")
        }

        // Get the method implementation
        guard let method = class_getInstanceMethod(object_getClass(uninitializedConfig), selector) else {
            fatalError("Failed to get method implementation for 'initWithEnforcesChildRestrictions:'")
        }

        let implementation = method_getImplementation(method)

        // Cast to the appropriate function signature
        // Signature: (id self, SEL _cmd, BOOL enforcesChildRestrictions) -> id
        typealias InitFunction = @convention(c) (AnyObject, Selector, Bool) -> AnyObject

        let initFunction = unsafeBitCast(implementation, to: InitFunction.self)

        // Call the function
        let initialized = initFunction(uninitializedConfig, selector, enforcesChildRestrictions)

        guard let config = initialized as? STScreenTimeConfiguration else {
            fatalError("Initialization returned unexpected type")
        }

        return config
    }

    // Calls the private `[STScreenTimeConfiguration setEnforcesChildRestrictions]`
    func set(enforcesChildRestrictions: Bool) {
        let selector = NSSelectorFromString("setEnforcesChildRestrictions:")
        guard self.responds(to: selector) else {
            fatalError("instances doesn't responds to selector: \(selector)")
        }

        self.perform(selector, with: enforcesChildRestrictions)
    }
}setEnforcesChildRestrictions