Kisani/Services/BackupStatusService.swift

import Foundation
import Photos
import UIKit
import os.log

private let log = Logger(subsystem: "com.albert.nasbackup", category: "BackupStatus")
private let signposter = OSSignposter(subsystem: "com.albert.nasbackup", category: "BackupStatus")

@MainActor
final class BackupStatusService: NSObject, ObservableObject {
    static let shared = BackupStatusService()

    @Published private(set) var snapshot: BackupStatusSnapshot = .empty
    @Published private(set) var isRefreshing: Bool = false
    @Published private(set) var refreshError: String?
    @Published private(set) var lastManifest: BackupManifest?

    private let photoService = PhotoLibraryService()
    private let cacheKey = "backupStatusSnapshot_v2"

    // One active reconciliation task at a time — cancelled before starting a new one.
    private var refreshTask: Task<Void, Never>?
    private var debounceTask: Task<Void, Never>?

    private var lastFullRefresh: Date?
    private static let debounceInterval: TimeInterval = 4
    private static let minRefreshInterval: TimeInterval = 30

    // MARK: — Init

    private override init() {
        super.init()
        loadCachedSnapshot()
        PHPhotoLibrary.shared().register(self)
        NotificationCenter.default.addObserver(
            self,
            selector: #selector(handleMemoryWarning),
            name: UIApplication.didReceiveMemoryWarningNotification,
            object: nil
        )
    }

    deinit {
        PHPhotoLibrary.shared().unregisterChangeObserver(self)
        NotificationCenter.default.removeObserver(self)
    }

    @objc private func handleMemoryWarning() {
        log.warning("Memory warning — cancelling in-flight refresh")
        cancelAll()
        isRefreshing = false
    }

    // MARK: — Public API

    func refresh(force: Bool = false) {
        cancelAll()
        refreshTask = Task { [weak self] in
            await self?.performRefresh(force: force)
        }
    }

    func refreshAndWait(force: Bool = true) async {
        cancelAll()
        await performRefresh(force: force)
    }

    func refreshAfterBackup(entries: [ManifestEntry], connection: NASConnection) {
        let safe = entries.count
        // Only count files actually uploaded (fileSize > 0) toward nasArchiveTotal.
        // Collision-skipped and fileExists-skipped entries have fileSize 0 and were
        // already counted in the last directory scan — adding them would double-count.
        let newlyOnNAS = entries.filter { $0.fileSize > 0 }.count
        snapshot.alreadySafe = min(snapshot.phoneTotal, snapshot.alreadySafe + safe)
        snapshot.nasArchiveTotal += newlyOnNAS
        saveSnapshot()
        log.info("Post-backup optimistic update: +\(safe) safe, +\(newlyOnNAS) NAS")
        Task {
            await writeManifest(entries: entries, connection: connection)
            refresh(force: true)
        }
    }

    // MARK: — Refresh lifecycle

    private func cancelAll() {
        refreshTask?.cancel()
        refreshTask = nil
        debounceTask?.cancel()
        debounceTask = nil
    }

    private func performRefresh(force: Bool) async {
        guard !isRefreshing else {
            log.debug("Refresh skipped — already in progress")
            return
        }

        if !force, let last = lastFullRefresh,
           Date().timeIntervalSince(last) < Self.minRefreshInterval {
            log.debug("Refresh debounced — updating phone count only")
            updatePhoneCountOnly()
            return
        }

        // Ensure LocalPhotoIndex is loaded from disk before reconcile.
        // loadOrBuild() is a fast no-op when already in memory; on cold launch it
        // reads the cached JSON so the fast path (skip NAS roundtrip) can fire.
        await LocalPhotoIndex.shared.loadOrBuild()

        lastFullRefresh = Date()
        isRefreshing = true
        refreshError = nil
        log.info("Refresh started (force=\(force))")

        do {
            try await reconcile()
            log.info("Refresh completed — phone=\(self.snapshot.phoneTotal) safe=\(self.snapshot.alreadySafe)")
        } catch is CancellationError {
            log.info("Refresh cancelled")
        } catch BackupError.timeout {
            // NAS did not respond within the timeout window.
            // Retain cached counts so the UI shows useful data instead of zeroes.
            log.warning("Refresh timed out — retaining cached state")
            refreshError = "NAS not responding — using last known state"
            snapshot.lastCheckedAt = Date()
            saveSnapshot()
        } catch {
            log.error("Refresh failed: \(error.localizedDescription)")
            refreshError = error.localizedDescription
            snapshot.connectionState = .offline
            snapshot.lastCheckedAt = Date()
            saveSnapshot()
        }

        isRefreshing = false
    }

