#!/usr/bin/ruby -w
#
# Copyright (C) 2007-2008 Thomer M. Gil [http://thomer.com/]
# Thanks to Brian Moore and Justin Payne for bugfixes and suggestions.
#
#
# This program is free software. You may distribute it under the terms of
# the GNU General Public License as published by the Free Software
# Foundation, version 2.
# 
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
# 
# This program converts video files to mp4, suitable to be played on an iPod.
# It is careful about maintaining the proper aspect ratio.
#

require 'getoptlong'

# will automatically try with -vcoded libxvid, also.
# will automatically try with -acodec libfaac, also.
DEFAULT_ARGS = "-f mp4 -vcodec xvid -maxrate 1000 -qmin 3 -qmax 5 -bufsize 4096 -g 300 -acodec aac"
DEFAULT_AUDIO_BITRATE = 128
DEFAULT_VIDEO_BITRATE = 400
WIDTH = 320.0
HEIGHT = 240.0

$options = {}
opts = GetoptLong.new(
  [ "-a", GetoptLong::REQUIRED_ARGUMENT ],      # audio bitrate
  [ "-h", GetoptLong::NO_ARGUMENT ],            # help
  [ "-b", GetoptLong::REQUIRED_ARGUMENT ],      # video bitrate
  [ "-v", GetoptLong::NO_ARGUMENT ]             # verbose
)
opts.each { |opt, arg| $options[opt] = arg }

if $options['-h']
  puts <<EOF
mp4ize - encode videos to mp4 for an iPod

Usage: mp4ize file1.avi [file2.mpg [file3.asf [...]]]


Options:

  -h       : this help
  -v       : verbose
  -a RATE  : override default audio bitrate (#{DEFAULT_AUDIO_BITRATE})
  -b RATE  : override default video bitrate (#{DEFAULT_VIDEO_BITRATE})
EOF
  exit
end

audio_bitrate = $options['-a'] || DEFAULT_AUDIO_BITRATE
video_bitrate = $options['-b'] || DEFAULT_VIDEO_BITRATE

ARGV.each do |infile|
  outfile = infile.dup
  ext = File.extname(outfile)
  outfile.sub!(/#{ext}$/, '.mp4')

  # open the file to figure out the aspect ratio
  duration, w, h = 0.0, nil, nil
  IO.popen("ffmpeg -i '#{infile}' 2>&1") do |pipe|
    pipe.each_line do |line|
      puts line
      if line.match(/Video:.+ (\d+)x(\d+)/)
        w, h = $1.to_f, $2.to_f
      elsif line.match(/Duration:\s+(\d+):(\d+):(\d+)\.(\d+)/)
        duration += $1.to_f * 3600
        duration += $2.to_f * 60
        duration += $3.to_f
        duration += $4.to_f / 10
      end
    end
  end

  begin
    aspect = w/h
  rescue
    puts "Couldn't figure out aspect ratio."
    exit
  end

  height = (WIDTH / aspect.to_f).to_i
  height -= 1 if height % 2 == 1
  pad = ((HEIGHT - height.to_f) / 2.0).to_i
  pad -= 1 if pad % 2 == 1

  File.unlink(outfile) if File.exists?(outfile)
  cmd = "ffmpeg -i '#{infile}' #{DEFAULT_ARGS} -s #{WIDTH.to_i}x#{height} -padtop #{pad} -padbottom #{pad} -ab #{audio_bitrate} -b #{video_bitrate} '#{outfile}' 2>&1"
  puts cmd if $options['-v']

  # We could just call "system cmd" here, but we want the exit code of mp4ize
  # to tell us whether the duration of the generated mp4 equals the duration
  # of the original movie.  Exits with a non-zero code if the two are not
  # within 1% of each other.

  time = 0
  STDOUT.sync = true

  # try with -vcodec libxvid and -vcodec xvid
  # try with -acodec libfaac and -acodec aac
  catch(:done) do
    3.times do
      catch(:retry) do
        IO.popen(cmd) do |f|
          # ffmpeg keeps refreshing one line, so separate on \r
          f.each_line("\r") do |line|
            printf line
            if line.match(/Unknown.*code.*xvid/)
              cmd.sub!('-vcodec xvid', '-vcodec libxvid')
              throw :retry
            elsif line.match(/Unknown.*code.*aac/)
              cmd.sub!('-acodec aac', '-acodec libfaac')
              throw :retry
            elsif line.match(/time=([^\s]+)/)
              time = $1.to_f
            end
          end
        end
        throw :done
      end
    end
  end


  # return completeness of mp4 file
  puts "expected duration: #{duration}" if $options['-v']
  puts "encoded duration: #{time}" if $options['-v']
  exit((time <= duration * 1.01) && (time >= duration * 0.99))
end