#!/usr/bin/ruby -w
#
# Copyright (C) 2007-2012 Thomer M. Gil [http://thomer.com/]
#
# Thanks to Brian Moore, Justin Payne, Matt Spitz, Martyn Parker,
# Jean-Francois Macaud, Thomas Hannigan, Anisse Astier, Juanma Hernández,
# Trung Huynh, Mark Ryan, William Lupton, John Taggart, Tobias Engel,
# Victor Yang, and Mike Jing for bugfixes and suggestions.
#
# Oct. 14, 2008: show percentage progress. add -t and -w flags.
# Jan. 11, 2009: switch to bit/s bitrates for newer ffmpeg versions.
#                add --iphone option.
#                add -y option to ffmpeg (overwrite).
# Jan. 20, 2009: don't exit early when processing multiple files.
# Feb. 17, 2009: deal with "Invalid pixel aspect ratio" error.
# Apr.  1, 2009: new --outdir parameter.
# May  22, 2009: handle filenames with quotes and whitespace.
# Oct   6, 2009: fix bug where we forget to read stderr
# Nov.  5, 2009: fix -v, -t, and -w command line options
#                removed bogus 'here' debug statement
# Oct. 27, 2010: assume ffmpeg 0.6: use libxvid and libfaac by default,
#                add "k" to -bufsize, -ab, and -b parameters.
# May   2, 2011: added iphone4 option
#                use ffmpeg, not /usr/bin/ffmpeg
# June 14, 2011: fix ffmpeg command line for Ubuntu 11.04 (Natty)
# Jan. 18, 2012: fix for newer ffmpeg version 0.7.8
# Feb.  1, 2012: don't clobber existing files, unless --force
# Mar. 10, 2012: backwards -pad compatibility
# Apr. 20, 2012: -aspect fix thanks to Mike Jing
#
# 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
# or an iPhone. It is careful about maintaining the proper aspect ratio.
#

require 'getoptlong'
require 'open3'

# see also http://code.google.com/p/mythpodcaster/wiki/TranscodingProfiles
DEFAULT_ARGS = "-f mp4 -y -vcodec libxvid -maxrate 1000 -mbd 2 -qmin 3 -qmax 5 -g 300 -acodec aac -strict experimental"
DEFAULT_BUFSIZE = "4096k"
DEFAULT_AUDIO_BITRATE = "128k"
DEFAULT_VIDEO_BITRATE = "400k"
IPOD_WIDTH = 320.0
IPOD_HEIGHT = 240.0
IPHONE_WIDTH = 480.0
IPHONE_HEIGHT = 320.0
IPHONE4_WIDTH = 960.0
IPHONE4_HEIGHT = 640.0

$options = {}
opts = GetoptLong.new(*[
  [ "--audio", "-a", GetoptLong::REQUIRED_ARGUMENT ],  # audio bitrate
  [ "--help", "-h", GetoptLong::NO_ARGUMENT ],         # help
  [ "--video", "-b", GetoptLong::REQUIRED_ARGUMENT ],  # video bitrate
  [ "--verbose", "-v", GetoptLong::NO_ARGUMENT ],      # verbose
  [ "--width", "-w", GetoptLong::REQUIRED_ARGUMENT ],  # override width
  [ "--height", "-t", GetoptLong::REQUIRED_ARGUMENT ], # override height
  [ "--iphone", "-i", GetoptLong::NO_ARGUMENT ],       # set width/height
  [ "--iphone4", "-4", GetoptLong::NO_ARGUMENT ],      # set width/height
  [ "--threads", "-n", GetoptLong::REQUIRED_ARGUMENT ],# number of enc threads
  [ "--outdir", "-o", GetoptLong::REQUIRED_ARGUMENT ], # dir where to write files
  [ "--force", "-f", GetoptLong::NO_ARGUMENT ],        # overwrite output files
])
opts.each { |opt, arg| $options[opt] = arg }

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

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