    // MARK: — Reconciliation

    private func reconcile() async throws {
        let spReconcile = signposter.beginInterval("Reconcile")
        defer { signposter.endInterval("Reconcile", spReconcile) }

        let store = ConnectionStore.shared
        guard let conn = store.savedConnection else { return }

        let filter = store.backupFilter

        log.info("""
            [reconcile] host=\(conn.host, privacy: .public) remotePath=\(conn.remotePath, privacy: .public) \
            manifestPath=\(conn.remotePath, privacy: .public)/\(BackupManifest.remoteFilename, privacy: .public) \
            filter=photos:\(filter.includePhotos) videos:\(filter.includeVideos) \
            screenshots:\(filter.includeScreenshots) raw:\(filter.includeRAW) \
            autoBackupEnabled=\(store.autoBackupEnabled) autoBackupOnOpen=\(store.autoBackupOnOpen)
            """)


        // ── Fast path: use local caches when both are warm ────────────────
        // Skips NAS connection + PHFetchResult batch enumeration entirely.
        // Triggered when: manifest cache < 5 min old AND LocalPhotoIndex has data.
        let manifestCacheIsStale = await NASManifestCache.shared.isStale
        let localIndexCount = await LocalPhotoIndex.shared.totalCount

        fastPath: if !manifestCacheIsStale, let cached = await NASManifestCache.shared.manifest, localIndexCount > 0 {
            let dirCount = await NASManifestCache.shared.directoryCount ?? 0
            let index: ManifestIndex
            if dirCount > cached.entries.count {
                // NAS has more files than the manifest tracks. Build a hybrid index that keeps
                // manifest localIdentifiers (for Kisani uploads) and adds all cached NAS
                // filenames as a filename-fallback. If filenames aren't cached yet, fall
                // through to a full NAS reconcile so they get populated.
                guard let filenames = await NASManifestCache.shared.directoryFilenames else {
                    log.info("Reconcile (cached): NAS dir(\(dirCount)) > manifest(\(cached.entries.count)) — no cached filenames, falling to full reconcile")
                    break fastPath
                }
                index = await Task.detached(priority: .utility) {
                    ManifestIndex(manifest: cached, nasFilenames: filenames, totalCount: dirCount)
                }.value
            } else {
                index = await Task.detached(priority: .utility) {
                    ManifestIndex(manifest: cached)
                }.value
            }
            let safe = await LocalPhotoIndex.shared.countSafe(against: index, filter: filter)
            let phoneTotal = await LocalPhotoIndex.shared.count(filter: filter)
            let nasTotal = max(index.totalCount, dirCount)
            var updated = BackupStatusSnapshot()
            updated.phoneTotal      = phoneTotal
            updated.alreadySafe     = min(safe, phoneTotal)
            updated.nasArchiveTotal = nasTotal
            updated.connectionState = .connected
            updated.lastCheckedAt   = Date()
            snapshot = updated
            saveSnapshot()
            log.info("Reconcile (cached): phone=\(phoneTotal) safe=\(safe) NAS=\(nasTotal) (manifest=\(index.totalCount) dir=\(dirCount))")
            return
        }

        // ── Full path: connect to NAS and download fresh manifest ─────────

        // Step 1: Phone count (O(1) PHFetchResult.count)
        let fetchResult = photoService.fetchResultForReconciliation(filter: filter)
        let phoneTotal = fetchResult.count
        log.info("PHFetchResult count: \(phoneTotal)")

        try Task.checkCancellation()

        // ── Step 2: NAS connection ────────────────────────────────────────
        let transfer: any NASTransferProtocol = conn.nasProtocol == .smb ? SMBService() : SFTPService()
        do {
            try await Self.withTimeout(15) {
                try await transfer.connect(
                    to: conn.host, port: conn.port,
                    username: conn.username, password: conn.password
                )
            }
        } catch {
            let reason: String
            if case BackupError.timeout = error {
                reason = "NAS not responding — using last known state"
            } else {
                reason = (error as? BackupError).flatMap { $0.errorDescription } ?? error.localizedDescription
            }
            log.warning("NAS connect failed: \(reason, privacy: .public) — retaining cached counts")
            refreshError = reason

            // Restore last-known NAS archive count so the dashboard
            // doesn't show 0 when the NAS is temporarily offline.
            let cachedCount = (await NASManifestCache.shared.manifest).map { ManifestIndex(manifest: $0).totalCount } ?? 0
            let cachedDirCount = await NASManifestCache.shared.directoryCount ?? 0
            snapshot.nasArchiveTotal = max(snapshot.nasArchiveTotal, cachedCount, cachedDirCount)
            log.info("NAS offline — restored nasArchiveTotal=\(self.snapshot.nasArchiveTotal) (manifest=\(cachedCount) dir=\(cachedDirCount))")
            snapshot.phoneTotal = phoneTotal
            snapshot.alreadySafe = min(snapshot.alreadySafe, phoneTotal)
            snapshot.connectionState = .offline
            snapshot.lastCheckedAt = Date()
            saveSnapshot()
            return
        }
        defer { transfer.disconnect() }

        try Task.checkCancellation()

        // ── Step 3: Load ManifestIndex ────────────────────────────────────
        let spManifest = signposter.beginInterval("ManifestLoad")
        let manifestPath = "\(conn.remotePath)/\(BackupManifest.remoteFilename)"
        let index = await buildManifestIndex(transfer: transfer, path: manifestPath, conn: conn)
        signposter.endInterval("ManifestLoad", spManifest,
                               "\(index.totalCount) entries, filename-only=\(index.isFilenameOnly)")
        log.info("Manifest index built — \(index.totalCount) NAS entries")

        try Task.checkCancellation()

        // ── Step 4: Count safe assets ─────────────────────────────────────
        // Prefer LocalPhotoIndex (pure in-memory, no CoreData) if populated;
        // otherwise fall back to batch PHFetchResult enumeration.
        let spCount = signposter.beginInterval("CountSafe")
        let localCount = await LocalPhotoIndex.shared.totalCount
        let safe: Int
        if localCount > 0 {
            safe = await LocalPhotoIndex.shared.countSafe(against: index, filter: filter)
        } else {
            safe = try await Task.detached(priority: .userInitiated) { [weak self, fetchResult, index] in
                guard let self else { throw CancellationError() }
                return try await self.countSafe(in: fetchResult, against: index)
            }.value
        }
        signposter.endInterval("CountSafe", spCount, "\(safe)/\(phoneTotal) safe")
        log.info("Comparison complete — \(safe)/\(phoneTotal) already safe")

        // ── Step 5: Commit ────────────────────────────────────────────────
        let dirCount = await NASManifestCache.shared.directoryCount ?? 0
        let nasTotal = max(index.totalCount, dirCount)
        var updated = BackupStatusSnapshot()
        updated.phoneTotal      = phoneTotal
        updated.alreadySafe     = min(safe, phoneTotal)
        updated.nasArchiveTotal = nasTotal
        updated.connectionState = .connected
        updated.lastCheckedAt   = Date()
        snapshot = updated
        saveSnapshot()
        log.info("Reconcile (full): phone=\(phoneTotal) safe=\(safe) NAS=\(nasTotal) (manifest=\(index.totalCount) dir=\(dirCount))")

        // Invariant check — alreadySafe + needBackup must always equal phoneTotal
        let inv = updated.alreadySafe + updated.needBackup
        if inv != updated.phoneTotal {
            log.error("Invariant broken: alreadySafe(\(updated.alreadySafe)) + needBackup(\(updated.needBackup)) = \(inv) ≠ phoneTotal(\(updated.phoneTotal))")
        }
    }

