stemify/ios/Stemify/Plex/PlexClient.swift

import Foundation
import UIKit

struct NotSignedIn: LocalizedError {
    var errorDescription: String? { "Not signed in to Plex" }
}

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

/// Plex Media Server and plex.tv calls. Library reads go through [http], whose cache
/// answers from disk when the server can't be reached (see AppGraph).
final class PlexClient: @unchecked Sendable {
    static let product = "Stemify"
    private static let page = 1000

    private let http: Http
    private let settings: Settings
    private let model: String
    private let osVersion: String

    init(http: Http, settings: Settings) {
        self.http = http
        self.settings = settings
        var sys = utsname()
        uname(&sys)
        let machine = withUnsafeBytes(of: &sys.machine) { String(decoding: $0.prefix { $0 != 0 }, as: UTF8.self) }
        model = machine.isEmpty ? "iPhone" : machine
        osVersion = ProcessInfo.processInfo.operatingSystemVersionString
            .replacingOccurrences(of: "Version ", with: "").components(separatedBy: " ").first ?? ""
    }

    var signedIn: Bool { !settings.plexToken.isEmpty }
    private var base: String { settings.plexUrl.trimmingCharacters(in: CharacterSet(charactersIn: "/")) }
    private var token: String { settings.plexToken }
    /// The chosen music libraries; listings pool them, in the order they were picked.
    private var sections: [String] { settings.musicSection.split(separator: ",").map(String.init) }

    private func request(_ url: URL, withToken: Bool = true, method: String = "GET") -> URLRequest {
        var r = URLRequest(url: url)
        r.httpMethod = method
        r.setValue("application/json", forHTTPHeaderField: "Accept")
        r.setValue(Self.product, forHTTPHeaderField: "X-Plex-Product")
        r.setValue("1.0", forHTTPHeaderField: "X-Plex-Version")
        r.setValue(settings.clientId, forHTTPHeaderField: "X-Plex-Client-Identifier")
        r.setValue("iOS", forHTTPHeaderField: "X-Plex-Platform")
        r.setValue(osVersion, forHTTPHeaderField: "X-Plex-Platform-Version")
        r.setValue(model, forHTTPHeaderField: "X-Plex-Device")
        r.setValue("\(Self.product) (\(model))", forHTTPHeaderField: "X-Plex-Device-Name")
        if withToken, !token.isEmpty { r.setValue(token, forHTTPHeaderField: "X-Plex-Token") }
        return r
    }

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

    private func decode(_ data: Data) throws -> MediaContainer {
        try JSONDecoder().decode(PlexResponse.self, from: data).container
    }

    private func get(_ path: String, _ params: [(String, Any?)] = [], start: Int? = nil, size: Int? = nil) async throws -> MediaContainer {
        guard signedIn else { throw NotSignedIn() }
        var req = request(try url(path, params))
        if let start { req.setValue(String(start), forHTTPHeaderField: "X-Plex-Container-Start") }
        if let size { req.setValue(String(size), forHTTPHeaderField: "X-Plex-Container-Size") }
        let body = try await http.data(req)
        if body.allSatisfy({ $0 == 0x20 || $0 == 0x0A || $0 == 0x0D || $0 == 0x09 }) { return MediaContainer() }
        return try decode(body)
    }

    @discardableResult
    private func send(_ method: String, _ path: String, _ params: [(String, Any?)] = []) async throws -> MediaContainer {
        guard signedIn else { throw NotSignedIn() }
        var req = request(try url(path, params), method: method)
        if method != "DELETE" { req.httpBody = Data() }
        let body = try await http.data(req)
        let text = String(decoding: body, as: UTF8.self).trimmingCharacters(in: .whitespacesAndNewlines)
        guard text.hasPrefix("{") else { return MediaContainer() }
        return (try? decode(body)) ?? MediaContainer()
    }

    /// Pages through a section listing; Plex caps a single response on big libraries.
    private func all(_ path: String, _ params: [(String, Any?)] = []) async throws -> [PlexItem] {
        var out: [PlexItem] = []
        var start = 0
        while true {
            let page = try await get(path, params, start: start, size: Self.page)
            out += page.metadata
            start += page.metadata.count
            let total = page.totalSize ?? page.size
            if page.metadata.isEmpty || start >= total { break }
        }
        return out
    }

    /// Every section's listing, one after another.
    private func pooled(_ params: [(String, Any?)]) async throws -> [PlexItem] {
        var out: [PlexItem] = []
        for sec in sections { out += try await all("/library/sections/\(sec)/all", params) }
        return out
    }

