stemify/ios/Stemify/Cache/MediaCache.swift

import CryptoKit
import Foundation
import Observation

struct CacheMeta: Codable {
    var length: Int64 = -1
    var bits = ""
    var pinned = false
    var access: Int64 = 0
}

struct MediaError: LocalizedError {
    let message: String
    init(_ m: String) { message = m }
    var errorDescription: String? { message }
}

/// A growable bit set, persisted as base64 of its little-endian bytes.
struct Bits {
    private(set) var words: [UInt64] = []
    private(set) var count = 0

    init() {}
    init(base64: String) {
        guard let d = Data(base64Encoded: base64) else { return }
        words = [UInt64](repeating: 0, count: (d.count + 7) / 8)
        for (i, byte) in d.enumerated() { words[i / 8] |= UInt64(byte) << UInt64((i % 8) * 8) }
        count = words.reduce(0) { $0 + $1.nonzeroBitCount }
    }

    subscript(i: Int) -> Bool {
        let w = i >> 6
        return w < words.count && words[w] & (1 << UInt64(i & 63)) != 0
    }

    mutating func set(_ i: Int) {
        let w = i >> 6
        if w >= words.count { words += [UInt64](repeating: 0, count: w - words.count + 1) }
        let m: UInt64 = 1 << UInt64(i & 63)
        if words[w] & m == 0 { words[w] |= m; count += 1 }
    }

    mutating func clear(_ i: Int) {
        let w = i >> 6
        guard w < words.count else { return }
        let m: UInt64 = 1 << UInt64(i & 63)
        if words[w] & m != 0 { words[w] &= ~m; count -= 1 }
    }

    /// Clears every bit at or after [from].
    mutating func truncate(_ from: Int) {
        var i = from
        while i < words.count * 64 { clear(i); i += 1 }
    }

    func nextClear(_ from: Int) -> Int {
        var i = max(0, from)
        while self[i] { i += 1 }
        return i
    }

    var base64: String {
        var bytes = [UInt8]()
        bytes.reserveCapacity(words.count * 8)
        for w in words { for b in 0..<8 { bytes.append(UInt8((w >> UInt64(b * 8)) & 0xFF)) } }
        while bytes.last == 0 { bytes.removeLast() }
        return Data(bytes).base64EncodedString()
    }
}

/// Disk cache for stems and original tracks. Each resource is a sparse file filled in
/// BLOCK-sized pieces with HTTP range requests: reads block only on the bytes they need,
/// a background prefetcher fills the rest in order, and seeks (FLAC without a seek table
/// bisects the file) fetch out of order. Pinned resources are downloads and never evicted.
@Observable
final class MediaCache: @unchecked Sendable {
    static let block = 64 * 1024
    static let onDemandRun = 8
    static let prefetchRun = 32

    struct Usage: Equatable { var cachedBytes: Int64; var pinnedBytes: Int64 }
    private final class CatalogEntry { var bytes: Int64; var access: Int64; var pinned: Bool
        init(_ b: Int64, _ a: Int64, _ p: Bool) { bytes = b; access = a; pinned = p } }

    @ObservationIgnored private let dir: URL
    @ObservationIgnored let session: URLSession
    @ObservationIgnored private let limitBytes: () -> Int64
    @ObservationIgnored private let lock = NSRecursiveLock()
    @ObservationIgnored private var open: [String: CachedResource] = [:]
    @ObservationIgnored private var refs: [String: Int] = [:]
    @ObservationIgnored private var catalog: [String: CatalogEntry] = [:]
    @ObservationIgnored private let prefetchQueue: OperationQueue = {
        let q = OperationQueue()
        q.maxConcurrentOperationCount = 6
        q.qualityOfService = .utility
        return q
    }()
    @ObservationIgnored private let trimQueue = DispatchQueue(label: "media-trim", qos: .utility)
    @ObservationIgnored private var trimQueued = false

    private(set) var usage = Usage(cachedBytes: 0, pinnedBytes: 0)

