stemify/ios/Stemify/Stems/StemServer.swift

import Foundation
import Observation

enum StemServerStatus: Equatable {
    case unknown
    case online(tracks: Int, transcode: Bool)
    case offline(reason: String, cachedTracks: Int)

    var isOnline: Bool { if case .online = self { true } else { false } }
}

private struct ServerInfo: Decodable {
    var tracks = 0
    var ffmpeg = false
    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: AnyKey.self)
        tracks = c.int("tracks") ?? 0
        ffmpeg = c.bool("ffmpeg") ?? false
    }
}

/// Client for server/stemserve.py: the stems index, stem file URLs and waveforms.
@Observable
final class StemServer: @unchecked Sendable {
    @ObservationIgnored private let http: Http
    @ObservationIgnored private let settings: Settings
    @ObservationIgnored private let indexFile: URL
    @ObservationIgnored private let etagFile: URL
    @ObservationIgnored private let refreshLock = AsyncLock()
    @ObservationIgnored private var transcode = false
    @ObservationIgnored private var indexListeners: [() -> Void] = []
    private static let checkSeconds: UInt64 = 60

    private(set) var index = StemIndex([])
    private(set) var status = StemServerStatus.unknown
    /// When the index was last checked against the server, nil before the first check.
    private(set) var checkedAt: Date?

    init(http: Http, settings: Settings, dir: URL) {
        self.http = http
        self.settings = settings
        indexFile = dir.appendingPathComponent("stem_index.json")
        etagFile = dir.appendingPathComponent("stem_index.etag")
        if let data = try? Data(contentsOf: indexFile),
           let parsed = try? JSONDecoder().decode(StemIndexFile.self, from: data) {
            index = StemIndex(parsed.tracks, roots: roots())
        }
        // stemd on the server separates new music around the clock; look for its new stems
        // every minute. Unchanged indexes cost one 304.
        Task.detached { [weak self] in
            while let self {
                await self.refresh()
                try? await Task.sleep(nanoseconds: Self.checkSeconds * 1_000_000_000)
            }
        }
    }

    private var base: String { settings.stemUrl.trimmingCharacters(in: CharacterSet(charactersIn: "/")) }
    private func roots() -> [String] { settings.stemRoots.split(separator: "\n").map(String.init).filter { !$0.isEmpty } }

    /// Main-thread callback after the index is replaced.
    func onIndexChange(_ f: @escaping () -> Void) { indexListeners.append(f) }

    private func req(_ url: URL, method: String = "GET") -> URLRequest {
        var r = URLRequest(url: url)
        r.httpMethod = method
        let t = settings.stemToken
        if !t.isEmpty { r.setValue(t, forHTTPHeaderField: "X-Stem-Token") }
        return r
    }

    private func url(_ path: String, _ params: [(String, Any?)] = []) throws -> URL {
        guard let u = URL(string: base + path) else { throw PlainError("Bad stemserve address") }
        return u.with(params)
    }

    func refresh(force: Bool = false) async {
        await refreshLock.run {
            do {
                let info = try JSONDecoder().decode(ServerInfo.self, from: try await http.data(req(try url("/api/info"))))
                transcode = info.ffmpeg
                let fm = FileManager.default
                let etag = (!force && fm.fileExists(atPath: indexFile.path)) ? try? String(contentsOf: etagFile, encoding: .utf8) : nil
                var r = req(try url("/api/index"))
                if let etag { r.setValue(etag, forHTTPHeaderField: "If-None-Match") }
                // Straight to the session: a 304 must not be confused with a cached 200.
                let (body, resp) = try await http.session.data(for: r)
                let code = (resp as? HTTPURLResponse)?.statusCode ?? 0
                var newIndex: StemIndex?
                if code == 304 {
                } else if (200..<300).contains(code) {
                    let parsed = try JSONDecoder().decode(StemIndexFile.self, from: body)
                    newIndex = StemIndex(parsed.tracks, roots: roots())
                    try? body.write(to: indexFile, options: .atomic)
                    if let tag = (resp as? HTTPURLResponse)?.value(forHTTPHeaderField: "ETag") {
                        try? tag.write(to: etagFile, atomically: true, encoding: .utf8)
                    }
                } else {
                    throw HttpError(code: code, message: "stemserve answered \(code)")
                }
                let fresh = newIndex
                await MainActor.run {
                    if let newIndex = fresh {
                        self.index = newIndex
                        self.indexListeners.forEach { $0() }
                    }
                    self.status = .online(tracks: self.index.size, transcode: info.ffmpeg)
                    self.checkedAt = Date()
                }
            } catch {
                let reason = error.localizedDescription
                await MainActor.run { self.status = .offline(reason: reason, cachedTracks: self.index.size) }
            }
        }
    }

