stemify/ios/Stemify/Cache/Downloads.swift

import Foundation
import Observation
#if canImport(UIKit) && !SWIFT_PACKAGE
import UIKit
#endif

struct DownloadRecord: Codable, Equatable {
    var track: Track
    var keys: [String] = []
    var stems = false
    var done = false
    var error: String?
    var added = Int64(Date().timeIntervalSince1970 * 1000)

    init(track: Track) { self.track = track }

    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: AnyKey.self)
        track = try c.decode(Track.self, forKey: AnyKey("track"))
        keys = (try? c.decode([String].self, forKey: AnyKey("keys"))) ?? []
        stems = c.bool("stems") ?? false
        done = c.bool("done") ?? false
        error = c.str("error")
        added = c.int64("added") ?? 0
    }
}

/// Offline copies: the track's stems (or the original file) fully fetched into
/// [MediaCache] and pinned, plus the metadata needed to browse them with no server.
@Observable
final class Downloads: @unchecked Sendable {
    @ObservationIgnored private let settings: Settings
    @ObservationIgnored private let plex: PlexClient
    @ObservationIgnored private let stems: StemServer
    @ObservationIgnored private let cache: MediaCache
    @ObservationIgnored private let net: Connectivity
    @ObservationIgnored private let file: URL
    @ObservationIgnored private let lock = NSCondition()
    @ObservationIgnored private var state: [String: DownloadRecord] = [:]
    @ObservationIgnored private var prog: [String: Float] = [:]
    @ObservationIgnored private var queue: [String] = []
    @ObservationIgnored private var cancelled = Set<String>()
    @ObservationIgnored private let writeQueue = DispatchQueue(label: "downloads-save", qos: .utility)
    #if canImport(UIKit) && !SWIFT_PACKAGE
    @ObservationIgnored private var bgTask = UIBackgroundTaskIdentifier.invalid
    #endif
    static let waveBuckets = 256

    private(set) var records: [String: DownloadRecord] = [:]
    /// ratingKey -> 0..1 for queued and running downloads.
    private(set) var progress: [String: Float] = [:]

    init(settings: Settings, plex: PlexClient, stems: StemServer, cache: MediaCache, net: Connectivity, dir: URL) {
        self.settings = settings
        self.plex = plex
        self.stems = stems
        self.cache = cache
        self.net = net
        file = dir.appendingPathComponent("downloads.json")
        if let data = try? Data(contentsOf: file), let list = try? JSONDecoder().decode([Lenient<DownloadRecord>].self, from: data) {
            for r in list.compactMap(\.value) { state[r.track.ratingKey] = r }
        }
        records = state
        for _ in 0..<2 {
            let t = Thread { [weak self] in self?.work() }
            t.qualityOfService = .utility
            t.start()
        }
        // Resume anything interrupted last time.
        state.values.filter { !$0.done && $0.error == nil }.forEach { enqueue($0.track.ratingKey) }
    }

    func isDownloaded(_ ratingKey: String) -> Bool { records[ratingKey]?.done == true }

    func download(_ tracks: [Track]) {
        lock.lock()
        let add = tracks.filter { t in
            guard let r = state[t.ratingKey] else { return true }
            return !r.done && r.error != nil
        }
        lock.unlock()
        if add.isEmpty { return }
        update { m in for t in add { m[t.ratingKey] = DownloadRecord(track: t) } }
        add.forEach { enqueue($0.ratingKey) }
    }

    func remove(_ tracks: [Track]) {
        let keys = Set(tracks.map(\.ratingKey))
        lock.lock()
        cancelled.formUnion(keys)
        let gone = state.filter { keys.contains($0.key) }.values
        queue.removeAll { keys.contains($0) }
        for k in keys { prog[k] = nil }
        lock.unlock()
        update { m in for k in keys { m[k] = nil } }
        publish()
        DispatchQueue.global(qos: .utility).async { [cache] in gone.forEach { $0.keys.forEach(cache.unpin) } }
    }

    func retryFailed() {
        lock.lock()
        let failed = state.values.filter { $0.error != nil }
        lock.unlock()
        update { m in for r in failed { m[r.track.ratingKey]?.error = nil } }
        failed.forEach { enqueue($0.track.ratingKey) }
    }