Options:

  -h/--help          : this help
  -v/--verbose       : verbose

  -a/--audio RATE    : override default audio bitrate (#{DEFAULT_AUDIO_BITRATE})
  -b/--video RATE    : override default video bitrate (#{DEFAULT_VIDEO_BITRATE})

  -w/--width WIDTH   : over default width (#{IPOD_WIDTH.to_i})
  -t/--height HEIGHT : over default height (#{IPOD_HEIGHT.to_i})
  -i/--iphone        : same as --width #{IPHONE_WIDTH.to_i} --height #{IPHONE_HEIGHT.to_i}
  -4/--iphone4       : same as --width #{IPHONE4_WIDTH.to_i} --height #{IPHONE4_HEIGHT.to_i}

  -n/--threads N     : number of encoding threads
  -o/--outdir O      : write files to given directory

  -f/--force         : overwrite .mp4 if it already exists (default off)
EOF
  exit
end

# not both --iphone and --iphone4
if $options['--iphone'] && $options['--iphone4']
  warn "You can't use both --iphone and --iphone4"
  exit 1
end

# --iphone sets --width and --height
if ($options['--iphone'] || $options['--iphone4']) && ($options['--width'] || $options['--height'])
  warn "You can't use --iphone or --iphone4 together with --width or --height."
  exit 1
end

if $options['--iphone']
  $options['--width'] = $options['-w'] = IPHONE_WIDTH
  $options['--height'] = $options['-t'] = IPHONE_HEIGHT
elsif $options['--iphone4']
  $options['--width'] = $options['-w'] = IPHONE4_WIDTH
  $options['--height'] = $options['-t'] = IPHONE4_HEIGHT
end

audio_bitrate = $options['--audio'] || DEFAULT_AUDIO_BITRATE
video_bitrate = $options['--video'] || DEFAULT_VIDEO_BITRATE

opt_threads = ""
if $options['--threads']
  opt_threads = "-threads #{$options['--threads']}"
end

ARGV.each do |infile|
  if !File.exists?(infile)
    warn "#{infile} does not exist"
    exit 1
  end

  outfile = infile.dup
  ext = File.extname(outfile)
  outfile.sub!(/#{ext}$/, '.mp4')
  if $options['--outdir']
    if !File.directory?($options['--outdir'])
      warn "#{$options['--outdir']} does not exist or is not a directory. exiting."
      exit 1
    end
    outfile = File.join($options['--outdir'], File.basename(outfile))
  end

  if File.exists?(outfile) && !$options['--force']
    warn "#{outfile} already exists. exiting."
    exit 1
  end

  # open the file to figure out the aspect ratio
  duration, w, h = 0.0, nil, nil
  Open3.popen3("ffmpeg", "-i", infile) do |stdin, stdout, stderr|
    [stdout, stderr].each do |io|
      io.each_line do |line|
        if line.match(/(Video|Stream).+ (\d+)x(\d+)/)
          w, h = $2.to_f, $3.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
  end

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

  user_width = $options['--width'] ? $options['--width'].to_i : IPOD_WIDTH
  user_height = $options['--height'] ? $options['--height'].to_i : IPOD_HEIGHT

  width = user_width.to_i
  height = (width / aspect.to_f).to_i
  height -= (height % 2)
  pad = ((user_height - height.to_f) / 2.0).to_i
  pad -= (pad % 2)
  p_width = width
  p_height = height + (pad * 2)
  p_x = 0
  p_y = pad
  padarg1, padarg2 = "padtop", "padbottom"

  # recalculate using the height as the baseline rather than the width
  if pad < 0
    height = user_height.to_i
    width = (height * aspect.to_f).to_i
    width -= (width % 2)
    pad = ((user_width - width.to_f)/2.0).to_i
    pad -= (pad % 2)
    p_width = width + (pad * 2)
    p_height = height
    p_x = pad
    p_y = 0
    padarg1, padarg2 = "padleft", "padright"
  end

  aspect_out = p_width.to_f/p_height.to_f

  File.unlink(outfile) if File.exists?(outfile)

  # use %infile% and %outfile% and replace those after the split() so that we
  # don't split() a filename that has spaces in it.
  cmd = "ffmpeg -i %infile% #{DEFAULT_ARGS} -bufsize #{DEFAULT_BUFSIZE} -s #{width}x#{height} -vf pad=#{p_width}:#{p_height}:#{p_x}:#{p_y} -ab #{audio_bitrate} -b #{video_bitrate} #{opt_threads} %outfile%"
  puts cmd if $options['--verbose']

  # 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

  # modify -ab and -b arguments to kbit if necessary
  catch(:done) do
    3.times do
      catch(:retry) do
        # split arguments, set %infile% and %outfile%
        cmd_array = cmd.split(/\s+/)
        cmd_array.collect! {|s| s.sub(/^%infile%$/, infile)}
        cmd_array.collect! {|s| s.sub(/^%outfile%$/, outfile)}

        puts "\n\n------------------------------------" if $options['--verbose']
        puts "cmdline: #{cmd_array.join(' ')}" if $options['--verbose']

        Open3.popen3(*cmd_array) do |stdin, stdout, stderr|
          io = select([stdout, stderr], nil, nil, 10)
          3.times do |std| # both stdout and stderr
            next if io[0][std].nil?
            io[0][std].each_line("\r") do |line|
              puts "the line is #{line}" if $options['--verbose']
              printf("\r%.2f%% | ", time / duration * 100.0)
              print line
              if line.match(/Invalid pixel aspect ratio/)
                cmd.sub!("-s #{width}x#{height}", "-s #{width}x#{height} -aspect #{aspect_out}")
                throw :retry
              elsif line.match(/Unrecognized option 'vf'/)
                cmd.sub!("-vf pad=#{p_width}:#{p_height}:#{p_x}:#{p_y}", "-#{padarg1} #{pad} -#{padarg2} #{pad}")
                throw :retry
              elsif line.match(/time=([^\s]+)/)
                time = $1.to_f
              end
            end
          end
        end
        throw :done
      end
    end
  end

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