LocalData/Sources/LocalData/Services/StorageRouter.swift

import Foundation

/// 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 catalogRegistries: [String: [AnyStorageKey]] = [:]
    private var registeredKeys: [String: AnyStorageKey] = [:]
    private var migrationHistory: [String: Date] = [:]
    private var storageConfiguration: StorageConfiguration = .default
    private let keychain: KeychainStoring
    private let encryption: EncryptionHelper
    private let file: FileStorageHelper
    private let defaults: UserDefaultsHelper
    private let sync: SyncHelper
    
    /// Initialize a new StorageRouter.
    /// Internal for testing isolation via @testable import.
    /// Consumers should use the `shared` singleton.
    internal init(
        keychain: KeychainStoring = KeychainHelper.shared,
        encryption: EncryptionHelper = .shared,
        file: FileStorageHelper = .shared,
        defaults: UserDefaultsHelper = .shared,
        sync: SyncHelper = .shared
    ) {
        self.keychain = keychain
        self.encryption = encryption
        self.file = file
        self.defaults = defaults
        self.sync = sync
    }

    // 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 encryption.updateConfiguration(configuration)
        await encryption.updateKeychainHelper(keychain)
    }

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

    /// Updates the file storage configuration.
    public func updateFileStorageConfiguration(_ configuration: FileStorageConfiguration) async {
        await file.updateConfiguration(configuration)
    }

    /// Updates the global storage configuration (defaults).
    public func updateStorageConfiguration(_ configuration: StorageConfiguration) {
        self.storageConfiguration = 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 encryption.updateKeychainHelper(keychain)
        await encryption.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.
    /// - Parameters:
    ///   - catalog: The catalog type to register.
    ///   - migrateImmediately: If true, triggers a proactive migration (sweep) for all keys in the catalog.
    public func registerCatalog(_ catalog: some StorageKeyCatalog, migrateImmediately: Bool = false) async throws {
        let entries = catalog.allKeys
        try validateDescription(entries)
        try validateUniqueKeys(entries)
        
        // Validate against existing registrations to prevent collisions across catalogs
        let catalogName = catalog.name
        Logger.info(">>> [STORAGE] Registering Catalog: \(catalogName) (\(entries.count) keys)")

        for entry in entries {
            if let existing = registeredKeys[entry.descriptor.name] {
                let existingCatalog = existing.descriptor.catalog ?? "Unknown"
                let errorMessage = "STORAGE KEY COLLISION: Key name '\(entry.descriptor.name)' in \(catalogName) is already registered by \(existingCatalog)."
                Logger.error(errorMessage)
                throw StorageError.duplicateRegisteredKeys([entry.descriptor.name])
            }
        }

        // Add to registry with catalog name context
        catalogRegistries[catalogName] = entries
        for entry in entries {
            registeredKeys[entry.descriptor.name] = entry.withCatalog(catalogName)
        }
        
        Logger.info("<<< [STORAGE] Catalog Registered Successfully: \(catalogName)")
        
        if migrateImmediately {
            try await migrateAllRegisteredKeys()
        }
    }
    
    /// Returns all currently registered storage keys, grouped by catalog name.
    public func allRegisteredCatalogs() -> [String: [AnyStorageKey]] {
        catalogRegistries
    }
    
    /// Returns all currently registered storage keys.
    public func allRegisteredEntries() -> [AnyStorageKey] {
        Array(registeredKeys.values)
    }
    
    /// Triggers a proactive migration (sweep) for all registered storage keys.
    /// This "drains" any legacy data into the modern storage locations.
    public func migrateAllRegisteredKeys() async throws {
        Logger.debug(">>> [STORAGE] STARTING GLOBAL MIGRATION SWEEP")
        for entry in registeredKeys.values {
            try await entry.migrate(on: self)
        }
        Logger.debug("<<< [STORAGE] GLOBAL MIGRATION SWEEP COMPLETE")
    }

    /// Returns the last migration date for a specific key, if available.
    public func migrationHistory<Key: StorageKey>(for key: Key) -> Date? {
        migrationHistory[key.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 {
        Logger.debug(">>> [STORAGE] SET: \(key.name) [Domain: \(key.domain)]")
        try validateCatalogRegistration(for: key)
        try validatePlatformAvailability(for: key)
        
        let data = try serialize(value, with: key.serializer)
        let securedData = try await applySecurity(data, for: .from(key), isEncrypt: true)
        
        try await store(securedData, for: key)
        try await handleSync(key, data: securedData)
        Logger.debug("<<< [STORAGE] SET SUCCESS: \(key.name)")
    }
    
    /// 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 {
        Logger.debug(">>> [STORAGE] GET: \(key.name)")
        try validateCatalogRegistration(for: key)
        try validatePlatformAvailability(for: key)
        
        // 1. Try primary location
        if let securedData = try await retrieve(for: .from(key)) {
            let data = try await applySecurity(securedData, for: .from(key), isEncrypt: false)
            let result = try deserialize(data, with: key.serializer)
            Logger.debug("<<< [STORAGE] GET SUCCESS: \(key.name)")
            return result
        }
        
        // 2. Attempt migration (legacy or advanced)
        if let migratedValue = try await attemptMigration(for: key) {
            return migratedValue
        }
        
        Logger.debug("<<< [STORAGE] GET NOT FOUND: \(key.name)")
        throw StorageError.notFound
    }
    
    /// Manually triggers migration for a specific key.
    /// This is useful for "draining" old keys at app startup even if the destination already has data.
    /// - Parameter key: The storage key to migrate.
    /// - Returns: The migration result.
    /// - Throws: Migration or storage errors.
    public func forceMigration<Key: StorageKey>(for key: Key) async throws -> MigrationResult {
        Logger.debug(">>> [STORAGE] MANUAL MIGRATION: \(key.name)")
        try validateCatalogRegistration(for: key)

        guard let migration = resolveMigration(for: key) else {
            Logger.debug("<<< [STORAGE] MANUAL MIGRATION: No migration configured for \(key.name)")
            return MigrationResult(success: true, migratedCount: 0)
        }

        let context = buildMigrationContext()
        guard try await shouldAllowMigration(for: key, context: context) else {
            Logger.debug("<<< [STORAGE] MANUAL MIGRATION: Skipped for \(key.name) on this platform")
            return MigrationResult(success: true, migratedCount: 0)
        }
        let result = try await migration.migrate(using: self, context: context)
        if result.success {
            recordMigration(for: .from(key))
        }
        Logger.debug("<<< [STORAGE] MANUAL MIGRATION COMPLETE: \(key.name)")
        return result
    }
    
    /// 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: .from(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 defaults.exists(forKey: key.name, suite: suite)
        case .appGroupUserDefaults(let identifier):
            let resolvedId = try resolveIdentifier(identifier)
            return try await defaults.exists(forKey: key.name, appGroupIdentifier: resolvedId)
        case .keychain(let service):
            let resolvedService = try resolveService(service)
            return try await keychain.exists(service: resolvedService, key: key.name)
        case .fileSystem(let directory), .encryptedFileSystem(let directory):
            return await file.exists(in: directory, fileName: key.name)
        case .appGroupFileSystem(let identifier, let directory):
            let resolvedId = try resolveIdentifier(identifier)
            return await file.exists(
                in: directory,
                fileName: key.name,
                appGroupIdentifier: resolvedId
            )
        }
    }

    internal func exists(descriptor: StorageKeyDescriptor) async throws -> Bool {
        switch descriptor.domain {
        case .userDefaults(let suite):
            return try await defaults.exists(forKey: descriptor.name, suite: suite)
        case .appGroupUserDefaults(let identifier):
            let resolvedId = try resolveIdentifier(identifier)
            return try await defaults.exists(forKey: descriptor.name, appGroupIdentifier: resolvedId)
        case .keychain(let service):
            let resolvedService = try resolveService(service)
            return try await keychain.exists(service: resolvedService, key: descriptor.name)
        case .fileSystem(let directory), .encryptedFileSystem(let directory):
            return await file.exists(in: directory, fileName: descriptor.name)
        case .appGroupFileSystem(let identifier, let directory):
            let resolvedId = try resolveIdentifier(identifier)
            return await file.exists(
                in: directory,
                fileName: descriptor.name,
                appGroupIdentifier: resolvedId
            )
        }
    }
    
    // MARK: - Platform Validation
    
    nonisolated 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 !registeredKeys.isEmpty else { return }
        guard registeredKeys[key.name] != nil else {
            let errorMessage = "UNREGISTERED STORAGE KEY: '\(key.name)' (Owner: \(key.owner)) accessed but not found in any registered catalog. Did you forget to call registerCatalog?"
            Logger.error(errorMessage)
            #if DEBUG
            if !isRunningTests {
                assertionFailure(errorMessage)
            }
            #endif
            throw StorageError.unregisteredKey(key.name)
        }
    }

    private var isRunningTests: Bool {
        // Broad check for any test-related environment variables or classes
        if ProcessInfo.processInfo.environment.keys.contains(where: { 
            $0.hasPrefix("XCTest") || $0.hasPrefix("SWIFT_TESTING") || $0.hasPrefix("SWIFT_DETERMINISTIC")
        }) {
            return true
        }
        return NSClassFromString("XCTestCase") != nil || NSClassFromString("Testing.Test") != nil
    }

    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: - Migration

    private func attemptMigration<Key: StorageKey>(for key: Key) async throws -> Key.Value? {
        guard let migration = resolveMigration(for: key) else { return nil }

        let context = buildMigrationContext()
        let shouldMigrate = try await migration.shouldMigrate(using: self, context: context)
        guard shouldMigrate else { return nil }

        let result = try await migration.migrate(using: self, context: context)
        if result.success {
            recordMigration(for: .from(key))
            if let securedData = try await retrieve(for: .from(key)) {
                let data = try await applySecurity(securedData, for: .from(key), isEncrypt: false)
                return try deserialize(data, with: key.serializer)
            }
            return nil
        }

        if let error = result.errors.first {
            throw error
        }
        return nil
    }

    private func resolveMigration<Key: StorageKey>(for key: Key) -> AnyStorageMigration? {
        key.migration
    }

    internal func buildMigrationContext() -> MigrationContext {
        MigrationContext(migrationHistory: migrationHistory)
    }

    internal func recordMigration(for descriptor: StorageKeyDescriptor) {
        migrationHistory[descriptor.name] = Date()
    }

    internal func shouldAllowMigration<Key: StorageKey>(
        for key: Key,
        context: MigrationContext
    ) async throws -> Bool {
        guard key.availability.isAvailable(on: context.deviceInfo.platform) else {
            return false
        }

        if key.syncPolicy != .never {
            let isSyncAvailable = await sync.isSyncAvailable()
            guard isSyncAvailable else { return false }
        }
        return true
    }
    
    // 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
        }
    }
    
    nonisolated internal 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
    
    internal func applySecurity(
        _ data: Data,
        for descriptor: StorageKeyDescriptor,
        isEncrypt: Bool
    ) async throws -> Data {
        switch descriptor.security {
        case .none:
            return data
            
        case .encrypted(let encryptionPolicy):
            await encryption.updateKeychainHelper(keychain)
            if isEncrypt {
                return try await encryption.encrypt(
                    data,
                    keyName: descriptor.name,
                    policy: encryptionPolicy
                )
            } else {
                return try await encryption.decrypt(
                    data,
                    keyName: descriptor.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 {
        try await store(data, for: .from(key))
    }
    
    private func store(_ data: Data, for descriptor: StorageKeyDescriptor) async throws {
        switch descriptor.domain {
        case .userDefaults(let suite):
            try await defaults.set(data, forKey: descriptor.name, suite: suite)
            
        case .appGroupUserDefaults(let identifier):
            let resolvedId = try resolveIdentifier(identifier)
            try await defaults.set(data, forKey: descriptor.name, appGroupIdentifier: resolvedId)
            
        case .keychain(let service):
            guard case let .keychain(accessibility, accessControl) = descriptor.security else {
                throw StorageError.securityApplicationFailed
            }
            let resolvedService = try resolveService(service)
            try await keychain.set(
                data,
                service: resolvedService,
                key: descriptor.name,
                accessibility: accessibility,
                accessControl: accessControl
            )
            
        case .fileSystem(let directory):
            try await file.write(
                data,
                to: directory,
                fileName: descriptor.name,
                useCompleteFileProtection: false
            )
            
        case .encryptedFileSystem(let directory):
            try await file.write(
                data,
                to: directory,
                fileName: descriptor.name,
                useCompleteFileProtection: true
            )
            
        case .appGroupFileSystem(let identifier, let directory):
            let resolvedId = try resolveIdentifier(identifier)
            try await file.write(
                data,
                to: directory,
                fileName: descriptor.name,
                appGroupIdentifier: resolvedId,
                useCompleteFileProtection: false
            )
        }
    }
    
    internal func retrieve(for descriptor: StorageKeyDescriptor) async throws -> Data? {
        switch descriptor.domain {
        case .userDefaults(let suite):
            return try await defaults.get(forKey: descriptor.name, suite: suite)
        case .appGroupUserDefaults(let identifier):
            let resolvedId = try resolveIdentifier(identifier)
            return try await defaults.get(forKey: descriptor.name, appGroupIdentifier: resolvedId)
            
        case .keychain(let service):
            let resolvedService = try resolveService(service)
            return try await keychain.get(service: resolvedService, key: descriptor.name)
            
        case .fileSystem(let directory), .encryptedFileSystem(let directory):
            return try await file.read(from: directory, fileName: descriptor.name)
        case .appGroupFileSystem(let identifier, let directory):
            let resolvedId = try resolveIdentifier(identifier)
            return try await file.read(
                from: directory,
                fileName: descriptor.name,
                appGroupIdentifier: resolvedId
            )
        }
    }
    
    internal func delete(for descriptor: StorageKeyDescriptor) async throws {
        switch descriptor.domain {
        case .userDefaults(let suite):
            try await defaults.remove(forKey: descriptor.name, suite: suite)
        case .appGroupUserDefaults(let identifier):
            let resolvedId = try resolveIdentifier(identifier)
            try await defaults.remove(forKey: descriptor.name, appGroupIdentifier: resolvedId)
            
        case .keychain(let service):
            let resolvedService = try resolveService(service)
            try await keychain.delete(service: resolvedService, key: descriptor.name)
            
        case .fileSystem(let directory), .encryptedFileSystem(let directory):
            try await file.delete(from: directory, fileName: descriptor.name)
        case .appGroupFileSystem(let identifier, let directory):
            let resolvedId = try resolveIdentifier(identifier)
            try await file.delete(
                from: directory,
                fileName: descriptor.name,
                appGroupIdentifier: resolvedId
            )
        }
    }
    
    // MARK: - Sync
    
    private func handleSync(_ key: any StorageKey, data: Data) async throws {
        try await sync.syncIfNeeded(
            data: data,
            keyName: key.name,
            availability: key.availability,
            syncPolicy: key.syncPolicy
        )
    }

    /// Attempts to sync any registered keys that already have stored values.
    /// This is useful for bootstrapping watch data after app launch or reconnection.
    public func syncRegisteredKeysIfNeeded() async {
        let isAvailable = await sync.isSyncAvailable()
        guard isAvailable else {
            Logger.debug("<<< [SYNC] Skipping bootstrap sync: WatchConnectivity unavailable")
            return
        }

        Logger.debug(">>> [SYNC] Starting bootstrap sync for registered keys")
        for entry in registeredKeys.values {
            let descriptor = entry.descriptor
            guard descriptor.availability == .all || descriptor.availability == .phoneWithWatchSync else {
                Logger.debug("<<< [SYNC] Skipping key \(descriptor.name): availability=\(descriptor.availability)")
                continue
            }
            guard descriptor.syncPolicy != .never else {
                Logger.debug("<<< [SYNC] Skipping key \(descriptor.name): syncPolicy=\(descriptor.syncPolicy)")
                continue
            }

            do {
                guard let storedData = try await retrieve(for: descriptor) else {
                    Logger.debug("<<< [SYNC] No stored data for key \(descriptor.name)")
                    continue
                }
                try await sync.syncIfNeeded(
                    data: storedData,
                    keyName: descriptor.name,
                    availability: descriptor.availability,
                    syncPolicy: descriptor.syncPolicy
                )
                Logger.debug("<<< [SYNC] Bootstrapped context for key: \(descriptor.name) (\(storedData.count) bytes)")
            } catch {
                Logger.error("Failed to bootstrap sync for key: \(descriptor.name)", error: error)
            }
        }
    }

    /// Builds a snapshot of syncable key data for immediate watch requests.
    public func syncSnapshot() async -> [String: Data] {
        let isAvailable = await sync.isSyncAvailable()
        guard isAvailable else {
            Logger.debug("<<< [SYNC] Skipping snapshot: WatchConnectivity unavailable")
            return [:]
        }

        let maxAutoSyncSize = await sync.maxAutoSyncSize()
        var payload: [String: Data] = [:]

        for entry in registeredKeys.values {
            let descriptor = entry.descriptor
            guard descriptor.availability == .all || descriptor.availability == .phoneWithWatchSync else {
                continue
            }
            guard descriptor.syncPolicy != .never else { continue }

            do {
                guard let storedData = try await retrieve(for: descriptor) else { continue }

                if descriptor.syncPolicy == .automaticSmall, storedData.count > maxAutoSyncSize {
                    Logger.debug("<<< [SYNC] Snapshot skip \(descriptor.name): size exceeds maxAutoSyncSize")
                    continue
                }

                payload[descriptor.name] = storedData
            } catch {
                Logger.error("Failed to build sync snapshot for key: \(descriptor.name)", error: error)
            }
        }

        Logger.debug("<<< [SYNC] Snapshot built with \(payload.count) keys")
        return payload
    }

    // MARK: - Internal Sync Handling

    /// Internal method to update storage from received sync data.
    /// This is called by SyncHelper when the paired device sends new context.
    func updateFromSync(keyName: String, data: Data) async throws {
        // Find the registered entry for this key
        guard let entry = registeredKeys[keyName] else {
            Logger.debug("Received sync data for unregistered or uncatalogued key: \(keyName)")
            return
        }
        
        // The data received is already 'secured' (encrypted if necessary) by the sender.
        // We can store it directly in our local domain.
        try await store(data, for: entry.descriptor)
        Logger.info("Successfully updated local storage from sync for key: \(keyName)")
    }

    // MARK: - Resolution Helpers

    private func resolveService(_ service: String?) throws -> String {
        guard let resolved = service ?? storageConfiguration.defaultKeychainService else {
            Logger.error("No keychain service provided and no default configured")
            throw StorageError.keychainError(errSecBadReq) // Or a more specific error
        }
        Logger.debug("Resolved Keychain Service: \(resolved)")
        return resolved
    }

    private func resolveIdentifier(_ identifier: String?) throws -> String {
        guard let resolved = identifier ?? storageConfiguration.defaultAppGroupIdentifier else {
            Logger.error("No App Group identifier provided and no default configured")
            throw StorageError.invalidAppGroupIdentifier("none")
        }
        Logger.debug("Resolved App Group ID: \(resolved)")
        return resolved
    }
}