    /// Builds a ManifestIndex from NAS.
    /// - The download suspends the main actor.
    /// - JSON decode and Set construction run in a background task.
    private func buildManifestIndex(
        transfer: any NASTransferProtocol,
        path: String,
        conn: NASConnection
    ) async -> ManifestIndex {
        log.info("buildManifestIndex: remotePath=\(conn.remotePath, privacy: .public) manifestFile=\(path, privacy: .public)")
        do {
            let data = try await Self.withTimeout(12) { try await transfer.downloadData(at: path) }
            log.info("buildManifestIndex: manifest downloaded — \(data.count) bytes")
            // Decode and build index off main actor — JSONDecoder is not cheap for large manifests.
            let decoded: (ManifestIndex, BackupManifest)? = await Task.detached(priority: .utility) {
                guard let manifest = try? JSONDecoder.kisani.decode(BackupManifest.self, from: data)
                else { return nil }
                return (ManifestIndex(manifest: manifest), manifest)
            }.value
            if let (index, manifest) = decoded {
                log.info("buildManifestIndex: decoded \(index.totalCount) entries (isFilenameOnly=\(index.isFilenameOnly))")
                lastManifest = manifest
                Task { await NASManifestCache.shared.update(manifest) }

                // Always scan directory so the NAS Archive count is accurate on every
                // full reconcile (not just when the manifest is new/sparse).
                // When the directory has more files than the manifest (uploaded by other
                // tools), switch to filename-based matching so countSafe reflects real
                // NAS coverage instead of only Kisani-tracked files.
                let items = (try? await Self.withTimeout(12) { try await transfer.listDirectory(at: conn.remotePath) }) ?? []
                let filteredItems = items.filter {
                    !$0.isDirectory && !$0.name.hasPrefix(".") && $0.name != BackupManifest.remoteFilename
                }
                let dirCount = filteredItems.count
                log.info("buildManifestIndex: dir=\(dirCount) manifest=\(index.totalCount)")
                if dirCount > index.totalCount {
                    // NAS has more files than manifest: keep manifest localIdentifiers
                    // for Kisani-tracked uploads and add all NAS filenames as fallback.
                    // This prevents discarding accumulated localIdentifier history.
                    let nasNames = filteredItems.map { $0.name }
                    await NASManifestCache.shared.setDirectoryFilenames(nasNames)
                    return await Task.detached(priority: .utility) {
                        ManifestIndex(manifest: manifest, nasFilenames: nasNames, totalCount: dirCount)
                    }.value
                } else {
                    await NASManifestCache.shared.setDirectoryCount(dirCount)
                    return index
                }
            }
            log.warning("buildManifestIndex: manifest downloaded but failed to decode (\(data.count) bytes)")
        } catch {
            log.info("buildManifestIndex: manifest not found at \(path, privacy: .public) — \(error.localizedDescription, privacy: .public) — falling back to directory listing")
        }

        // Bootstrap: list NAS directory and build filename-only index in background.
        let items = (try? await Self.withTimeout(12) { try await transfer.listDirectory(at: conn.remotePath) }) ?? []
        log.info("buildManifestIndex: bootstrap from \(items.count) items in \(conn.remotePath, privacy: .public)")
        let bootstrapIndex = await Task.detached(priority: .utility) {
            ManifestIndex(nasListing: items)
        }.value
        await NASManifestCache.shared.setDirectoryCount(bootstrapIndex.totalCount)
        return bootstrapIndex
    }