    /// A page of a listing that runs on through each section in turn.
    private func pooledPage(_ params: [(String, Any?)], start: Int, size: Int) async throws -> ([PlexItem], Int) {
        if sections.count == 1 {
            let c = try await get("/library/sections/\(sections[0])/all", params, start: start, size: size)
            return (c.metadata, c.totalSize ?? c.size)
        }
        var out: [PlexItem] = []
        var offset = 0, total = 0
        for sec in sections {
            let n = try await get("/library/sections/\(sec)/all", params, start: 0, size: 0)
            let count = n.totalSize ?? n.size
            let from = max(start + out.count - offset, 0)
            if out.count < size, from < count {
                out += try await get("/library/sections/\(sec)/all", params, start: from, size: size - out.count).metadata
            }
            offset += count; total += count
        }
        return (out, total)
    }

    // ---- server -------------------------------------------------------------------

    func machineId() async throws -> String {
        let body = try await http.data(request(try url("/identity"), withToken: false))
        guard let id = try decode(body).machineIdentifier else { throw PlainError("Not a Plex server") }
        return id
    }

    func musicSections() async throws -> [MusicSection] {
        try await get("/library/sections").directories
            .filter { $0.type == "artist" }
            .map { MusicSection(key: $0.key, title: $0.title, locations: $0.locations.map(\.path)) }
    }

    // ---- library ------------------------------------------------------------------

    func artists(sort: String = "titleSort") async throws -> [Artist] {
        try await pooled([("type", 8), ("sort", sort)]).map { $0.toArtist() }
    }

    func albums(sort: String = "titleSort") async throws -> [Album] {
        try await pooled([("type", 9), ("sort", sort)]).map { $0.toAlbum() }
    }

    func recentlyAdded(count: Int = 20) async throws -> [Album] {
        var out: [Album] = []
        for sec in sections {
            out += try await get("/library/sections/\(sec)/all", [("type", 9), ("sort", "addedAt:desc")], start: 0, size: count)
                .metadata.map { $0.toAlbum() }
        }
        return out
    }

    /// One page of every track. Returns the page and the library's track count.
    func tracksPage(start: Int, size: Int, sort: String = "titleSort") async throws -> ([Track], Int) {
        let (page, total) = try await pooledPage([("type", 10), ("sort", sort)], start: start, size: size)
        return (page.map { $0.toTrack() }, total)
    }

    /// Random tracks from the whole library, for shuffling it without loading every page.
    func randomTracks(_ count: Int) async throws -> [Track] {
        var out: [Track] = []
        for sec in sections {
            out += try await get("/library/sections/\(sec)/all", [("type", 10), ("sort", "random")], start: 0, size: count / max(sections.count, 1) + 1)
                .metadata.map { $0.toTrack() }
        }
        return Array(out.shuffled().prefix(count))
    }

    func album(_ key: String) async throws -> Album? {
        try await get("/library/metadata/\(key)").metadata.first?.toAlbum()
    }

    func albumTracks(_ key: String) async throws -> [Track] {
        try await get("/library/metadata/\(key)/children").metadata.map { $0.toTrack() }
    }

    // Not /library/metadata/<artist>/children: that leaves out compilations, singles and live
    // albums (an artist with only "The Best of Kansas" came back empty).
    func artistAlbums(_ key: String) async throws -> [Album] {
        try await pooled([("type", 9), ("artist.id", key)]).map { $0.toAlbum() }
    }

    func artistTracks(_ key: String) async throws -> [Track] {
        try await get("/library/metadata/\(key)/allLeaves").metadata.map { $0.toTrack() }
    }

    func track(_ key: String) async throws -> PlexItem? {
        try await get("/library/metadata/\(key)").metadata.first
    }

    /// Title search inside the music section; used to give stem-only entries Plex metadata.
    func searchTracks(_ title: String) async throws -> [Track] {
        var out: [Track] = []
        for sec in sections { out += try await get("/library/sections/\(sec)/search", [("type", 10), ("query", title)]).metadata.map { $0.toTrack() } }
        return out
    }

    /// Plex's own loudness analysis (what Plexamp draws its waveform from), in dB.
    func loudness(_ trackKey: String, samples: Int) async throws -> [Float]? {
        guard let streamId = try await track(trackKey)?.media.first?.parts.first?.streams.first(where: { $0.streamType == 2 })?.id
        else { return nil }
        let levels: [PlexLevel]
        do {
            levels = try await get("/library/streams/\(streamId)/levels", [("subsample", samples)]).levels
        } catch is HttpError {
            return nil // not analysed yet
        }
        return levels.isEmpty ? nil : levels.map(\.v)
    }

    func search(_ query: String) async throws -> SearchResults {
        let hubs = try await get("/hubs/search", [("query", query), ("limit", 12), ("sectionId", sections.count == 1 ? sections[0] : nil)]).hubs
        func hub(_ type: String) -> [PlexItem] { hubs.first { $0.type == type }?.metadata ?? [] }
        return SearchResults(
            artists: hub("artist").map { $0.toArtist() },
            albums: hub("album").map { $0.toAlbum() },
            tracks: hub("track").map { $0.toTrack() },
            playlists: hub("playlist").filter { $0.playlistType == "audio" }.map { $0.toPlaylist() }
        )
    }

