#!/usr/bin/ruby -w

require 'fileutils'
require "musicbrainz"
require "yaml"
require "ostruct"
require 'iconv'
require "tempfile"

class Track
	attr_reader :id, :name

	def initialize(mb)
		@name = mb.result(MusicBrainz::Query::TrackGetTrackName)

		tmp = mb.result(MusicBrainz::Query::TrackGetTrackId)
		@id = mb.id_from_url(tmp)
	end
end

class Artist
	attr_reader :id, :name

	def initialize(mb)
		tmp = mb.result(MusicBrainz::Query::AlbumGetAlbumArtistId)
		@id = mb.id_from_url(tmp)

		@name = mb.result(MusicBrainz::Query::AlbumGetArtistName, 1)
	end
end

class Album
	attr_reader :id, :name, :artist, :tracks
	attr_accessor :genre, :year

	def initialize(mb)
		@genre = nil
		@year = nil

		@artist = Artist.new(mb)

		tmp = mb.result(MusicBrainz::Query::AlbumGetAlbumId)
		@id = mb.id_from_url(tmp)

		@name = mb.result(MusicBrainz::Query::AlbumGetAlbumName)

		num_tracks = mb.result(MusicBrainz::Query::AlbumGetNumTracks).to_i
		@tracks = Array.new(num_tracks)

		1.upto(num_tracks) do |j|
			mb.select(MusicBrainz::Query::SelectTrack, j)

			@tracks[j - 1] = Track.new(mb)

			mb.select(MusicBrainz::Query::Back)
		end
	end
end

class Encoder
	def initialize(config)
		@cfg = config
	end

	def encode(dest, metadata)
	end

	def apply_tags(dest, metadata)
	end

	def compute_replaygain
	end

	def suffix
	end
end

class VorbisEncoder < Encoder
	def encode(dest, m)
		<<EOF
		oggenc -q #{@cfg.vorbis_quality} -Q -o "#{dest}" -
EOF
	end

	def apply_tags(dest, m)
		Tempfile.open("rbrip.vorbiscomments") do |tf|
			tf.puts("artist=" + m[:artist])
			tf.puts("title=" + m[:track_name])
			tf.puts("album=" + m[:disc])
			tf.puts("tracknumber=" + m[:track_no])
			tf.puts("genre=" + m[:genre])
			tf.puts("date=" + m[:year].to_s)
			tf.puts("musicbrainz_albumid=" + m[:disc_id])
			tf.puts("musicbrainz_artistid=" + m[:artist_id])
			tf.puts("musicbrainz_trackid=" + m[:track_id])
			tf.flush

			FileUtils.mv(dest, dest + ".untagged")
			`vorbiscomment -w "#{dest}.untagged" "#{dest}" -c #{tf.path}`
			FileUtils.rm(dest + ".untagged")
		end
	end

	def compute_replaygain(files)
		"vorbisgain -a " + files.map { |f| f + suffix }.join(" ")
	end

	def suffix
		".ogg"
	end
end

class FlacEncoder < Encoder
	def encode(dest, m)
		s =<<EOF
flac --no-ogg
  --sign=signed --endian=little --channels=2 --bps=16 --sample-rate=44100
  -o "#{dest}" -
EOF

		s.gsub!(/\n/, " ")
		s
	end

	def apply_tags(dest, m)
		Tempfile.open("rbrip.vorbiscomments") do |tf|
			tf.puts("artist=" + m[:artist])
			tf.puts("title=" + m[:track_name])
			tf.puts("album=" + m[:disc])
			tf.puts("tracknumber=" + m[:track_no])
			tf.puts("genre=" + m[:genre])
			tf.puts("date=" + m[:year].to_s)
			tf.puts("musicbrainz_albumid=" + m[:disc_id])
			tf.puts("musicbrainz_artistid=" + m[:artist_id])
			tf.puts("musicbrainz_trackid=" + m[:track_id])
			tf.flush

			`metaflac --import-tags-from=#{tf.path} #{dest}`
		end
	end

	def compute_replaygain(files)
		"metaflac --add-replay-gain " + files.map { |f| f + suffix }.join(" ")
	end

	def suffix
		".flac"
	end
end

class String
	def sanitize
		# cut leading and trailing "..."
		if self[0..2] == "..."
			tmp = self[3..-1]
		else
			tmp = self.dup
		end

		if tmp[-3..-1] == "..."
			tmp = tmp[0..-4]
		end

		tmp.gsub(/\s+/, "_").gsub(/&/, "and").delete(",'()[]/\\!?\":").
		tr("-", "_").
		squeeze(" ").squeeze(".").squeeze("_")
	end

	def sanitize!
		replace(sanitize)
	end

	def rpad(total)
		n = [total - length, 0].max
		self + (" " * n)
	end