    /// Enumerates a PHFetchResult in batches using autoreleasepool.
    /// Marked nonisolated so it can be called from a background Task without
    /// requiring main-actor scheduling — it accesses no @MainActor state.
    private nonisolated func countSafe(
        in result: PHFetchResult<PHAsset>,
        against index: ManifestIndex,
        batchSize: Int = 150
    ) async throws -> Int {
        var safe = 0
        let total = result.count
        var processed = 0

        while processed < total {
            try Task.checkCancellation()

            let batchEnd = min(processed + batchSize, total)

            // autoreleasepool releases PHAsset objects and PHAssetResource arrays
            // created during this batch before moving to the next.
            let batchSafe: Int = autoreleasepool {
                var count = 0
                for i in processed..<batchEnd {
                    let asset = result.object(at: i)
                    if index.matches(localIdentifier: asset.localIdentifier) {
                        count += 1
                    } else {
                        let resources = PHAssetResource.assetResources(for: asset)
                        if let name = resources.first?.originalFilename,
                           index.matches(filename: name) {
                            count += 1
                        }
                    }
                }
                return count
            }

            safe += batchSafe
            processed = batchEnd
            await Task.yield()
        }

        return safe
    }

    // MARK: — Phone-only lightweight update (no NAS connection)

    private func updatePhoneCountOnly() {
        let result = photoService.fetchResultForReconciliation(
            filter: ConnectionStore.shared.backupFilter
        )
        let phoneTotal = result.count
        guard phoneTotal != snapshot.phoneTotal else { return }
        snapshot.phoneTotal = phoneTotal
        snapshot.alreadySafe = min(snapshot.alreadySafe, phoneTotal)
        saveSnapshot()
        log.debug("Phone-only update: \(phoneTotal) assets")
    }

    // MARK: — Manifest write (non-blocking, background only)