    private func enqueue(_ key: String) {
        lock.lock()
        cancelled.remove(key)
        prog[key] = 0
        if !queue.contains(key) { queue.append(key) }
        lock.broadcast()
        lock.unlock()
        publish()
    }

    private func isCancelled(_ key: String) -> Bool {
        lock.lock(); defer { lock.unlock() }
        return cancelled.contains(key)
    }

    private func work() {
        while true {
            lock.lock()
            while queue.isEmpty { lock.wait() }
            let key = queue.removeFirst()
            lock.unlock()
            run(key)
        }
    }

    private func setProgress(_ key: String, _ f: Float?) {
        lock.lock()
        if f == nil || prog[key] != nil { prog[key] = f }
        lock.unlock()
        publish()
    }

    private func run(_ key: String) {
        lock.lock()
        let rec = state[key]
        lock.unlock()
        guard let rec, !rec.done else { setProgress(key, nil); return }
        defer { setProgress(key, nil) }
        do {
            while settings.downloadsWifiOnly && !net.unmetered {
                if isCancelled(key) { return }
                Thread.sleep(forTimeInterval: 2)
            }
            let t = rec.track
            let entry = stems.find(t)
            let sources: [(String, URL?)]
            if let entry {
                let q = stems.effectiveQuality(settings.stemQuality)
                stems.prepare(entry, q)
                sources = entry.stems.map { (PlaybackManager.stemKey(entry, $0.file, $0.size, q), stems.fileUrl(entry, $0, q)) }
            } else {
                guard let part = t.partKey else { throw PlainError("No file on the server") }
                sources = [(PlaybackManager.plexKey(t), plex.partUrl(part))]
            }
            update { m in
                m[key]?.keys = sources.map(\.0)
                m[key]?.stems = entry != nil
            }
            for (i, (cacheKey, url)) in sources.enumerated() {
                let res = cache.acquire(cacheKey) { url }
                defer { cache.release(res) }
                res.pinned = true
                var lastPublish = Date.distantPast
                try res.downloadAll(progress: { got, len in
                    let f: Float = len > 0 ? Float(got) / Float(len) : 0
                    if Date().timeIntervalSince(lastPublish) > 0.25 {
                        lastPublish = Date()
                        self.setProgress(key, (Float(i) + f) / Float(sources.count))
                    }
                }, cancelled: { self.isCancelled(key) })
            }
            if isCancelled(key) { return }
            if let entry { Task { _ = await stems.peaks(entry, buckets: Self.waveBuckets) } } // warms the HTTP cache for offline
            if let art = plex.artUrl(t.thumb, px: 600) { ArtLoader.shared.prefetch(art) }
            update { m in
                m[key]?.done = true
                m[key]?.error = nil
            }
        } catch {
            if isCancelled(key) { return }
            update { m in m[key]?.error = error.localizedDescription.isEmpty ? "Download failed" : error.localizedDescription }
        }
    }

    private func update(_ f: (inout [String: DownloadRecord]) -> Void) {
        lock.lock()
        f(&state)
        let snapshot = state
        lock.unlock()
        DispatchQueue.main.async { self.records = snapshot }
        writeQueue.async { [file] in
            let list = snapshot.values.sorted { $0.added < $1.added }
            if let data = try? JSONEncoder().encode(list) { try? data.write(to: file, options: .atomic) }
        }
    }

    private func publish() {
        lock.lock()
        let p = prog
        lock.unlock()
        DispatchQueue.main.async {
            if self.progress != p { self.progress = p }
            self.holdBackground(!p.isEmpty)
        }
    }

    /// iOS has no foreground service; ask for background time while downloads run. Playback
    /// keeps the app alive anyway, and anything cut off resumes on next launch.
    private func holdBackground(_ hold: Bool) {
        #if canImport(UIKit) && !SWIFT_PACKAGE
        if hold && bgTask == .invalid {
            bgTask = UIApplication.shared.beginBackgroundTask(withName: "downloads") { [weak self] in
                guard let self else { return }
                UIApplication.shared.endBackgroundTask(self.bgTask)
                self.bgTask = .invalid
            }
        } else if !hold && bgTask != .invalid {
            UIApplication.shared.endBackgroundTask(bgTask)
            bgTask = .invalid
        }
        #endif
    }
}