    func find(_ track: Track) -> StemEntry? {
        let idx = index
        if let k = track.stemKey { return idx.entries.first { $0.key == k } }
        let hit = idx.find(track.partFile, durationMs: track.durationMs)
        let learned = idx.roots.joined(separator: "\n")
        if learned != settings.stemRoots { settings.stemRoots = learned }
        return hit
    }

    /// Format the server will actually send: Opus needs ffmpeg there.
    func effectiveQuality(_ q: StemQuality) -> StemQuality {
        q.fmt == "opus" && !transcode && status.isOnline ? .Lossless : q
    }

    func fileUrl(_ entry: StemEntry, _ stem: StemFile, _ quality: StemQuality) -> URL? {
        let path = (entry.key.split(separator: "/", omittingEmptySubsequences: false).map(String.init) + [stem.file])
            .map(urlEncode).joined(separator: "/")
        guard let u = URL(string: base + "/file/" + path) else { return nil }
        var params: [(String, Any?)] = []
        if quality.fmt == "opus" { params += [("fmt", "opus"), ("br", quality.kbps)] }
        if !settings.stemToken.isEmpty { params.append(("token", settings.stemToken)) }
        return u.with(params)
    }

    /// Asks the server to start the Opus transcodes now, so the stems are ready when we need them.
    func prepare(_ entry: StemEntry, _ quality: StemQuality) {
        guard quality.fmt == "opus" else { return }
        Task.detached { [self] in
            guard let u = try? url("/api/prepare", [("key", entry.key), ("fmt", "opus"), ("br", quality.kbps)]) else { return }
            var r = req(u, method: "POST")
            r.httpBody = Data()
            _ = try? await http.session.data(for: r)
        }
    }

    /// "Stem Now": asks stemd (through stemserve) to separate these songs before the rest of the
    /// library. Returns how many it queued: the ones with no stems yet.
    func stemNow(_ tracks: [Track]) async throws -> Int {
        var seen = Set<String>()
        let paths = tracks.filter { find($0) == nil }.compactMap(\.partFile).filter { seen.insert($0).inserted }
        if paths.isEmpty { return 0 }
        var r = req(try url("/api/stem"), method: "POST")
        r.setValue("application/json", forHTTPHeaderField: "Content-Type")
        r.httpBody = try JSONSerialization.data(withJSONObject: ["paths": paths])
        let (body, resp) = try await http.session.data(for: r)
        let code = (resp as? HTTPURLResponse)?.statusCode ?? 0
        if !(200..<300).contains(code) {
            let msg = (try? JSONSerialization.jsonObject(with: body) as? [String: Any])?["error"] as? String
            throw HttpError(code: code, message: code == 404 ? "stemserve needs updating for Stem Now" : msg ?? "stemserve answered \(code)")
        }
        return paths.count
    }

    /// Per-stem RMS envelopes (0..1 of full scale), [buckets] values each.
    func peaks(_ entry: StemEntry, buckets: Int) async -> [String: [Float]]? {
        guard let u = try? url("/api/peaks", [("key", entry.key), ("n", buckets)]),
              let body = try? await http.data(req(u)),
              let obj = try? JSONSerialization.jsonObject(with: body) as? [String: Any],
              let stems = obj["stems"] as? [String: [Any]] else { return nil }
        return stems.mapValues { arr in arr.map { Float(($0 as? NSNumber)?.intValue ?? 0) / 1000 } }
    }
}

/// A simple async mutex: one closure at a time, in arrival order.
actor AsyncLock {
    private var busy = false
    private var waiters: [CheckedContinuation<Void, Never>] = []

    func run<T>(_ body: () async throws -> T) async rethrows -> T {
        if busy {
            await withCheckedContinuation { waiters.append($0) }
        }
        busy = true
        defer {
            if waiters.isEmpty { busy = false } else { waiters.removeFirst().resume() }
        }
        return try await body()
    }
}