#!/usr/bin/env ruby
#
# omptagger [version 0.6.1]
# http://dev.gentoo.org/~omp/omptagger/
#
# Copyright 2007 David Shakaryan <omp@gentoo.org>
# Distributed under the terms of the GNU General Public License v3.
# See http://www.gnu.org/licenses/gpl.txt for the full license text.
#
# Program for modifying and displaying tags for various formats of audio files.
# Behaviour changes between acting as a wrapper tool or using a library, based
# on the format of the file being worked with.
#
# Dependencies:
# * ruby [ http://www.ruby-lang.org/ ]
# * flac [ http://flac.sourceforge.net/ ]
# * id3lib-ruby [ http://id3lib-ruby.rubyforge.org/ ]
# * vorbis-tools [ http://www.vorbis.com/ ]

require 'getoptlong'
require 'id3lib'

# List of valid options.
getopt = GetoptLong.new(
  ['--view',       '-v', GetoptLong::NO_ARGUMENT],
  ['--view-tag',   '-t', GetoptLong::REQUIRED_ARGUMENT],
  ['--set-tag',    '-s', GetoptLong::REQUIRED_ARGUMENT],
  ['--generate',   '-g', GetoptLong::NO_ARGUMENT],
  ['--preview',    '-p', GetoptLong::NO_ARGUMENT],
  ['--remove',     '-r', GetoptLong::NO_ARGUMENT],
  ['--remove-tag', '-d', GetoptLong::REQUIRED_ARGUMENT],
  ['--rename',     '-m', GetoptLong::NO_ARGUMENT],
  ['--scheme',     '-n', GetoptLong::REQUIRED_ARGUMENT],
  ['--no-colour',  '-c', GetoptLong::NO_ARGUMENT],
  ['--help',       '-h', GetoptLong::NO_ARGUMENT]
)

# Store any options set in a hash. Create a tracker variable to determine
# whether an action has been set.
action = false
$opts = Hash['scheme' => '%a - %t']
getopt.each do |opt, arg|
  opt.sub!('--', '')

  case opt
  # Options which may be used only once.
  when 'view', 'generate', 'preview', 'remove', 'rename', 'help'
    $opts[opt] = arg
    action = true
  when 'no-colour', 'scheme'
    $opts[opt] = arg

  # Options which may be used more than once.
  when 'view-tag', 'set-tag', 'remove-tag'
    if $opts.has_key?(opt)
      $opts[opt].push(arg)
    else
      $opts[opt] = [arg]
    end
    action = true
  end
end

# Define colours unless the no-colour option is set.
if $opts.has_key?('no-colour')
  def colred(str); return str; end
  def colgrn(str); return str; end
  def colyel(str); return str; end
  def colcyn(str); return str; end
else
  def colred(str); return "\e[31m#{str}\e[0m"; end
  def colgrn(str); return "\e[32m#{str}\e[0m"; end
  def colyel(str); return "\e[33m#{str}\e[0m"; end
  def colcyn(str); return "\e[36m#{str}\e[0m"; end
end

# Methods for outputting warning and status messages.
def warn(message)
  puts '  ' + colred('ERROR:') + ' ' + message
end
def status(message)
  puts '  ' + colgrn(message)
end

# Method for escaping single quotes for use within single quotes in shell
# commands. Substitutes all instances of ' with '\''.
def esc(str)
  return str.gsub("'", "'\\\\''")
end

# Method for dynamic spacing; used to correctly align tag and value output.
def spacing(tags, tag)
  return ' ' * (tags.map do |str| str.size end.max + 2 - tag.length)
end

# Method for splitting filename based on a naming scheme.
def namingscheme(scheme, name)
  regexp = /\A#{Regexp.escape(scheme).gsub(/%[a-z]/, '(.+?)')}\Z/
  tags = scheme.scan(regexp).flatten
  values = name.scan(regexp).flatten

  unless tags.length == values.length
    raise 'Filename does not match naming scheme.'
  end

  result = Hash.new
  tags.zip(values).each do |tag, value|
    result[tag.delete('%')] = value
  end

  return result
end

