update-itunes-track-count.rb

require 'json'
require 'plist'
require 'tempfile'
require 'terminal-table'


class Hash
  def without(*keys)
    copy = self.dup
    keys.each { |key| copy.delete(key) }
    copy
  end

  def contains?(target)
    target.keys.all? do |key|
      self.include?(key) && self[key] == target[key]
    end
  end
end

class TrackCounts
  def initialize(xml_path)
    @xml = Plist.parse_xml(xml_path)
  end

  def tracks
    @tracks ||= begin
      tracks =  tracks_without_track_counts
      albums = albums_without_complete_track_counts(tracks)
      tracks_from_albums_without_complete_track_counts(albums)
    end
  end

  def tracks_without_track_counts
    @xml['Tracks'].select do |_id, track|
      track['Track Number'] && track['Track Count'].nil? && track['Kind'] !~ /video/
    end
  end

  def albums_without_complete_track_counts(tracks)
    albums = []
    tracks.each do |_id, track|
      album = {}
      album['Artist'] = track['Artist']
      album['Album'] = track['Album']
      album['Year'] = track['Year']
      album['Disc Number'] = track['Disc Number']
      albums << album
    end
    albums.uniq
  end

  def tracks_from_albums_without_complete_track_counts(albums)
    tracks = {}
    albums.each do |album|
      catch(:skip_album) do
        highest_track = 0
        album_tracks = {}
        @xml['Tracks'].each do |id, track|
          if track.contains?(album.without('Disc Number'))
            throw :skip_album if track['Track Number'].nil?
            highest_track = track['Track Number'] if track['Track Number'] > highest_track
            album_tracks[id] = track
          end
        end
        track_numbers = album_tracks.map{ |k, v| v['Track Number']}.sort
        throw :skip_album unless track_numbers[0] == 1 && track_numbers.each_cons(2).all? { |x, y| y == x + 1 }
        album_tracks.map { |k, v| v['Track Count'] = highest_track }
        tracks.merge!(album_tracks)
      end
    end
    tracks
  end

  def tracks_for_osa
    @tracks.map { |id, track| { id => track['Track Count'] } }
  end
end

class TerminalIO
  class << self
    def print_table(tracks)
      rows = []
      tracks.each do |_id, track|
        rows << [track['Name'], track['Artist'], track['Album'], track['Track Number'], "\e[38;5;21m#{track['Track Count']}\e[0m"]
      end
      Terminal::Table.new headings: ['Name', 'Artist', 'Album', 'Track Number', 'Track Count'], rows: rows
    end

    def confirm?(tracks)
      puts print_table(tracks)
      puts "Update the above track counts? (y/N)"
      if gets.chomp == 'y'
        return true
      end
    end
  end
end

class OSAScript
  class << self
    def script(tracks)
      <<~HEREDOC
        (function() {
          const tracks = #{tracks.to_json}
          const iTunes = Application('iTunes')

          tracks.forEach( track => {
            for (let key in track) {
              const iTunesTrack = iTunes.playlists.whose( { name: 'Library' } )[0].tracks.whose( { databaseID: { _equals: key } } )
              iTunesTrack.trackCount.set(track[key])
            }
          })
        })()
      HEREDOC
    end

    def update_tracks(tracks)
      tmp = Tempfile.new
      tmp << script(tracks)
      tmp.flush
      system("osascript -l JavaScript #{tmp.path}")
    end
  end
end

track_counts = TrackCounts.new("#{Dir.home}/Music/iTunes/iTunes Library.xml")

if TerminalIO.confirm?(track_counts.tracks)
  OSAScript.update_tracks(track_counts.tracks_for_osa)
end