end

def locale_from_utf8(s)
	# FIXME: Encoding is only available in Ruby 1.9
	to_charset = Encoding.locale_charmap
	utf8 = 'UTF-8'.freeze

	if to_charset == utf8
		s.dup
	else
		Iconv.iconv to_charset, utf8, s
	end
end

cfg_file = ARGV.first || File.expand_path("~/.rbrip.yaml")

begin
	config = YAML.load(File.open(cfg_file))
rescue
	config = OpenStruct.new(
		:device => "/dev/hdd",
		:destdir => ".",
		:dirmask => "%a/%y-%l",
		:filemask => "%n-%a-%t",
		:sanitize_filenames => true,
		:vorbis_quality => 5
	)

	File.open(cfg_file, "w") { |f| YAML.dump(config, f) }
end

mb = MusicBrainz::Client.new
mb.device = config.device

unless mb.query(MusicBrainz::Query::GetCDInfo)
	STDERR.puts "Error: " + mb.error
	exit 1
end

num_albums = mb.result(MusicBrainz::Query::GetNumAlbums).to_i

if num_albums < 1
	STDERR.puts "Album not found. Guess you need to submit it:"
	puts mb.url
	exit 1
elsif num_albums > 1
	STDERR.puts "Multiple albums found. What now?"
	exit 1
end

mb.select(MusicBrainz::Query::Rewind)
mb.select(MusicBrainz::Query::SelectAlbum, 1)

album = Album.new(mb)

if album.id == ""
	STDERR.puts "Album not found. Guess you need to submit it:"
	puts mb.url
	exit 1
end

puts "Disc info:\n
 Artist: #{album.artist.name.rpad(20)} (#{album.artist.id})
 Name:   #{album.name.rpad(20)} (#{album.id})

 Tracks:\n"

album.tracks.each_with_index do |t, i|
	tmp = locale_from_utf8 t.name
	puts "%3s." % (i + 1) + " #{tmp}"
end

print "\nContinue [Y/n]? "
STDOUT.flush
exit if STDIN.getc == ?n

print "\nEnter year of release: "
STDOUT.flush

album.year = STDIN.gets.strip.to_i

print "Enter genre: "
STDOUT.flush

album.genre = STDIN.gets.strip

# where to put the encoded files?
destdir = config.dirmask.dup

dir_map = {
	"%A" => album.artist.name.sanitize,
	"%L" => album.name.sanitize,
	"%Y" => album.year.to_s
}

if config.sanitize_filenames
	dir_map.each_value { |v| v.sanitize! }
end

dir_map.each do |k, v|
	destdir.gsub!(/#{k}/, v)
	destdir.gsub!(/#{k.downcase}/, v.downcase)
end

destdir = File.join(config.destdir, destdir)
FileUtils.makedirs(destdir) unless File.directory?(destdir)

puts "\nRipping to #{destdir}\n\n"

files = []
enc = VorbisEncoder.new(config)
#enc = FlacEncoder.new(config)

metadata = {
	:artist => album.artist.name,
	:artist_id => album.artist.id,
	:disc => album.name,
	:disc_id => album.id,
	:genre => album.genre,
	:year => album.year
}
file_map = dir_map.dup

album.tracks.each_with_index do |t, i|
	t = album.tracks[i]

	track_no = i + 1
	file = config.filemask.dup

	file_map["%N"] = "%02i" % track_no

	if config.sanitize_filenames
		file_map["%T"] = t.name.sanitize
	else
		file_map["%T"] = t.name
	end

	metadata[:track_no] = track_no.to_s
	metadata[:track_name] = t.name
	metadata[:track_id] = t.id

	file_map.each do |k, v|
		file.gsub!(/#{k}/, v)
		file.gsub!(/#{k.downcase}/, v.downcase)
	end

	if config.sanitize_filenames
		if file[-1] == ?.
			file = file[0..-2]
		end
	end

	file = File.join(destdir, file)
	files << file

	rip_cmd = "cdparanoia -q -d \"#{config.device}\" -w -- #{track_no} -"
	encode_cmd = enc.encode(file + enc.suffix, metadata)

	puts "Ripping: #{track_no}. #{t.name} => #{file + enc.suffix}"
	`#{rip_cmd} | #{encode_cmd}`

	puts "Applying tags..."
	enc.apply_tags(file + enc.suffix, metadata)
end

rplgain_cmd = enc.compute_replaygain(files)

unless rplgain_cmd.nil?
	puts "Computing replaygain..."
	`#{rplgain_cmd} > /dev/null`
end