    // ---- playlists ----------------------------------------------------------------

    func playlists(sort: String = "titleSort") async throws -> [Playlist] {
        try await get("/playlists", [("playlistType", "audio"), ("sort", sort)]).metadata.map { $0.toPlaylist() }
    }

    func playlistTracks(_ key: String) async throws -> [Track] {
        try await all("/playlists/\(key)/items").map { $0.toTrack() }
    }

    private func libraryUri(_ keys: [String]) async throws -> String {
        var machine = settings.plexMachineId
        if machine.isEmpty {
            machine = try await machineId()
            settings.plexMachineId = machine
        }
        return "server://\(machine)/com.plexapp.plugins.library/library/metadata/\(keys.joined(separator: ","))"
    }

    func createPlaylist(_ title: String, trackKeys: [String]) async throws -> Playlist? {
        let uri = try await libraryUri(trackKeys)
        return try await send("POST", "/playlists", [("type", "audio"), ("title", title), ("smart", 0), ("uri", uri)])
            .metadata.first?.toPlaylist()
    }

    func addToPlaylist(_ key: String, trackKeys: [String]) async throws {
        let uri = try await libraryUri(trackKeys)
        try await send("PUT", "/playlists/\(key)/items", [("uri", uri)])
    }

    func removeFromPlaylist(_ key: String, playlistItemId: Int64) async throws {
        try await send("DELETE", "/playlists/\(key)/items/\(playlistItemId)")
    }

    /// Moves an item to just after [afterItemId], or to the top when it's nil.
    func movePlaylistItem(_ key: String, playlistItemId: Int64, afterItemId: Int64?) async throws {
        try await send("PUT", "/playlists/\(key)/items/\(playlistItemId)/move", [("after", afterItemId)])
    }

    func renamePlaylist(_ key: String, title: String) async throws {
        try await send("PUT", "/playlists/\(key)", [("title", title)])
    }

    func deletePlaylist(_ key: String) async throws {
        try await send("DELETE", "/playlists/\(key)")
    }

    // ---- playback reporting ---------------------------------------------------------

    func scrobble(_ trackKey: String) async {
        _ = try? await get("/:/scrobble", [("key", trackKey), ("identifier", "com.plexapp.plugins.library")])
    }

    func timeline(_ track: Track, state: String, positionMs: Int64) async {
        _ = try? await get("/:/timeline", [
            ("ratingKey", track.ratingKey), ("key", "/library/metadata/\(track.ratingKey)"),
            ("state", state), ("time", positionMs), ("duration", track.durationMs), ("hasMDE", 1),
        ])
    }

    // ---- URLs -----------------------------------------------------------------------

    func artUrl(_ thumb: String?, px: Int) -> URL? {
        guard let thumb, !thumb.isEmpty, signedIn, let u = URL(string: base + "/photo/:/transcode") else { return nil }
        return u.with([("width", px), ("height", px), ("minSize", 1), ("upscale", 1), ("url", thumb), ("X-Plex-Token", token)])
    }

    func partUrl(_ partKey: String) -> URL? {
        URL(string: base + partKey + "?X-Plex-Token=" + urlEncode(token))
    }

    // ---- plex.tv sign-in ------------------------------------------------------------

    func createPin() async throws -> PlexPin {
        var req = request(URL(string: "https://plex.tv/api/v2/pins?strong=false")!, withToken: false, method: "POST")
        req.httpBody = Data()
        return try JSONDecoder().decode(PlexPin.self, from: try await http.data(req))
    }

    func checkPin(_ id: Int64) async throws -> PlexPin {
        let req = request(URL(string: "https://plex.tv/api/v2/pins/\(id)")!, withToken: false)
        return try JSONDecoder().decode(PlexPin.self, from: try await http.data(req))
    }

    func user(token: String) async throws -> PlexUser {
        var req = request(URL(string: "https://plex.tv/api/v2/user")!, withToken: false)
        req.setValue(token, forHTTPHeaderField: "X-Plex-Token")
        return try JSONDecoder().decode(PlexUser.self, from: try await http.data(req))
    }

    /// True if [url] answers as a Plex server with the given token.
    func probe(_ url: String, token: String) async -> Bool {
        let trimmed = url.trimmingCharacters(in: CharacterSet(charactersIn: "/"))
        guard let u = URL(string: trimmed + "/library/sections") else { return false }
        var req = request(u, withToken: false)
        req.setValue(token, forHTTPHeaderField: "X-Plex-Token")
        return await http.ok(req)
    }
}