#!/usr/bin/env ruby
#
# omptagger
# http://github.com/omp/omptagger
#
# Copyright 2007-2010 David Vazgenovich Shakaryan <dvshakaryan@gmail.com>
# 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'

# Define colours.
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

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

# Output a tag name and value.
def output(tag, val)
  # Alter capitalisation of tag name for aesthetic purposes.
  tag = tag.downcase.capitalize unless tag == 'ISRC'
  tag = 'TrackNumber' if tag == 'Tracknumber'

  # The number statically defined here should be two greater than the length of
  # the longest possible tag name.
  puts '    ' + colcyn(tag) + ' ' * (14 - tag.length) + val
end

# Split a filename based on a naming scheme. An asterisk is a wildcard key.
def schemesplit(scheme, file, keystr)
  # Determine keys and values using a regular expression.
  regexp = /\A#{Regexp.escape(scheme).gsub(/%([#{keystr}]|\\\*)/, '(.*?)')}\Z/
  keys = scheme.scan(regexp).flatten
  vals = file.scan(regexp).flatten

  raise 'Filename does not match naming scheme.' if keys.length != vals.length

  # Map scheme keys to values.
  result = Hash.new
  keys.zip(vals).each do |key, val|
    result[key[1,1]] = val unless key == '%*'
  end

  return result
end

# Display help information.
def help
  options = [
    ['--view        -v', 'View all tags'],
    ['--view-tag    -t', 'View a tag'],
    ['--set-tag     -s', 'Set a tag'],
    ['--remove      -r', 'Remove all tags'],
    ['--remove-tag  -d', 'Remove a tag'],
    ['--generate    -g', 'Generate tags based on filename'],
    ['--rename      -m', 'Generate filename based on tags'],
    ['--scheme      -n', 'Specify a file naming scheme'],
    ['--pretend     -p', 'Do not make any actual changes'],
    ['--no-colour   -c', 'Disable use of colour in output'],
    ['--list        -l', 'Display list of available tags'],
    ['--help        -h', 'Display help information']]
  notes = [
    'The default file naming scheme is Artist - Title.',
    'Schemes must be specified prior to actions that use them.',
    'Underscores in filenames are converted to spaces in tags.',
    'Backslashes in filenames are converted to forward slashes in tags.']

  puts colyel('Usage:') + ' omptagger ' + colgrn('[options] [files]')
  puts
  puts colyel('Options:')
  options.each do |option, description|
    puts '  ' + colgrn(option) + '  ' + description
  end
  puts
  puts colyel('Notes:')
  notes.each do |note|
    puts '  ' + colgrn('*') + ' ' + note
  end
  exit
end

# Display list of available tags.
def list
  puts '┌───────────────────────────────────┐',
       '│ Keys   Vorbis Comments   ID3 Tags │',
       '├───────────────────────────────────┤',
       '│ %a     Artist            Artist   │',
       '│ %b     Album             Album    │',
       '│ %d     Date              Year     │',
       '│ %n     TrackNumber       Track    │',
       '│ %t     Title             Title    │',
       '│ %*     Wildcard          Wildcard │',
       '├───────────────────────────────────┤',
       '│        Contact                    │',
       '│        Copyright                  │',
       '│        Description                │',
       '│        Genre                      │',
       '│        ISRC                       │',
       '│        License                    │',
       '│        Location                   │',
       '│        Organization               │',
       '│        Performer                  │',
       '│        Version                    │',
       '└───────────────────────────────────┘'
  exit
end

