LocalData/Sources/LocalData/Services/StorageRouter.swift

import Foundation
#if os(iOS) || os(watchOS)
import WatchConnectivity
#endif

/// The main storage router that coordinates all storage operations.
/// Uses specialized helper actors for each storage domain.
public actor StorageRouter: StorageProviding {
    
    public static let shared = StorageRouter()
    
    private var registeredKeyNames: Set<String> = []
    
    private init() {}

    // MARK: - Configuration

    /// Updates the encryption configuration.
    /// > [!WARNING]
    /// > Changing these constants in an existing app will cause the app to look for the master key 
    /// > under a new name. Previously encrypted data will be lost.
    public func updateEncryptionConfiguration(_ configuration: EncryptionConfiguration) async {
        await EncryptionHelper.shared.updateConfiguration(configuration)
    }

    /// Updates the sync configuration.
    public func updateSyncConfiguration(_ configuration: SyncConfiguration) async {
        await SyncHelper.shared.updateConfiguration(configuration)
    }

    // MARK: - Key Material Providers

    /// Registers a key material provider for external encryption policies.
    public func registerKeyMaterialProvider(
        _ provider: any KeyMaterialProviding,
        for source: KeyMaterialSource
    ) async {
        await EncryptionHelper.shared.registerKeyMaterialProvider(provider, for: source)
    }

    /// Registers a catalog of known storage keys for audit and validation.
    /// When registered, all storage operations will verify keys are listed.
    public func registerCatalog<C: StorageKeyCatalog>(_ catalog: C.Type) throws {
        let entries = catalog.allKeys
        try validateDescription(entries)
        try validateUniqueKeys(entries)
        registeredKeyNames = Set(entries.map { $0.descriptor.name })
    }
    
    // MARK: - StorageProviding Implementation
    
    /// Stores a value for the given key.
    /// - Parameters:
    ///   - value: The value to store.
    ///   - key: The storage key defining where and how to store.
    /// - Throws: Various errors depending on the storage domain and security policy.
    public func set<Key: StorageKey>(_ value: Key.Value, for key: Key) async throws {
        try validateCatalogRegistration(for: key)
        try validatePlatformAvailability(for: key)
        
        let data = try serialize(value, with: key.serializer)
        let securedData = try await applySecurity(data, for: key, isEncrypt: true)
        
        try await store(securedData, for: key)
        try await handleSync(key, data: securedData)
    }
    
    /// Retrieves a value for the given key.
    /// - Parameter key: The storage key to retrieve.
    /// - Returns: The stored value.
    /// - Throws: `StorageError.notFound` if no value exists, plus domain-specific errors.
    public func get<Key: StorageKey>(_ key: Key) async throws -> Key.Value {
        try validateCatalogRegistration(for: key)
        try validatePlatformAvailability(for: key)
        
        guard let securedData = try await retrieve(for: key) else {
            throw StorageError.notFound
        }
        
        let data = try await applySecurity(securedData, for: key, isEncrypt: false)
        return try deserialize(data, with: key.serializer)
    }
    
    /// Removes the value for the given key.
    /// - Parameter key: The storage key to remove.
    /// - Throws: Domain-specific errors if removal fails.
    public func remove<Key: StorageKey>(_ key: Key) async throws {
        try validateCatalogRegistration(for: key)
        try validatePlatformAvailability(for: key)
        try await delete(for: key)
    }
    
    /// Checks if a value exists for the given key.
    /// - Parameter key: The storage key to check.
    /// - Returns: True if a value exists.
    public func exists<Key: StorageKey>(_ key: Key) async throws -> Bool {
        try validateCatalogRegistration(for: key)
        try validatePlatformAvailability(for: key)
        
        switch key.domain {
        case .userDefaults(let suite):
            return try await UserDefaultsHelper.shared.exists(forKey: key.name, suite: suite)
        case .appGroupUserDefaults(let identifier):
            return try await UserDefaultsHelper.shared.exists(forKey: key.name, appGroupIdentifier: identifier)
        case .keychain(let service):
            return try await KeychainHelper.shared.exists(service: service, key: key.name)
        case .fileSystem(let directory), .encryptedFileSystem(let directory):
            return await FileStorageHelper.shared.exists(in: directory, fileName: key.name)
        case .appGroupFileSystem(let identifier, let directory):
            return await FileStorageHelper.shared.exists(
                in: directory,
                fileName: key.name,
                appGroupIdentifier: identifier
            )
        }
    }
    
    // MARK: - Platform Validation
    
    private func validatePlatformAvailability<Key: StorageKey>(for key: Key) throws {
        #if os(watchOS)
        if key.availability == .phoneOnly {
            throw StorageError.phoneOnlyKeyAccessedOnWatch(key.name)
        }
        #elseif os(iOS)
        if key.availability == .watchOnly {
            throw StorageError.watchOnlyKeyAccessedOnPhone(key.name)
        }
        #endif
    }

    private func validateCatalogRegistration<Key: StorageKey>(for key: Key) throws {
        guard !registeredKeyNames.isEmpty else { return }
        guard registeredKeyNames.contains(key.name) else {
#if DEBUG
            assertionFailure("StorageKey not registered in catalog: \(key.name)")
#endif
            throw StorageError.unregisteredKey(key.name)
        }
    }

    private func validateUniqueKeys(_ entries: [AnyStorageKey]) throws {
        var exactNames: [String: Int] = [:]
        var duplicates: [String] = []

        for entry in entries {
            exactNames[entry.descriptor.name, default: 0] += 1
        }

        for (name, count) in exactNames where count > 1 {
            duplicates.append(name)
        }

        guard duplicates.isEmpty else {
            throw StorageError.duplicateRegisteredKeys(duplicates.sorted())
        }
    }

    private func validateDescription(_ entries: [AnyStorageKey]) throws {
        let missing = entries
            .map(\.descriptor)
            .filter { $0.description.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty }
            .map(\.name)

        guard missing.isEmpty else {
            throw StorageError.missingDescription(missing.sorted().joined(separator: ", "))
        }
    }
    
    // MARK: - Serialization
    
    private func serialize<Value: Codable & Sendable>(
        _ value: Value,
        with serializer: Serializer<Value>
    ) throws -> Data {
        do {
            return try serializer.encode(value)
        } catch {
            throw StorageError.serializationFailed
        }
    }
    
    private func deserialize<Value: Codable & Sendable>(
        _ data: Data,
        with serializer: Serializer<Value>
    ) throws -> Value {
        do {
            return try serializer.decode(data)
        } catch {
            throw StorageError.deserializationFailed
        }
    }
    
    // MARK: - Security
    
    private func applySecurity(
        _ data: Data,
        for key: any StorageKey,
        isEncrypt: Bool
    ) async throws -> Data {
        switch key.security {
        case .none:
            return data
            
        case .encrypted(let encryptionPolicy):
            if isEncrypt {
                return try await EncryptionHelper.shared.encrypt(
                    data,
                    keyName: key.name,
                    policy: encryptionPolicy
                )
            } else {
                return try await EncryptionHelper.shared.decrypt(
                    data,
                    keyName: key.name,
                    policy: encryptionPolicy
                )
            }
            
        case .keychain:
            // Keychain security is handled in store/retrieve
            return data
        }
    }
    
    // MARK: - Storage Operations
    
    private func store(_ data: Data, for key: any StorageKey) async throws {
        switch key.domain {
        case .userDefaults(let suite):
            try await UserDefaultsHelper.shared.set(data, forKey: key.name, suite: suite)
            
        case .appGroupUserDefaults(let identifier):
            try await UserDefaultsHelper.shared.set(data, forKey: key.name, appGroupIdentifier: identifier)
            
        case .keychain(let service):
            guard case let .keychain(accessibility, accessControl) = key.security else {
                throw StorageError.securityApplicationFailed
            }
            try await KeychainHelper.shared.set(
                data,
                service: service,
                key: key.name,
                accessibility: accessibility,
                accessControl: accessControl
            )
            
        case .fileSystem(let directory):
            try await FileStorageHelper.shared.write(
                data,
                to: directory,
                fileName: key.name,
                useCompleteFileProtection: false
            )
            
        case .encryptedFileSystem(let directory):
            try await FileStorageHelper.shared.write(
                data,
                to: directory,
                fileName: key.name,
                useCompleteFileProtection: true
            )
            
        case .appGroupFileSystem(let identifier, let directory):
            try await FileStorageHelper.shared.write(
                data,
                to: directory,
                fileName: key.name,
                appGroupIdentifier: identifier,
                useCompleteFileProtection: false
            )
        }
    }
    
    private func retrieve(for key: any StorageKey) async throws -> Data? {
        switch key.domain {
        case .userDefaults(let suite):
            return try await UserDefaultsHelper.shared.get(forKey: key.name, suite: suite)
        case .appGroupUserDefaults(let identifier):
            return try await UserDefaultsHelper.shared.get(forKey: key.name, appGroupIdentifier: identifier)
            
        case .keychain(let service):
            return try await KeychainHelper.shared.get(service: service, key: key.name)
            
        case .fileSystem(let directory), .encryptedFileSystem(let directory):
            return try await FileStorageHelper.shared.read(from: directory, fileName: key.name)
        case .appGroupFileSystem(let identifier, let directory):
            return try await FileStorageHelper.shared.read(
                from: directory,
                fileName: key.name,
                appGroupIdentifier: identifier
            )
        }
    }
    
    private func delete(for key: any StorageKey) async throws {
        switch key.domain {
        case .userDefaults(let suite):
            try await UserDefaultsHelper.shared.remove(forKey: key.name, suite: suite)
        case .appGroupUserDefaults(let identifier):
            try await UserDefaultsHelper.shared.remove(forKey: key.name, appGroupIdentifier: identifier)
            
        case .keychain(let service):
            try await KeychainHelper.shared.delete(service: service, key: key.name)
            
        case .fileSystem(let directory), .encryptedFileSystem(let directory):
            try await FileStorageHelper.shared.delete(from: directory, fileName: key.name)
        case .appGroupFileSystem(let identifier, let directory):
            try await FileStorageHelper.shared.delete(
                from: directory,
                fileName: key.name,
                appGroupIdentifier: identifier
            )
        }
    }
    
    // MARK: - Sync
    
    private func handleSync(_ key: any StorageKey, data: Data) async throws {
        try await SyncHelper.shared.syncIfNeeded(
            data: data,
            keyName: key.name,
            availability: key.availability,
            syncPolicy: key.syncPolicy
        )
    }
}