# Method for outputting program help.
def help
  puts colyel('Usage:') + ' omptagger ' + colgrn('[options]') + ' ' +
    colgrn('[files]')
  puts
  puts colyel('Options:')
  puts '  ' + colgrn('--view') + '        ' + colgrn('-v') +
    '  Display all tags'
  puts '  ' + colgrn('--view-tag') + '    ' + colgrn('-t') +
    '  Display a tag ' + colcyn('[required argument: tag]')
  puts '  ' + colgrn('--set-tag') + '     ' + colgrn('-s') +
    '  Set a tag ' + colcyn('[required argument: tag=string]')
  puts '  ' + colgrn('--generate') + '    ' + colgrn('-g') +
    '  Generate tags based on filename'
  puts '  ' + colgrn('--preview') + '     ' + colgrn('-p') +
    '  Preview tags that will be set with generate option'
  puts '  ' + colgrn('--remove') + '      ' + colgrn('-r') +
    '  Remove all tags'
  puts '  ' + colgrn('--remove-tag') + '  ' + colgrn('-d') +
    '  Remove a tag ' + colcyn('[required argument: tag]')
  puts '  ' + colgrn('--rename') + '      ' + colgrn('-m') +
    '  Rename files based on tags'
  puts '  ' + colgrn('--scheme') + '      ' + colgrn('-n') +
    '  Allow a file naming scheme to be specified'
  puts '  ' + colgrn('--no-colour') + '   ' + colgrn('-c') +
    '  Disable use of colour in program output'
  puts '  ' + colgrn('--help') + '        ' + colgrn('-h') +
    '  Display program help'
  puts
  puts colyel('Notes:')
  puts '  ' + colgrn('*') + ' Default file naming scheme is ' +
    colcyn('Artist - Title') + '.'
  puts '  ' + colgrn('*') + ' Underscores in filenames are converted to ' +
    'spaces in tags.'
  puts '  ' + colgrn('*') + ' Scheme keys: ' + colcyn('%a') + ' - ARTIST, ' +
    colcyn('%b') + ' - ALBUM, ' + colcyn('%d') + ' - DATE/YEAR,'
  puts '  ' + '  ' + colcyn('%n') + ' - TRACKNUMBER/TRACK, ' +
    colcyn('%t') + ' - TITLE.'
end