# Class for Vorbis comments; used by FLAC and Vorbis files.
class VorbisComments
  # Possible tags.
  TAGS = ['ALBUM',
          'ARTIST',
          'CONTACT',
          'COPYRIGHT',
          'DATE',
          'DESCRIPTION',
          'GENRE',
          'ISRC',
          'LICENSE',
          'LOCATION',
          'ORGANIZATION',
          'PERFORMER',
          'TITLE',
          'TRACKNUMBER',
          'VERSION']

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

  def initialize(file)
    @file = @origfile = file
    @tags = Hash.new
    @write = false

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

      @tags[tag] = val.gsub(/^#{tag}=/i, '').split("\n")
    end
  end

  def view
    raise 'No tags are set.' if @tags.empty?

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

  def view_tag(arg)
    tag = arg.upcase
    raise tag + ' is not a valid tag.' unless TAGS.include?(tag)

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

    val.each do |val|
      output(tag, val)
    end
  end

  def set_tag(arg)
    arg = arg.split('=', 2)
    tag = arg.shift.upcase
    raise tag + ' is not a valid tag.' unless TAGS.include?(tag)

    @tags[tag] = arg
    view_tag(tag)

    @write = true
  end

  def remove
    raise 'No tags are set.' if @tags.empty?

    @tags.clear
    @write = true
  end

  def remove_tag(arg)
    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)
    @write = true
  end

  def generate(scheme)
    val = File.basename(@file, File.extname(@file))
    val = val.gsub('_', ' ').gsub('\\', '/')
    scheme = scheme.gsub('_', ' ').gsub('\\', '/')
    scheme = schemesplit(scheme, val, KEYS.keys.join)

    tagval = Array.new

    scheme.each do |tag, val|
      tagval << [KEYS[tag], val]
    end

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

    @write = true
  end

  def rename(scheme)
    file = scheme

    file.scan(/%([#{KEYS.keys.join}])/).flatten.each do |key|
      raise 'Missing ' + KEYS[key] + ' tag.' unless @tags[KEYS[key]]
      file = file.gsub('%' + key, @tags[KEYS[key]].first)
    end

    file = File.dirname(@file) + '/' + file.gsub('/', '_') +
      File.extname(@file)

    raise 'Generated filename and current filename are identical.' if
      File.basename(@file) == File.basename(file)

    output('FILENAME', file.sub(/^\.\//, ''))
    raise 'File with generated name already exists.' if File.exist?(file)

    @file = file
  end

  def finalise
    if @write
      clear_tags

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

    File.rename(@origfile, @file) unless @file == @origfile
  end
end

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

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

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

# Subclass of VorbisComments containing methods unique to Vorbis files.
class Vorbis < VorbisComments
  def read_tag(tag)
    %x(vorbiscomment -l -- '#{esc(@origfile)}').scan(/#{tag}=.*/i).join("\n")
  end

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

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

# Class for ID3 tags; used by MP3 files.
class ID3
  # Possible tags.
  TAGS = Hash['ALBUM'   => :TALB,
              'ARTIST'  => :TPE1,
              'COMMENT' => :COMM,
              'TITLE'   => :TIT2,
              'TRACK'   => :TRCK,
              'YEAR'    => :TYER]

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

  def initialize(file)
    @file = @origfile = file
    @tags = ID3Lib::Tag.new(@file)
    @write = false
  end

  def view
    raise 'No tags are set.' if @tags.empty?

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

  def view_tag(arg)
    tag = arg.upcase
    raise tag + ' is not a valid tag.' unless TAGS.has_key?(tag)

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

    output(tag, val)
  end

  def set_tag(arg)
    arg = arg.split('=', 2)
    tag = arg.shift.upcase
    raise tag + ' is not a valid tag.' unless TAGS.has_key?(tag)

    @tags.set_frame_text(TAGS[tag], arg.to_s)
    view_tag(tag)

    @write = true
  end

  def generate(scheme)
    val = File.basename(@file, File.extname(@file))
    val = val.gsub('_', ' ').gsub('\\', '/')
    scheme = scheme.gsub('_', ' ').gsub('\\', '/')
    scheme = schemesplit(scheme, val, KEYS.keys.join)

    tagval = Array.new

    scheme.each do |tag, val|
      tagval << [KEYS[tag], val]
    end

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

    @write = true
  end

  def remove
    raise 'No tags are set.' if @tags.empty?

    @tags.clear
    @write = true
  end

  def remove_tag(arg)
    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])
    @write = true
  end

  def rename(scheme)
    file = scheme

    file.scan(/%([#{KEYS.keys.join}])/).flatten.each do |key|
      raise 'Missing ' + KEYS[key] + ' tag.' if
        @tags.frame_text(TAGS[KEYS[key]]).nil?

      file = file.gsub('%' + key, @tags.frame_text(TAGS[KEYS[key]]))
    end

    file = File.dirname(@file) + '/' + file.gsub('/', '_') +
      File.extname(@file)

    raise 'Generated filename and current filename are identical.' if
      File.basename(@file) == File.basename(file)

    output('FILENAME', file.sub(/^\.\//, ''))
    raise 'File with generated name already exists.' if File.exist?(file)

    @file = file
  end

  def finalise
    if @write
      if @tags.empty? or (@tags.length == 1 and @tags.frame_text(:TLEN))
        @tags.strip!
      else
        @tags.update!
      end
    end

    File.rename(@origfile, @file) unless @file == @origfile
  end
end

# Parse options.
actions = Array.new
options = Array.new
scheme = '%a - %t'
GetoptLong.new(
  ['--view',       '-v', GetoptLong::NO_ARGUMENT],
  ['--view-tag',   '-t', GetoptLong::REQUIRED_ARGUMENT],
  ['--set-tag',    '-s', GetoptLong::REQUIRED_ARGUMENT],
  ['--remove',     '-r', GetoptLong::NO_ARGUMENT],
  ['--remove-tag', '-d', GetoptLong::REQUIRED_ARGUMENT],
  ['--generate',   '-g', GetoptLong::NO_ARGUMENT],
  ['--rename',     '-m', GetoptLong::NO_ARGUMENT],
  ['--scheme',     '-n', GetoptLong::REQUIRED_ARGUMENT],
  ['--pretend',    '-p', GetoptLong::NO_ARGUMENT],
  ['--no-colour',  '-c', GetoptLong::NO_ARGUMENT],
  ['--list',       '-l', GetoptLong::NO_ARGUMENT],
  ['--help',       '-h', GetoptLong::NO_ARGUMENT]
).each do |opt, arg|
  opt = opt.sub(/^--/, '').gsub('-', '_').intern

  case opt
  when :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
  when :scheme
    scheme = arg
  when :pretend, :help, :list
    options << opt
  when :view_tag, :set_tag, :remove_tag
    actions << [opt, arg]
  when :generate, :rename
    actions << [opt, scheme]
  else
    actions << [opt]
  end
end

# Deal with help and list options. Set a default action if one was not
# specified. Display error if an action was specified, but no arguments are
# present.
if options.include?(:help)
  help
elsif options.include?(:list)
  list
elsif actions.empty? and ARGV.empty?
  help
elsif actions.empty?
  actions = [[:view]]
elsif ARGV.empty?
  puts colred('ERROR:') + ' No files specified.'
end

# Treat each argument as a file.
ARGV.each do |file|
  # Continue with next file if an error is raised.
  begin
    # Output filename and verify that file exists.
    puts colyel(file)
    raise 'File does not exist.' unless File.exist?(file)

    # Determine file format.
    case File.extname(file).downcase
    when '.flac'
      track = FLAC.new(file)
    when '.ogg'
      track = Vorbis.new(file)
    when '.mp3'
      track = ID3.new(file)
    else
      raise 'File extension not recognised.'
    end

    # Run specified actions.
    actions.each do |action|
      # Continue with next action if an error is raised.
      begin
        # Output action name and arguments.
        puts '  ' + colgrn('Action: ' + action.join(' '))

        # Run action.
        case action.first
        when :view, :remove
          track.send(action.first)
        when :view_tag, :set_tag, :remove_tag, :generate, :rename
          track.send(action.first, action.last)
        end
      rescue RuntimeError => message
        puts '    ' + colred('ERROR:') + ' ' + message
      end
    end

    # Finalise changes.
    track.finalise unless options.include?(:pretend)
  rescue RuntimeError => message
    puts '  ' + colred('ERROR:') + ' ' + message
  end
end