    func writeManifest(entries: [ManifestEntry], connection: NASConnection) async {
        guard !entries.isEmpty else { return }

        // Capture cache before connecting — used as fallback if NAS download fails.
        // This prevents discarding accumulated history when the connection is flaky.
        let cacheBase = await NASManifestCache.shared.manifest

        let transfer: any NASTransferProtocol = connection.nasProtocol == .smb ? SMBService() : SFTPService()
        do {
            try await transfer.connect(
                to: connection.host, port: connection.port,
                username: connection.username, password: connection.password
            )
            defer { transfer.disconnect() }

            let manifestPath = "\(connection.remotePath)/\(BackupManifest.remoteFilename)"
            log.info("writeManifest: path=\(manifestPath, privacy: .public) entries=\(entries.count)")

            // Try NAS download — authoritative if successful.
            let existingData = try? await transfer.downloadData(at: manifestPath)

            // Merge + JSON encode off main actor — can be slow for large manifests.
            // Use the manifest with more entries as the merge base so history is never lost
            // if the NAS manifest has fewer entries than the local cache (e.g. 1 vs 264).
            let result: (data: Data, manifest: BackupManifest)? = await Task.detached(priority: .utility) {
                let nasManifest = existingData.flatMap { try? JSONDecoder.kisani.decode(BackupManifest.self, from: $0) }
                let nasCount   = nasManifest?.entries.count ?? 0
                let cacheCount = cacheBase?.entries.count ?? 0

                var manifest: BackupManifest
                if nasCount > 0, nasCount >= cacheCount {
                    manifest = nasManifest!   // NAS at least as current — authoritative
                } else if cacheCount > 0 {
                    manifest = cacheBase!     // Cache has more history — preserve it
                } else {
                    manifest = BackupManifest()
                }
                manifest.merge(entries: entries)
                guard let encoded = try? JSONEncoder.kisani.encode(manifest) else { return nil }
                return (encoded, manifest)
            }.value

            guard let (encoded, mergedManifest) = result else { return }

            // Update local cache BEFORE the NAS write — decouples reconcile correctness
            // from NAS write success. If the write fails/times out, the next fast-path
            // reconcile still reads the correct merged manifest from local cache.
            await NASManifestCache.shared.update(mergedManifest)
            lastManifest = mergedManifest

            try await transfer.writeData(encoded, to: manifestPath)
            log.info("writeManifest: done — \(entries.count) new, total=\(mergedManifest.entries.count)")
        } catch {
            log.error("writeManifest failed (non-fatal): \(error.localizedDescription)")
        }
    }

    // MARK: — Timeout helper

    /// Runs `work` and throws `BackupError.timeout` if it doesn't finish within `seconds`.
    /// Uses a racing task group so the timeout fires even when the underlying SMB/SFTP
    /// stack doesn't propagate Swift cooperative cancellation.
    private nonisolated static func withTimeout<T: Sendable>(
        _ seconds: TimeInterval,
        work: @Sendable @escaping () async throws -> T
    ) async throws -> T {
        try await withThrowingTaskGroup(of: T.self) { group in
            group.addTask(priority: .userInitiated) { try await work() }
            group.addTask(priority: .utility) {
                try await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000))
                throw BackupError.timeout
            }
            defer { group.cancelAll() }
            guard let result = try await group.next() else { throw BackupError.timeout }
            return result
        }
    }

    // MARK: — Cache

    private func loadCachedSnapshot() {
        guard let data = UserDefaults.standard.data(forKey: cacheKey),
              let cached = try? JSONDecoder().decode(BackupStatusSnapshot.self, from: data)
        else { return }
        snapshot = cached
    }

    private func saveSnapshot() {
        let encoder = JSONEncoder()
        encoder.dateEncodingStrategy = .iso8601
        if let data = try? encoder.encode(snapshot) {
            UserDefaults.standard.set(data, forKey: cacheKey)
        }
    }
}

// MARK: — PHPhotoLibraryChangeObserver

extension BackupStatusService: PHPhotoLibraryChangeObserver {
    nonisolated func photoLibraryDidChange(_ changeInstance: PHChange) {
        // Update LocalPhotoIndex incrementally — avoids full rebuild on every change.
        // We need the fetch result for change details; rebuild from scratch if unavailable.
        Task(priority: .utility) {
            let opts = PHFetchOptions()
            let allAssets = PHAsset.fetchAssets(with: opts)
            if let details = changeInstance.changeDetails(for: allAssets) {
                await LocalPhotoIndex.shared.apply(
                    inserted: details.insertedObjects,
                    removed:  details.removedObjects,
                    changed:  details.changedObjects
                )
                await LocalPhotoIndex.shared.flush()
            }
        }
        Task { @MainActor [weak self] in
            guard let self else { return }
            self.debounceTask?.cancel()
            self.debounceTask = Task {
                do {
                    try await Task.sleep(nanoseconds: UInt64(BackupStatusService.debounceInterval * 1_000_000_000))
                    self.refresh()
                } catch {
                    // Cancelled — a newer change event took over
                }
            }
        }
    }
}