    init(dir: URL, session: URLSession, limitBytes: @escaping () -> Int64) {
        self.dir = dir
        self.session = session
        self.limitBytes = limitBytes
        try? FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
        trimQueue.async { [self] in
            let files = (try? FileManager.default.contentsOfDirectory(atPath: dir.path)) ?? []
            for name in files where name.hasSuffix(".meta") {
                guard let m = readMeta(dir.appendingPathComponent(name)) else { continue }
                let id = String(name.dropLast(5))
                lock.lock()
                if catalog[id] == nil {
                    catalog[id] = CatalogEntry(Int64(Bits(base64: m.bits).count) * Int64(Self.block), m.access, m.pinned)
                }
                lock.unlock()
            }
            publishUsage()
        }
    }

    /// Opens (or reuses) a resource; pair every call with [release].
    func acquire(_ key: String, url: @escaping () -> URL?) -> CachedResource {
        let id = Self.idFor(key)
        lock.lock()
        defer { lock.unlock() }
        refs[id, default: 0] += 1
        if let r = open[id] { return r }
        let r = CachedResource(id: id, dataFile: dir.appendingPathComponent("\(id).data"),
                               metaFile: dir.appendingPathComponent("\(id).meta"), url: url, cache: self)
        open[id] = r
        return r
    }

    func release(_ res: CachedResource) {
        lock.lock()
        let n = (refs[res.id] ?? 1) - 1
        if n > 0 {
            refs[res.id] = n
            lock.unlock()
            return
        }
        refs[res.id] = nil
        open[res.id] = nil
        lock.unlock()
        res.close()
        scheduleTrim()
    }

    func isComplete(_ key: String) -> Bool {
        let id = Self.idFor(key)
        lock.lock()
        if let r = open[id] { lock.unlock(); return r.isComplete }
        lock.unlock()
        guard let m = readMeta(dir.appendingPathComponent("\(id).meta")) else { return false }
        return m.length >= 0 && Bits(base64: m.bits).count >= Self.blocksFor(m.length)
            && FileManager.default.fileExists(atPath: dir.appendingPathComponent("\(id).data").path)
    }

    func isPinned(_ key: String) -> Bool {
        lock.lock(); defer { lock.unlock() }
        return catalog[Self.idFor(key)]?.pinned == true
    }

    /// Unpins a download so normal eviction can reclaim it.
    func unpin(_ key: String) {
        let id = Self.idFor(key)
        lock.lock()
        let r = open[id]
        lock.unlock()
        if let r { r.pinned = false; return }
        let f = dir.appendingPathComponent("\(id).meta")
        guard var m = readMeta(f) else { return }
        m.pinned = false
        writeMeta(f, m)
        lock.lock()
        catalog[id]?.pinned = false
        lock.unlock()
        publishUsage()
        scheduleTrim()
    }

    /// Deletes every resource that isn't pinned or in use.
    func clearUnpinned() {
        trimQueue.async { self.evict(keepBytes: 0) }
    }

    func submitPrefetch(_ task: @escaping () -> Void) { prefetchQueue.addOperation(task) }

    func update(_ id: String, bytes: Int64, access: Int64, pinned: Bool) {
        lock.lock()
        let e = catalog[id] ?? CatalogEntry(0, 0, false)
        catalog[id] = e
        e.bytes = bytes; e.access = access; e.pinned = pinned
        lock.unlock()
        publishUsage()
    }

    private func publishUsage() {
        lock.lock()
        var all: Int64 = 0
        var pinned: Int64 = 0
        for e in catalog.values {
            all += e.bytes
            if e.pinned { pinned += e.bytes }
        }
        lock.unlock()
        let u = Usage(cachedBytes: all, pinnedBytes: pinned)
        DispatchQueue.main.async { if self.usage != u { self.usage = u } }
    }

    func scheduleTrim() {
        lock.lock()
        if trimQueued { lock.unlock(); return }
        trimQueued = true
        lock.unlock()
        trimQueue.async { [self] in
            lock.lock(); trimQueued = false; lock.unlock()
            evict(keepBytes: limitBytes())
        }
    }

