stemify/ios/Stemify/Stems/StemIndex.swift

import Foundation

struct StemFile: Codable, Hashable {
    var name: String
    var file: String
    var size: Int64
    var mtime: Int64 = 0

    init(name: String, file: String, size: Int64, mtime: Int64 = 0) {
        self.name = name; self.file = file; self.size = size; self.mtime = mtime
    }

    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: AnyKey.self)
        name = c.str("name") ?? ""
        file = c.str("file") ?? ""
        size = c.int64("size") ?? 0
        mtime = c.int64("mtime") ?? 0
    }
}

struct StemEntry: Codable, Hashable {
    var key: String
    var source: String
    var duration: Double?
    var samplerate: Int?
    var model: String?
    var finished: Double?
    var stems: [StemFile] = []

    init(key: String, source: String, duration: Double? = nil, samplerate: Int? = nil, model: String? = nil,
         finished: Double? = nil, stems: [StemFile] = []) {
        self.key = key; self.source = source; self.duration = duration; self.samplerate = samplerate
        self.model = model; self.finished = finished; self.stems = stems
    }

    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: AnyKey.self)
        key = c.str("key") ?? ""
        source = c.str("source") ?? ""
        duration = c.double("duration")
        samplerate = c.int("samplerate")
        model = c.str("model")
        finished = c.double("finished")
        stems = c.list("stems")
    }

    /// "Artist/Album/01 Song" -> "01 Song" etc., for stem-only browsing.
    var title: String { key.split(separator: "/", omittingEmptySubsequences: false).last.map(String.init) ?? key }
    var folder: String {
        guard let i = key.lastIndex(of: "/") else { return "" }
        return String(key[..<i])
    }
}

struct StemIndexFile: Codable {
    var version = 1
    var tracks: [StemEntry] = []

    init(version: Int = 1, tracks: [StemEntry] = []) { self.version = version; self.tracks = tracks }
    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: AnyKey.self)
        version = c.int("version") ?? 1
        tracks = c.list("tracks")
    }
}

/// Maps Plex file paths to stem folders. The stems folder mirrors the music folder, so the
/// source path relative to the music root identifies a track; Plex reports absolute paths
/// as its container sees them, so the root is unknown up front.
///
/// Matching tries every suffix of the Plex path, longest first. A suffix that is only a
/// file name, or whose duration disagrees, is accepted only under a root prefix learned
/// from an earlier, longer match, so "Intro.flac" in one album can't grab another's stems.
final class StemIndex: @unchecked Sendable {
    let entries: [StemEntry]
    private var bySource: [String: StemEntry] = [:]
    private var learned: [String] = []
    private let lock = NSLock()

    init(_ entries: [StemEntry], roots: [String] = []) {
        self.entries = entries
        bySource.reserveCapacity(entries.count * 2)
        for e in entries {
            let k = Self.norm(e.source)
            if bySource[k] == nil { bySource[k] = e }
        }
        for r in roots where !learned.contains(r) { learned.append(r) }
    }

    var size: Int { entries.count }
    var roots: [String] { lock.lock(); defer { lock.unlock() }; return learned }

    /// Returns the entry for a Plex part path, learning the path's root on a confident match.
    func find(_ plexPath: String?, durationMs: Int64 = 0) -> StemEntry? {
        guard let plexPath, !plexPath.isEmpty, !entries.isEmpty else { return nil }
        lock.lock()
        defer { lock.unlock() }
        let path = Self.norm(plexPath)
        for root in learned where path.hasPrefix(root) {
            if let hit = bySource[String(path.dropFirst(root.count))] { return hit }
        }
        let parts = path.split(separator: "/", omittingEmptySubsequences: false).map(String.init)
        guard parts.count > 1 else { return nil }
        for i in 1..<parts.count {
            let rel = parts[i...].joined(separator: "/")
            guard let hit = bySource[rel] else { continue }
            let deep = parts.count - i >= 2
            if deep && durationMatches(hit, durationMs) {
                let root = parts[0..<i].joined(separator: "/") + "/"
                if !learned.contains(root) { learned.append(root) }
                return hit
            }
        }
        return nil
    }

    private func durationMatches(_ e: StemEntry, _ durationMs: Int64) -> Bool {
        guard let d = e.duration, durationMs > 0 else { return true }
        return abs(d * 1000 - Double(durationMs)) < 3000
    }

    /// NFC, forward slashes, no extension, lower case.
    static func norm(_ path: String) -> String {
        var p = path.precomposedStringWithCanonicalMapping.replacingOccurrences(of: "\\", with: "/")
        let slash = p.lastIndex(of: "/").map { p.distance(from: p.startIndex, to: $0) } ?? -1
        if let dotIdx = p.lastIndex(of: ".") {
            let dot = p.distance(from: p.startIndex, to: dotIdx)
            if dot > slash + 1 { p = String(p[..<dotIdx]) }
        }
        return p.lowercased()
    }
}