# Class for Vorbis comments; used by FLAC and Vorbis files.
class VorbisComments
  attr_reader :hash, :file

  # Possible tags.
  TAGS = ['TITLE', 'VERSION', 'ALBUM', 'TRACKNUMBER', 'ARTIST', 'PERFORMER',
          'COPYRIGHT', 'LICENSE', 'ORGANIZATION', 'DESCRIPTION', 'GENRE',
          'DATE', 'LOCATION', 'CONTACT', 'ISRC']

  # Map scheme keys to tags.
  KEYS = Hash['a' => 'ARTIST',
              'b' => 'ALBUM',
              'd' => 'DATE',
              'n' => 'TRACKNUMBER',
              't' => 'TITLE']

  # Method for outputting a tag and its corresponding value.
  def output(tag, value)
    puts '  ' + colcyn(tag) + spacing(TAGS, tag) + value
  end

  # Method for initialising a new object.
  def initialize(file)
    @file = file
    @tags = Hash.new

    TAGS.each do |tag|
      value = read_tag(tag)
      next if value.empty?

      @tags[tag] = value.split("\n").map do |str| str.sub(tag + '=', '') end
    end
  end

  # Method for writing tags.
  def write_tags
    clear_tags

    @tags.each do |tag, value|
      value.each do |value|
        write_tag(tag, value)
      end
    end
  end

  # Method for displaying all tags.
  def view
    raise 'No tags are set.' if @tags.empty?

    @tags.each do |tag, value|
      value.each do |value|
        output(tag, value)
      end
    end
  end

  # Method for displaying certain tags.
  def view_tag(arg)
    arg.each do |arg|
      # Process exceptions within the argument loop, as using the exception in
      # the file loop may end the argument loop prematurely.
      begin
        tag = arg.upcase
        raise tag + ' is not a valid tag.' unless TAGS.include?(tag)

        value = @tags[tag]
        raise tag + ' tag is not set.' if value.nil?

        value.each do |value|
          output(tag, value)
        end
      rescue RuntimeError => message
        warn(message)
      end
    end
  end

  # Method for setting tags.
  def set_tag(arg)
    status('Setting tags...')

    arg.each do |arg|
      # Process exceptions within the argument loop, as using the exception in
      # the file loop may end the argument loop prematurely.
      begin
        arg = arg.split('=', 2)
        tag = arg.first.upcase
        raise tag + ' is not a valid tag.' unless TAGS.include?(tag)

        if arg.length == 1
          value = ''
        else
          value = arg.last
        end

        @tags[tag] = [value]
        view_tag(tag)
      rescue RuntimeError => message
        warn(message)
      end
    end
  end

  # Method for generating new tags based on filename.
  def generate
    # Use only the basename of the file.
    value = File.basename(@file)
    # Substitute all underscores with a space and remove the file extension.
    value = value.gsub('_', ' ').sub(/\.(flac|ogg)$/, '')

    scheme = namingscheme($opts['scheme'].gsub('_', ' '), value)

    tag = Array.new
    value = Array.new

    scheme.each do |stag, svalue|
      unless KEYS.has_key?(stag)
        raise 'Naming scheme contains an invalid key.'
      end

      tag = tag.push(KEYS[stag])
      value = value.push(svalue)
    end

    status('Generating tags...')

    tag.zip(value).each do |arr|
      @tags[arr.first] = [arr.last]
      view_tag(arr.first)
    end
  end

  # Method for removing all tags.
  def remove
    raise 'No tags are set.' if @tags.empty?
    status('Removing tags...')

    @tags.clear
  end

  # Method for removing certain tags.
  def remove_tag(arg)
    status('Removing tags...')

    arg.each do |arg|
      # Process exceptions within the argument loop, as using the exception in
      # the file loop may end the argument loop prematurely.
      begin
        tag = arg.upcase
        raise tag + ' is not a valid tag.' unless TAGS.include?(tag)
        raise tag + ' tag is not set.' unless @tags.include?(tag)

        @tags.delete(tag)
      rescue RuntimeError => message
        warn(message)
      end
    end
  end

  # Method for renaming files based on tags.
  def rename
    unless @tags.has_key?('ARTIST')
      raise 'Missing ARTIST tag.'
    end
    unless @tags.has_key?('TITLE')
      raise 'Missing TITLE tag.'
    end

    # Determine what the new filename should be.
    ext = file.split('.').last.downcase
    file = File.dirname(@file) + '/'
    if @tags.has_key?('TRACKNUMBER')
      file = file + @tags['TRACKNUMBER'] + ' - '
    end
    file = file + @tags['ARTIST'] + ' - ' +
      @tags['TITLE'] + '.' + ext

    # Raise an error if the new and old filenames are identical.
    if File.basename(@file) == File.basename(file)
      raise 'Generated filename and current filename are identical.'
    end

    status('Renaming to ' + file.sub(/^\.\//, '') + '...')
    File.rename(@file, file)
  end
end

# Subclass of VorbisComments containing methods unique to FLAC files.
class FLAC < VorbisComments
  def read_tag(tag)
    %x(metaflac --show-tag=#{tag} -- '#{esc(@file)}').chomp
  end

  def clear_tags
    %x{metaflac --remove-all-tags -- '#{esc(@file)}'}
  end

  def write_tag(tag, value)
    %x{metaflac --set-tag=#{tag}='#{esc(value)}' -- '#{esc(@file)}'}
  end
end

# Subclass of VorbisComments containing methods unique to Vorbis files.
class Vorbis < VorbisComments
  def read_tag(tag)
    %x(vorbiscomment -l -- '#{esc(@file)}' | grep '#{tag}').chomp
  end

  def clear_tags
    %x{vorbiscomment -w -t '' -- '#{esc(@file)}' 2>/dev/null}
  end

  def write_tag(tag, value)
    %x{vorbiscomment -a -t #{tag}='#{esc(value)}' -- '#{esc(@file)}'}
  end
end

# Class for ID3 tags; used by MP3 files.
class ID3
  attr_reader :file, :tags

  # Possible tags.
  TAGS = Hash['TITLE'   => :TIT2,
              'ALBUM'   => :TALB,
              'TRACK'   => :TRCK,
              'ARTIST'  => :TPE1,
              'YEAR'    => :TYER,
              'COMMENT' => :COMM]

  # Map scheme keys to tags.
  KEYS = Hash['a' => 'ARTIST',
              'b' => 'ALBUM',
              'd' => 'YEAR',
              'n' => 'TRACK',
              't' => 'TITLE']

  # Method for outputting a tag and its corresponding value.
  def output(tag, value)
    puts '  ' + colcyn(tag) + spacing(TAGS.keys, tag) + colcyn(value)
  end

  # Method for initialising a new object.
  def initialize(file)
    @file = file
    @tags = ID3Lib::Tag.new(@file)
  end

  # Method for writing tags.
  def write_tags
    if @tags.empty? or (@tags.length == 1 and @tags.frame_text(:TLEN))
      @tags.strip!
    else
      @tags.update!
    end
  end

  # Method for displaying all tags.
  def view
    raise 'No tags are set.' if @tags.empty?

    TAGS.each do |tag, frame|
      value = @tags.frame_text(frame)
      output(tag, value) unless value.nil?
    end
  end

  # Method for displaying certain tags.
  def view_tag(arg)
    arg.each do |arg|
      # Process exceptions within the argument loop, as using the exception in
      # the file loop may end the argument loop prematurely.
      begin
        tag = arg.upcase
        raise tag + ' is not a valid tag.' unless TAGS.has_key?(tag)

        value = @tags.frame_text(TAGS[tag])
        raise tag + ' tag is not set.' if value.nil?

        output(tag, value)
      rescue RuntimeError => message
        warn(message)
      end
    end
  end

  # Method for setting tags.
  def set_tag(arg)
    status('Setting tags...')

    arg.each do |arg|
      # Process exceptions within the argument loop, as using the exception in
      # the file loop may end the argument loop prematurely.
      begin
        arg = arg.split('=', 2)
        tag = arg.first.upcase
        raise tag + ' is not a valid tag.' unless TAGS.has_key?(tag)

        if arg.length == 1
          value = ''
        else
          value = arg.last
        end

        @tags.set_frame_text(TAGS[tag], value)
        view_tag(tag)
      rescue RuntimeError => message
        warn(message)
      end
    end
  end

  # Method for generating new tags based on filename.
  def generate
    # Use only the basename of the file.
    value = File.basename(@file)
    # Substitute all underscores with a space and remove the file extension.
    value = value.gsub('_', ' ').sub(/\.mp3$/, '')

    scheme = namingscheme($opts['scheme'].gsub('_', ' '), value)

    tag = Array.new
    value = Array.new

    scheme.each do |stag, svalue|
      unless KEYS.has_key?(stag)
        raise 'Naming scheme contains an invalid key.'
      end

      tag = tag.push(KEYS[stag])
      value = value.push(svalue)
    end

    status('Generating tags...')

    tag.zip(value).each do |arr|
      @tags.set_frame_text(TAGS[arr.first], arr.last)
      view_tag(arr.first)
    end
  end

  # Method for removing all tags.
  def remove
    raise 'No tags are set.' if @tags.empty?
    status('Removing tags...')

    @tags.clear
  end

  # Method for removing certain tags.
  def remove_tag(arg)
    status('Removing tags...')

    arg.each do |arg|
      # Process exceptions within the argument loop, as using the exception in
      # the file loop may end the argument loop prematurely.
      begin
        tag = arg.upcase
        raise tag + ' is not a valid tag.' unless TAGS.has_key?(tag)
        raise tag + ' tag is not set.' if @tags.frame_text(TAGS[tag]).nil?

        @tags.remove_frame(TAGS[tag])
      rescue RuntimeError => message
        warn(message)
      end
    end
  end

  # Method for renaming files based on tags.
  def rename
    if @tags.frame_text(TAGS['ARTIST']).nil?
      raise 'Missing ARTIST tag.'
    end
    if @tags.frame_text(TAGS['TITLE']).nil?
      raise 'Missing TITLE tag.'
    end

    # Determine what the new filename should be.
    ext = file.split('.').last.downcase
    file = File.dirname(@file) + '/'
    unless @tags.frame_text(TAGS['TRACK']).nil?
      file = file + @tags.frame_text(TAGS['TRACK']) + ' - '
    end
    file = file + @tags.frame_text(TAGS['ARTIST']) + ' - ' +
      @tags.frame_text(TAGS['TITLE']) + '.' + ext

    # Raise an error if the new and old filenames are identical.
    if File.basename(@file) == File.basename(file)
      raise 'Generated filename and current filename are identical.'
    end

    status('Renaming to ' + file.sub(/^\.\//, '') + '...')
    File.rename(@file, file)
  end
end

# Display program help if help action is set. If no actions are set, default to
# either help or view, depending on whether an argument was passed.
if $opts.has_key?('help') or (!action and ARGV.empty?)
  help
  exit 0
elsif !action
  $opts['view'] = ''
end

# Treat all remaining arguments as files.
warn('No files specified.') if ARGV.length.zero?
ARGV.each do |file|
  # Process exceptions in order to produce error messages in a custom format
  # without an excessive number of nested if statements.
  begin
    # Output name of the current file.
    puts colyel(file)

    # Verify that the file exists.
    raise 'File does not exist.' unless File.exist?(file)

    # Determine whether or not the file is supported.
    raise 'File has no extension.' if file.split('.').length == 1
    ext = file.split('.').last.downcase
    raise 'Not a supported file format.' unless ext =~ /flac|ogg|mp3/

    # Create track variable based on the file format.
    if ext == 'flac'
      track = FLAC.new(file)
    elsif ext == 'ogg'
      track = Vorbis.new(file)
    elsif ext == 'mp3'
      track = ID3.new(file)
    end

    # Call methods based on the actions set.
    if $opts.has_key?('set-tag')
      track.set_tag($opts['set-tag'])
      track.write_tags
    end

    if $opts.has_key?('remove')
      track.remove
      track.write_tags
    elsif $opts.has_key?('remove-tag')
      track.remove_tag($opts['remove-tag'])
      track.write_tags
    end

    if $opts.has_key?('preview')
      track.generate
    elsif $opts.has_key?('generate')
      track.generate
      track.write_tags
    elsif $opts.has_key?('rename')
      track.rename
    end

    if $opts.has_key?('view')
      track.view
    elsif $opts.has_key?('view-tag')
      track.view_tag($opts['view-tag'])
    end
  rescue RuntimeError => message
    warn(message)
  end
end