Amphlify/Packages/AmphlifyKit/Sources/Store/MemoryLibrary.swift

import Domain
import Foundation

public actor MemoryLibrary: Library {
    private var artistsByID: [EntityID: Artist] = [:]
    private var albumsByID: [EntityID: Album] = [:]
    private var tracksByID: [EntityID: Track] = [:]
    private var playlistsByID: [EntityID: Playlist] = [:]
    private var entriesByPlaylist: [EntityID: [PlaylistEntry]] = [:]
    private var stationsByID: [EntityID: RadioStation] = [:]
    private var annotations: [EntityID: Annotation] = [:]
    private var cursors: [SourceID: Cursor] = [:]
    private var observers: [UUID: AsyncStream<Void>.Continuation] = [:]

    public init() {}

    public func apply(_ changes: [Change]) {
        for change in changes {
            switch change {
            case let .upsert(.artist(artist)): artistsByID[artist.id] = artist
            case let .upsert(.album(album)): albumsByID[album.id] = album
            case let .upsert(.track(track)): tracksByID[track.id] = track
            case let .upsert(.playlist(playlist, entries)):
                playlistsByID[playlist.id] = playlist
                entriesByPlaylist[playlist.id] = entries
            case let .upsert(.station(station)): stationsByID[station.id] = station
            case let .remove(kind, id): remove(kind, id)
            case let .annotate(id, annotation): annotations[id] = annotation
            case .checkpoint: break
            }
        }
        notify()
    }

    private func remove(_ kind: EntityKind, _ id: EntityID) {
        switch kind {
        case .artist: artistsByID[id] = nil
        case .album: albumsByID[id] = nil
        case .track: tracksByID[id] = nil
        case .playlist:
            playlistsByID[id] = nil
            entriesByPlaylist[id] = nil
        case .station: stationsByID[id] = nil
        }
        annotations[id] = nil
    }

    public func cursor(for source: SourceID) -> Cursor? {
        cursors[source]
    }

    public func setCursor(_ cursor: Cursor, for source: SourceID) {
        cursors[source] = cursor
    }

    public func updates() -> AsyncStream<Void> {
        let key = UUID()
        let (stream, continuation) = AsyncStream<Void>.makeStream(bufferingPolicy: .bufferingNewest(1))
        observers[key] = continuation
        continuation.onTermination = { [weak self] _ in
            Task { await self?.removeObserver(key) }
        }
        return stream
    }

    private func removeObserver(_ key: UUID) {
        observers[key] = nil
    }

    private func notify() {
        observers.values.forEach { $0.yield() }
    }

    public func artists() -> [Artist] {
        artistsByID.values.sorted { $0.sortName.localizedStandardCompare($1.sortName) == .orderedAscending }
    }

    public func artist(_ id: EntityID) -> Artist? {
        artistsByID[id]
    }

    public func albums(sortedBy sort: AlbumSort) -> [Album] {
        let all = Array(albumsByID.values)
        switch sort {
        case .recentlyAdded: return all.sorted { $0.addedAt > $1.addedAt }
        case .title: return all.sorted { $0.title.localizedStandardCompare($1.title) == .orderedAscending }
        case .artist:
            return all.sorted {
                let order = $0.artistName.localizedStandardCompare($1.artistName)
                return order == .orderedSame ? ($0.year ?? 0) < ($1.year ?? 0) : order == .orderedAscending
            }
        case .year: return all.sorted { ($0.year ?? 0) > ($1.year ?? 0) }
        }
    }

    public func albums(byArtist id: EntityID) -> [Album] {
        albumsByID.values.filter { $0.artistID == id }.sorted { ($0.year ?? 0) > ($1.year ?? 0) }
    }

    public func albums(inGenre genre: String) -> [Album] {
        albums(sortedBy: .artist).filter { $0.genre == genre }
    }

    public func album(_ id: EntityID) -> Album? {
        albumsByID[id]
    }

    public func tracks() -> [Track] {
        tracksByID.values.sorted { $0.title.localizedStandardCompare($1.title) == .orderedAscending }
    }

    public func tracks(inAlbum id: EntityID) -> [Track] {
        tracksByID.values.filter { $0.albumID == id }.sorted { ($0.discNumber, $0.trackNumber) < ($1.discNumber, $1.trackNumber) }
    }

    public func tracks(byArtist id: EntityID) -> [Track] {
        tracksByID.values.filter { $0.artistID == id }
            .sorted { (annotations[$0.id]?.playCount ?? 0, $1.title) > (annotations[$1.id]?.playCount ?? 0, $0.title) }
    }

    public func tracks(inGenre genre: String) -> [Track] {
        tracksByID.values.filter { $0.genre == genre }
    }

    public func tracks(inPlaylist id: EntityID) -> [Track] {
        (entriesByPlaylist[id] ?? []).compactMap { tracksByID[$0.track] }
    }

    public func playlists() -> [Playlist] {
        playlistsByID.values.sorted { $0.updatedAt > $1.updatedAt }
    }

    public func stations() -> [RadioStation] {
        stationsByID.values.sorted { $0.name < $1.name }
    }

    public func genres() -> [Genre] {
        Dictionary(grouping: albumsByID.values.filter { $0.genre != nil }, by: { $0.genre ?? "" })
            .map { Genre(name: $0.key, albumCount: $0.value.count) }
            .sorted { $0.name < $1.name }
    }

    private func albumsRanked(limit: Int, by score: (Annotation) -> Double?) -> [Album] {
        var best: [EntityID: Double] = [:]
        for (id, annotation) in annotations {
            guard let value = score(annotation), let album = tracksByID[id]?.albumID else { continue }
            best[album] = max(best[album] ?? -.infinity, value)
        }
        return best.sorted { $0.value > $1.value }.prefix(limit).compactMap { albumsByID[$0.key] }
    }

    public func recentlyPlayed(limit: Int) -> [Album] {
        albumsRanked(limit: limit) { $0.lastPlayedAt?.timeIntervalSince1970 }
    }

    public func mostPlayed(limit: Int) -> [Album] {
        albumsRanked(limit: limit) { $0.playCount > 0 ? Double($0.playCount) : nil }
    }

    public func favoriteTracks() -> [Track] {
        tracksByID.values
            .filter { annotations[$0.id]?.favoritedAt != nil }
            .sorted { (annotations[$0.id]?.favoritedAt ?? .distantPast) > (annotations[$1.id]?.favoritedAt ?? .distantPast) }
    }

    public func favorites() -> Set<EntityID> {
        Set(annotations.filter { $0.value.favoritedAt != nil }.keys)
    }

    public func setFavorite(_ isFavorite: Bool, for id: EntityID) {
        annotations[id, default: Annotation()].favoritedAt = isFavorite ? Date() : nil
        notify()
    }

    public func recordPlay(of track: Track, at date: Date) {
        annotations[track.id, default: Annotation()].playCount += 1
        annotations[track.id, default: Annotation()].lastPlayedAt = date
        notify()
    }

    public func search(_ query: String) -> SearchResults {
        let needle = query.trimmingCharacters(in: .whitespaces)
        guard !needle.isEmpty else { return SearchResults() }
        func matches(_ text: String) -> Bool {
            text.range(of: needle, options: [.caseInsensitive, .diacriticInsensitive]) != nil
        }
        var results = SearchResults()
        results.artists = artists().filter { matches($0.name) }
        results.albums = albums(sortedBy: .title).filter { matches($0.title) }
        results.tracks = Array(tracks().filter { matches($0.title) }.prefix(50))
        results.playlists = playlists().filter { matches($0.name) }
        return results
    }
}