    private func evict(keepBytes: Int64) {
        lock.lock()
        var total = catalog.values.reduce(Int64(0)) { $0 + $1.bytes }
        var victims: [String] = []
        for (id, e) in catalog.sorted(by: { $0.value.access < $1.value.access }) {
            if total <= keepBytes { break }
            if e.pinned || open[id] != nil { continue }
            victims.append(id)
            total -= e.bytes
        }
        victims.forEach { catalog[$0] = nil }
        lock.unlock()
        let fm = FileManager.default
        for id in victims {
            try? fm.removeItem(at: dir.appendingPathComponent("\(id).data"))
            try? fm.removeItem(at: dir.appendingPathComponent("\(id).meta"))
        }
        if !victims.isEmpty { publishUsage() }
    }

    func readMeta(_ f: URL) -> CacheMeta? {
        (try? Data(contentsOf: f)).flatMap { try? JSONDecoder().decode(CacheMeta.self, from: $0) }
    }

    func writeMeta(_ f: URL, _ m: CacheMeta) {
        if let d = try? JSONEncoder().encode(m) { try? d.write(to: f, options: .atomic) }
    }

    static func idFor(_ key: String) -> String {
        Insecure.SHA1.hash(data: Data(key.utf8)).map { String(format: "%02x", $0) }.joined()
    }

    static func blocksFor(_ length: Int64) -> Int { Int((length + Int64(block) - 1) / Int64(block)) }
}

/// One cached remote file. Thread-safe; readers block until their bytes have arrived.
final class CachedResource: @unchecked Sendable {
    let id: String
    private let dataFile: URL
    private let metaFile: URL
    private let url: () -> URL?
    private unowned let cache: MediaCache
    private let cond = NSCondition()
    private var bits: Bits
    private var claimed = Bits()
    private var handle: FileHandle?
    private let ioLock = NSLock()
    private var failure: Error?
    private var unsaved = 0
    private var prefetching = false
    private var lastRead = 0
    private var closed = false
    private(set) var length: Int64 = -1
    private var pinnedFlag = false
    private static let block = MediaCache.block

    var pinned: Bool {
        get { cond.lock(); defer { cond.unlock() }; return pinnedFlag }
        set { cond.lock(); pinnedFlag = newValue; cond.unlock(); saveMeta() }
    }

    init(id: String, dataFile: URL, metaFile: URL, url: @escaping () -> URL?, cache: MediaCache) {
        self.id = id
        self.dataFile = dataFile
        self.metaFile = metaFile
        self.url = url
        self.cache = cache
        if let m = cache.readMeta(metaFile), FileManager.default.fileExists(atPath: dataFile.path) {
            length = m.length
            bits = Bits(base64: m.bits)
            pinnedFlag = m.pinned
        } else {
            bits = Bits()
            try? FileManager.default.removeItem(at: dataFile)
        }
    }

    /// Caller holds cond.
    private var blockCount: Int { length < 0 ? Int.max : MediaCache.blocksFor(length) }

    var isComplete: Bool {
        cond.lock(); defer { cond.unlock() }
        return length >= 0 && bits.count >= MediaCache.blocksFor(length)
    }

    var cachedBytes: Int64 {
        cond.lock(); defer { cond.unlock() }
        return min(Int64(bits.count) * Int64(Self.block), length < 0 ? Int64.max : length)
    }

    private func file() throws -> FileHandle {
        ioLock.lock(); defer { ioLock.unlock() }
        if closed { throw MediaError("closed") }
        if let h = handle { return h }
        if !FileManager.default.fileExists(atPath: dataFile.path) {
            FileManager.default.createFile(atPath: dataFile.path, contents: nil)
        }
        let h = try FileHandle(forUpdating: dataFile)
        handle = h
        return h
    }

    private func pread(_ at: Int64, _ buf: UnsafeMutableRawPointer, _ n: Int) throws -> Int {
        let fd = try file().fileDescriptor
        let r = Darwin.pread(fd, buf, n, off_t(at))
        if r < 0 { throw MediaError("read failed") }
        return r
    }

    private func pwrite(_ at: Int64, _ buf: UnsafeRawPointer, _ n: Int) throws {
        let fd = try file().fileDescriptor
        var done = 0
        while done < n {
            let w = Darwin.pwrite(fd, buf + done, n - done, off_t(at) + off_t(done))
            if w <= 0 { throw MediaError("write failed") }
            done += w
        }
    }

    /// Total size, fetching the first run if it isn't known yet.
    func size() throws -> Int64 {
        cond.lock()
        let known = length >= 0
        cond.unlock()
        if !known { try awaitBlock(0) }
        return length
    }

