OpenAppLock/Severed/Services/ShieldController.swift

//
//  ShieldController.swift
//  Severed
//

import FamilyControls
import Foundation
import ManagedSettings

/// Applies and clears app shields for rules. One implementation talks to
/// ManagedSettings; the mock records calls for tests.
protocol ShieldApplying: AnyObject {
    func applyShield(
        ruleID: UUID, selectionData: Data?, mode: SelectionMode, blockAdultContent: Bool
    )
    /// Clears every shield except those for the given rule IDs. Covers rules
    /// that were deleted or expired while the app was not running.
    func clearShields(except activeRuleIDs: Set<UUID>)
}

/// Real shield enforcement via per-rule `ManagedSettingsStore`s. Store names
/// are tracked in UserDefaults because ManagedSettings cannot enumerate them.
final class ManagedSettingsShieldController: ShieldApplying {
    private static let trackedIDsKey = "shieldedRuleIDs"
    private let defaults: UserDefaults

    init(defaults: UserDefaults = .standard) {
        self.defaults = defaults
    }

    func applyShield(
        ruleID: UUID, selectionData: Data?, mode: SelectionMode, blockAdultContent: Bool
    ) {
        let store = store(for: ruleID)
        let selection = AppSelectionCodec.decode(selectionData)
        switch mode {
        case .block:
            store.shield.applications =
                selection.applicationTokens.isEmpty ? nil : selection.applicationTokens
            store.shield.applicationCategories =
                selection.categoryTokens.isEmpty ? nil : .specific(selection.categoryTokens)
            store.shield.webDomains =
                selection.webDomainTokens.isEmpty ? nil : selection.webDomainTokens
        case .allowOnly:
            store.shield.applicationCategories = .all(except: selection.applicationTokens)
            store.shield.webDomainCategories = .all(except: selection.webDomainTokens)
        }
        // Screen Time's "Limit Adult Websites" filter for the rule's lifetime.
        store.webContent.blockedByFilter = blockAdultContent ? .auto() : nil
        track(ruleID: ruleID)
    }

    func clearShields(except activeRuleIDs: Set<UUID>) {
        for ruleID in trackedIDs.subtracting(activeRuleIDs) {
            store(for: ruleID).clearAllSettings()
            untrack(ruleID: ruleID)
        }
    }

    private func store(for ruleID: UUID) -> ManagedSettingsStore {
        ManagedSettingsStore(named: ManagedSettingsStore.Name("rule-\(ruleID.uuidString)"))
    }

    private var trackedIDs: Set<UUID> {
        Set((defaults.stringArray(forKey: Self.trackedIDsKey) ?? []).compactMap(UUID.init))
    }

    private func track(ruleID: UUID) {
        let ids = trackedIDs.union([ruleID])
        defaults.set(ids.map(\.uuidString).sorted(), forKey: Self.trackedIDsKey)
    }

    private func untrack(ruleID: UUID) {
        let ids = trackedIDs.subtracting([ruleID])
        defaults.set(ids.map(\.uuidString).sorted(), forKey: Self.trackedIDsKey)
    }
}

/// Records shield operations without touching the system. Used by tests and
/// UI-test launches.
final class MockShieldController: ShieldApplying {
    private(set) var shieldedRuleIDs: Set<UUID> = []
    private(set) var appliedModes: [UUID: SelectionMode] = [:]
    private(set) var appliedAdultContentFlags: [UUID: Bool] = [:]

    func applyShield(
        ruleID: UUID, selectionData: Data?, mode: SelectionMode, blockAdultContent: Bool
    ) {
        shieldedRuleIDs.insert(ruleID)
        appliedModes[ruleID] = mode
        appliedAdultContentFlags[ruleID] = blockAdultContent
    }

    func clearShields(except activeRuleIDs: Set<UUID>) {
        shieldedRuleIDs.formIntersection(activeRuleIDs)
        appliedModes = appliedModes.filter { activeRuleIDs.contains($0.key) }
        appliedAdultContentFlags = appliedAdultContentFlags.filter {
            activeRuleIDs.contains($0.key)
        }
    }
}

/// Encodes/decodes `FamilyActivitySelection` for persistence on the rule model.
enum AppSelectionCodec {
    static func encode(_ selection: FamilyActivitySelection) -> Data? {
        try? JSONEncoder().encode(selection)
    }

    static func decode(_ data: Data?) -> FamilyActivitySelection {
        guard let data,
              let selection = try? JSONDecoder().decode(FamilyActivitySelection.self, from: data)
        else { return FamilyActivitySelection() }
        return selection
    }

    static func count(of selection: FamilyActivitySelection) -> Int {
        selection.applicationTokens.count
            + selection.categoryTokens.count
            + selection.webDomainTokens.count
    }
}