Swiftgram/TelegramCore/CollectCacheUsageStats.swift

import Foundation
#if os(macOS)
    import PostboxMac
    import SwiftSignalKitMac
#else
    import Postbox
    import SwiftSignalKit
#endif

public enum PeerCacheUsageCategory: Int32 {
    case image = 0
    case video
    case audio
    case file
}

public struct CacheUsageStats {
    public let media: [PeerId: [PeerCacheUsageCategory: [MediaId: Int64]]]
    public let mediaResourceIds: [MediaId: [MediaResourceId]]
    public let peers: [PeerId: Peer]
    public let otherSize: Int64
    public let otherPaths: [String]
    public let cacheSize: Int64
    
    public init(media: [PeerId: [PeerCacheUsageCategory: [MediaId: Int64]]], mediaResourceIds: [MediaId: [MediaResourceId]], peers: [PeerId: Peer], otherSize: Int64, otherPaths: [String], cacheSize: Int64) {
        self.media = media
        self.mediaResourceIds = mediaResourceIds
        self.peers = peers
        self.otherSize = otherSize
        self.otherPaths = otherPaths
        self.cacheSize = cacheSize
    }
}

public enum CacheUsageStatsResult {
    case progress(Float)
    case result(CacheUsageStats)
}

private enum CollectCacheUsageStatsError {
    case done(CacheUsageStats)
    case generic
}

private final class CacheUsageStatsState {
    var media: [PeerId: [PeerCacheUsageCategory: [MediaId: Int64]]] = [:]
    var mediaResourceIds: [MediaId: [MediaResourceId]] = [:]
    var allResourceIds = Set<WrappedMediaResourceId>()
    var lowerBound: MessageIndex?
}

public func collectCacheUsageStats(account: Account) -> Signal<CacheUsageStatsResult, NoError> {
    let state = Atomic<CacheUsageStatsState>(value: CacheUsageStatsState())
    
    let fetch = account.postbox.modify { modifier -> ([PeerId : Set<MediaId>], [MediaId : Media], MessageIndex?) in
        return modifier.enumerateMedia(lowerBound: state.with { $0.lowerBound }, limit: 1000)
    } |> mapError { _ -> CollectCacheUsageStatsError in preconditionFailure() }
    
    let process: ([PeerId : Set<MediaId>], [MediaId : Media], MessageIndex?) -> Signal<CacheUsageStatsResult, CollectCacheUsageStatsError> = { mediaByPeer, mediaRefs, updatedLowerBound in
        var mediaIdToPeerId: [MediaId: PeerId] = [:]
        for (peerId, mediaIds) in mediaByPeer {
            for id in mediaIds {
                mediaIdToPeerId[id] = peerId
            }
        }
        
        var resourceIdToMediaId: [WrappedMediaResourceId: (MediaId, PeerCacheUsageCategory)] = [:]
        var mediaResourceIds: [MediaId: [MediaResourceId]] = [:]
        var resourceIds: [MediaResourceId] = []
        for (id, media) in mediaRefs {
            mediaResourceIds[id] = []
            switch media {
                case let image as TelegramMediaImage:
                    for representation in image.representations {
                        resourceIds.append(representation.resource.id)
                        resourceIdToMediaId[WrappedMediaResourceId(representation.resource.id)] = (id, .image)
                        mediaResourceIds[id]!.append(representation.resource.id)
                    }
                case let file as TelegramMediaFile:
                    var category: PeerCacheUsageCategory = .file
                    loop: for attribute in file.attributes {
                        switch attribute {
                            case .Video:
                                category = .video
                                break loop
                            case .Audio:
                                category = .audio
                                break loop
                            default:
                                break
                        }
                    }
                    for representation in file.previewRepresentations {
                        resourceIds.append(representation.resource.id)
                        resourceIdToMediaId[WrappedMediaResourceId(representation.resource.id)] = (id, category)
                        mediaResourceIds[id]!.append(representation.resource.id)
                    }
                    resourceIds.append(file.resource.id)
                    resourceIdToMediaId[WrappedMediaResourceId(file.resource.id)] = (id, category)
                    mediaResourceIds[id]!.append(file.resource.id)
                default:
                    break
            }
        }
        return account.postbox.mediaBox.collectResourceCacheUsage(resourceIds)
            |> mapError { _ -> CollectCacheUsageStatsError in preconditionFailure() }
            |> mapToSignal { result -> Signal<CacheUsageStatsResult, CollectCacheUsageStatsError> in
                state.with { state -> Void in
                    state.lowerBound = updatedLowerBound
                    for (wrappedId, size) in result {
                        if let (id, category) = resourceIdToMediaId[wrappedId] {
                            if let peerId = mediaIdToPeerId[id] {
                                if state.media[peerId] == nil {
                                    state.media[peerId] = [:]
                                }
                                if state.media[peerId]![category] == nil {
                                    state.media[peerId]![category] = [:]
                                }
                                var currentSize: Int64 = 0
                                if let current = state.media[peerId]![category]![id] {
                                    currentSize = current
                                }
                                state.media[peerId]![category]![id] = currentSize + size
                            }
                        }
                    }
                    for (id, ids) in mediaResourceIds {
                        state.mediaResourceIds[id] = ids
                        for resourceId in ids {
                            state.allResourceIds.insert(WrappedMediaResourceId(resourceId))
                        }
                    }
                }
                if updatedLowerBound == nil {
                    let (finalMedia, finalMediaResourceIds, allResourceIds) = state.with { state -> ([PeerId: [PeerCacheUsageCategory: [MediaId: Int64]]], [MediaId: [MediaResourceId]], Set<WrappedMediaResourceId>) in
                        return (state.media, state.mediaResourceIds, state.allResourceIds)
                    }
                    
                    return account.postbox.mediaBox.collectOtherResourceUsage(excludeIds: allResourceIds)
                    |> mapError { _ in return CollectCacheUsageStatsError.generic }
                    |> mapToSignal { otherSize, otherPaths, cacheSize in
                        return account.postbox.modify { modifier -> CacheUsageStats in
                            var peers: [PeerId: Peer] = [:]
                            for peerId in finalMedia.keys {
                                if let peer = modifier.getPeer(peerId) {
                                    peers[peer.id] = peer
                                }
                            }
                            return CacheUsageStats(media: finalMedia, mediaResourceIds: finalMediaResourceIds, peers: peers, otherSize: otherSize, otherPaths: otherPaths, cacheSize: cacheSize)
                        } |> mapError { _ -> CollectCacheUsageStatsError in preconditionFailure() }
                        |> mapToSignal { stats -> Signal<CacheUsageStatsResult, CollectCacheUsageStatsError> in
                            return .fail(.done(stats))
                        }
                    }
                } else {
                    return .complete()
                }
            }
    }
    
    let signal = (fetch |> mapToSignal { mediaByPeer, mediaRefs, updatedLowerBound -> Signal<CacheUsageStatsResult, CollectCacheUsageStatsError> in
        return process(mediaByPeer, mediaRefs, updatedLowerBound)
    }) |> restart
    
    return signal |> `catch` { error in
        switch error {
            case let .done(result):
                return .single(.result(result))
            case .generic:
                return .complete()
        }
    }
}

public func clearCachedMediaResources(account: Account, mediaResourceIds: Set<WrappedMediaResourceId>) -> Signal<Void, NoError> {
    return account.postbox.mediaBox.removeCachedResources(mediaResourceIds)
}