    /// Reads up to [size] bytes at [pos] into [buf]; -1 at end of file. Blocks on the network if needed.
    func readAt(_ pos: Int64, _ buf: UnsafeMutableRawPointer, _ size: Int) throws -> Int {
        let total = try self.size()
        if pos >= total { return -1 }
        let n = Int(min(Int64(size), total - pos))
        if n <= 0 { return 0 }
        let first = Int(pos / Int64(Self.block))
        let last = Int((pos + Int64(n) - 1) / Int64(Self.block))
        for b in first...last { try awaitBlock(b) }
        cond.lock(); lastRead = last; cond.unlock()
        var got = 0
        while got < n {
            let r = try pread(pos + Int64(got), buf + got, n - got)
            if r <= 0 { break }
            got += r
        }
        return got
    }

    func readAt(_ pos: Int64, _ buf: inout [UInt8], _ off: Int, _ size: Int) throws -> Int {
        try buf.withUnsafeMutableBytes { try readAt(pos, $0.baseAddress! + off, size) }
    }

    /// Starts filling the rest of the file in the background, from wherever reading is.
    func prefetch() {
        cond.lock()
        if prefetching || closed { cond.unlock(); return }
        prefetching = true
        cond.unlock()
        cache.submitPrefetch { [self] in
            // A reader that needs these bytes will retry and surface the error.
            try? fillAll(fromBlock: { self.lastRead }, cancelled: { self.closed })
            cond.lock(); prefetching = false; cond.unlock()
        }
    }

    /// Downloads every missing block on the calling thread.
    func downloadAll(progress: (Int64, Int64) -> Void, cancelled: () -> Bool) throws {
        try fillAll(fromBlock: { 0 }, cancelled: cancelled, progress: progress)
    }

    private func fillAll(fromBlock: () -> Int, cancelled: () -> Bool, progress: (Int64, Int64) -> Void = { _, _ in }) throws {
        _ = try size()
        while !cancelled() {
            var run: ClosedRange<Int>?
            var waitFor = -1
            cond.lock()
            let count = blockCount
            if bits.count >= count { cond.unlock(); return }
            let start = min(max(fromBlock(), 0), count - 1)
            var b = bits.nextClear(start)
            if b >= count { b = bits.nextClear(0) }
            if b >= count { cond.unlock(); return }
            if claimed[b] {
                // Someone else is fetching here; look for a free block after it.
                var c = b
                while c < count && (bits[c] || claimed[c]) { c += 1 }
                if c >= count { waitFor = b } else { run = claimRun(c, MediaCache.prefetchRun) }
            } else {
                run = claimRun(b, MediaCache.prefetchRun)
            }
            cond.unlock()
            if let run { try fetch(run) }
            if waitFor >= 0 { try awaitBlock(waitFor) }
            progress(cachedBytes, length)
        }
    }

    private func awaitBlock(_ b: Int) throws {
        var waited: TimeInterval = 0
        while true {
            var run: ClosedRange<Int>?
            cond.lock()
            if closed { cond.unlock(); throw MediaError("closed") }
            if length >= 0 && b >= blockCount { cond.unlock(); return }
            if bits[b] { cond.unlock(); return }
            if !claimed[b] {
                run = claimRun(b, MediaCache.onDemandRun)
            } else {
                let before = Date()
                _ = cond.wait(until: Date().addingTimeInterval(1))
                waited += Date().timeIntervalSince(before)
                if waited > 60 {
                    let f = failure
                    cond.unlock()
                    throw f ?? MediaError("timed out waiting for data")
                }
            }
            cond.unlock()
            if let run { try fetch(run) }
        }
    }

    /// Claims up to [max] consecutive missing, unclaimed blocks from [b]. Caller holds the lock.
    private func claimRun(_ b: Int, _ max: Int) -> ClosedRange<Int> {
        var end = b
        let count = blockCount
        while end + 1 < count && end + 1 - b < max && !bits[end + 1] && !claimed[end + 1] { end += 1 }
        for i in b...end { claimed.set(i) }
        return b...end
    }

    /// Synchronous range GET; the body of one run is at most a couple of MB.
    private func get(_ req: URLRequest) throws -> (Data, HTTPURLResponse) {
        let sem = DispatchSemaphore(value: 0)
        var out: (Data?, URLResponse?, Error?)
        let task = cache.session.dataTask(with: req) { d, r, e in out = (d, r, e); sem.signal() }
        task.resume()
        sem.wait()
        if let e = out.2 { throw e }
        guard let resp = out.1 as? HTTPURLResponse else { throw MediaError("no response") }
        return (out.0 ?? Data(), resp)
    }

    private func fetch(_ run: ClosedRange<Int>) throws {
        var block = run.lowerBound
        let B = Int64(Self.block)
        defer {
            cond.lock()
            // Release whatever this run didn't deliver, including blocks past EOF.
            if block <= run.upperBound { for b in block...run.upperBound { claimed.clear(b) } }
            cond.broadcast()
            cond.unlock()
        }
        do {
            guard let u = url() else { throw MediaError("no URL") }
            let start = Int64(run.lowerBound) * B
            cond.lock()
            let len = length
            cond.unlock()
            let endExclusive = len >= 0 ? min(Int64(run.upperBound + 1) * B, len) : Int64(run.upperBound + 1) * B
            var req = URLRequest(url: u)
            req.setValue("bytes=\(start)-\(endExclusive - 1)", forHTTPHeaderField: "Range")
            let (body, resp) = try get(req)
            var skip: Int64 = 0
            var lastBlock = run.upperBound
            switch resp.statusCode {
            case 206:
                if let cr = resp.value(forHTTPHeaderField: "Content-Range"),
                   let slash = cr.lastIndex(of: "/"), let total = Int64(cr[cr.index(after: slash)...]) {
                    learnLength(total)
                }
            case 200:
                // No range support: the body is the whole file, so keep all of it.
                learnLength(Int64(body.count))
                skip = 0
                block = 0
                lastBlock = MediaCache.blocksFor(Int64(body.count)) - 1
            default:
                throw MediaError("HTTP \(resp.statusCode) fetching media")
            }
            cond.lock()
            let length = self.length
            let count = blockCount
            cond.unlock()
            if length < 0 { throw MediaError("server sent no length") }
            let bodyOffset = resp.statusCode == 200 ? Int64(0) : start
            lastBlock = min(lastBlock, count - 1)
            try body.withUnsafeBytes { (raw: UnsafeRawBufferPointer) in
                while block <= lastBlock {
                    let blockStart = Int64(block) * B
                    let want = Int(min(blockStart + B, length) - blockStart)
                    let from = Int(blockStart - bodyOffset + skip)
                    if from < 0 || from + want > raw.count { throw MediaError("connection ended early") }
                    cond.lock()
                    let have = bits[block]
                    cond.unlock()
                    if !have { try pwrite(blockStart, raw.baseAddress! + from, want) }
                    cond.lock()
                    bits.set(block)
                    claimed.clear(block)
                    failure = nil
                    cond.broadcast()
                    unsaved += 1
                    let save = unsaved >= 32
                    cond.unlock()
                    block += 1
                    if save || (block > lastBlock && isComplete) { saveMeta() }
                }
            }
            if resp.statusCode == 200 { block = run.upperBound + 1 }
        } catch {
            cond.lock()
            failure = error
            cond.broadcast()
            cond.unlock()
            throw error
        }
    }

    private func learnLength(_ total: Int64) {
        cond.lock()
        if length != total {
            length = total
            bits.truncate(MediaCache.blocksFor(total))
        }
        cond.unlock()
    }

    private func saveMeta() {
        cond.lock()
        unsaved = 0
        let m = CacheMeta(length: length, bits: bits.base64, pinned: pinnedFlag, access: Int64(Date().timeIntervalSince1970 * 1000))
        let bytes = min(Int64(bits.count) * Int64(Self.block), length < 0 ? 0 : length)
        cond.unlock()
        cache.writeMeta(metaFile, m)
        cache.update(id, bytes: bytes, access: m.access, pinned: m.pinned)
    }

    func close() {
        cond.lock()
        closed = true
        cond.broadcast()
        cond.unlock()
        saveMeta()
        ioLock.lock()
        try? handle?.close()
        handle = nil
        ioLock.unlock()
    }
}