#!/usr/bin/env ruby -W
#
# convert-video
#
# Copyright (c) 2013-2020 Don Melton
#

$LOAD_PATH.unshift File.expand_path(File.dirname(__FILE__) + '/../lib')

require 'video_transcoding/cli'

module VideoTranscoding
  class Command
    include CLI

    def about
      <<HERE
convert-video #{VERSION}
#{COPYRIGHT}
HERE
    end

    def usage
      <<HERE
Convert video file from Matroska to MP4 format or from MP4 to Matroksa format
WITHOUT TRANSCODING VIDEO.

Usage: #{$PROGRAM_NAME} [OPTION]... [FILE]...

-o, --output DIRECTORY
                    set output path
                      (default: input filename with output format extension
                        in current working directory)
    --use-m4v       use `.m4v` extension instead of `.mp4` for MP4 output
    --no-double     don't swap order of first two audio tracks
                      or add stereo copy of main surround track
    --ac3-encoder ac3|eac3
                    set AC-3 audio encoder (default: ac3)

-v, --verbose       increase diagnostic information
-q, --quiet         decrease     "           "

-h, --help          display this help and exit
    --version       output version information and exit

Requires `HandBrakeCLI`, `mp4track`, `ffmpeg` and `mkvpropedit`.
HERE
    end

    def initialize
      super
      @output       = nil
      @use_m4v      = false
      @double       = true
      @ac3_encoder  = 'ac3'
    end

    def define_options(opts)
      opts.on('-o', '--output ARG') { |arg| @output = arg }
      opts.on('--use-m4v')          { @use_m4v = true }
      opts.on('--no-double')        { @double = false }

      opts.on '--ac3-encoder ARG' do |arg|
        @ac3_encoder = case arg
        when 'ac3', 'eac3'
          arg
        else
          fail UsageError, "invalid AC-3 audio encoder: #{arg}"
        end
      end
    end

    def configure
      unless @output.nil? or File.directory? @output
        fail UsageError, "not a directory: #{@output}"
      end

      HandBrake.setup
      MP4track.setup
      FFmpeg.setup
      MKVpropedit.setup
    end

    def process_input(arg)
      Console.info "Processing: #{arg}..."
      media = Media.new(path: arg, allow_directory: false)
      Console.debug media.info.inspect
      output = resolve_output(media)
      fail "no compatible video track: #{arg}" unless media.info[:h264] or media.info[:hevc]
      fail "no video stream: #{arg}" unless media.info.has_key? :stream

      if media.info[:mkv]
        convert_to_mp4(media, output)
      else
        convert_to_mkv(media, output)
      end
    end

    def resolve_output(media)
      if media.info[:mkv]
        ext = @use_m4v ? '.m4v' : '.mp4'
      elsif media.info[:mp4]
        ext = '.mkv'
      else
        fail "unsupported input container format: #{media.path}"
      end

      output = File.basename(media.path, '.*') + ext

      unless @output.nil?
        output = File.absolute_path(@output) + File::SEPARATOR + output
      end

      fail "output file exists: #{output}" if File.exist? output
      output
    end

    def convert_to_mp4(media, output)
      map_options = [
        '-map', "0:#{media.info[:stream]}"
      ]
      copy_options = [
        '-c:v', 'copy'
      ]

      unless media.info[:audio].empty?
        track_order = media.info[:audio].keys
        stream = 0

        if @double and media.info[:audio][1][:channels] > 2.0
          if  track_order.size > 1 and
              media.info[:audio][1][:language] == media.info[:audio][2][:language] and
              media.info[:audio][2][:format] =~ /^AAC/ and
              media.info[:audio][2][:channels] <= 2.0
            first = track_order[0]
            track_order[0] = track_order[1]
            track_order[1] = first
          else
            map_options.concat([
              '-map', "0:#{media.info[:audio][1][:stream]}"
            ])
            copy_options.concat([
              '-ac:a:0', '2',
              '-c:a:0', FFmpeg.aac_encoder,
              '-b:a:0', '160k'
            ])
            stream += 1
          end
        end

        track_order.each do |track|
          map_options.concat([
            '-map', "0:#{media.info[:audio][track][:stream]}"
          ])

          if  media.info[:audio][track][:channels] > 2.0 and
              not ac3_format?(media.info[:audio][track][:format])
            copy_options.concat([
              "-ac:a:#{stream}", '6',
              "-c:a:#{stream}", @ac3_encoder,
              "-b:a:#{stream}", '640k'
            ])
          elsif media.info[:audio][track][:channels] <= 2.0 and
                not media.info[:audio][track][:format] =~ /^AAC/ and
                not ac3_format?(media.info[:audio][track][:format])
            copy_options.concat([
              "-ac:a:#{stream}", '2',
              "-c:a:#{stream}", FFmpeg.aac_encoder,
              "-b:a:#{stream}", '160k'
            ])
          else
            copy_options.concat([
              "-c:a:#{stream}", 'copy'
            ])
          end

          stream += 1
        end
      end

      forced_subtitle_track = 0

      unless media.info[:subtitle].empty?
        track_order = media.info[:subtitle].keys
        stream = 0

        track_order.each do |track|
          break unless media.info[:subtitle][track].has_key? :stream

          subtitle_encoder = case media.info[:subtitle][track][:format]
          when 'Text'
            'mov_text'
          when 'Bitmap'
            if media.info[:subtitle][track][:encoding] == 'VOBSUB'
              'copy'
            else
              nil
            end
          else
            nil
          end

          unless subtitle_encoder.nil?
            map_options.concat([
              '-map', "0:#{media.info[:subtitle][track][:stream]}"
            ])
            copy_options.concat([
              "-c:s:#{stream}", subtitle_encoder
            ])
            stream += 1

            if forced_subtitle_track == 0 and media.info[:subtitle][track][:forced]
              forced_subtitle_track = stream
            end
          end
        end
      end

      convert media, output, ['-movflags', 'disable_chpl', *map_options, *copy_options]
      mp4_media = Media.new(path: output, allow_directory: false)
      Console.debug mp4_media.info.inspect

      last_track = 0

      mp4_media.info[:audio].each do |track, info|
        if track == 1 and not info[:default]
          enabled = 'true'
        elsif track != 1 and info[:default]
          enabled = 'false'
        else
          enabled = nil
        end

        mp4_track_enable output, track, enabled
        last_track += 1
      end

      mp4_media.info[:subtitle].each do |track, info|
        if track == forced_subtitle_track and not info[:default]
          enabled = 'true'
        elsif track != forced_subtitle_track and info[:default]
          enabled = 'false'
        else
          enabled = nil
        end

        mp4_track_enable output, last_track + track, enabled
      end
    end

    def ac3_format?(track_format)
      if  (@ac3_encoder == 'ac3'  and track_format == 'AC3') or
          (@ac3_encoder == 'eac3' and track_format =~ /^(?:E-)?AC3$/)
        true
      else
        false
      end
    end

    def mp4_track_enable(output, track, enabled)
      unless enabled.nil?
        begin
          IO.popen([
            MP4track.command_name,
            '--track-index', track.to_s,
            '--enabled', enabled,
            output,
          ], 'rb', :err=>[:child, :out]) do |io|
            io.each do |line|
              Console.debug line
            end
          end
        rescue SystemCallError => e
          raise "adjusting audio enabled failed: #{e}"
        end

        fail "adjusting audio enabled failed: #{output}" unless $CHILD_STATUS.exitstatus == 0
      end
    end

    def convert_to_mkv(media, output)
      map_options = [
        '-map', "0:#{media.info[:stream]}"
      ]
      copy_options = [
        '-c:v', 'copy'
      ]
      audio_track_names = {}

      unless media.info[:audio].empty?
        track_order = media.info[:audio].keys
        stream = 0

        if  @double and
            track_order.size > 1 and
            media.info[:audio][1][:language] == media.info[:audio][2][:language] and
            media.info[:audio][1][:format] =~ /^AAC/ and
            media.info[:audio][1][:channels] <= 2.0 and
            media.info[:audio][2][:channels] > 2.0
          first = track_order[0]
          track_order[0] = track_order[1]
          track_order[1] = first
        end

        track_order.each do |track|
          map_options.concat([
            '-map', "0:#{media.info[:audio][track][:stream]}"
          ])
          copy_options.concat([
            "-c:a:#{stream}", 'copy'
          ])
          stream += 1
          audio_track_names[stream] = media.info[:audio][track][:name]
        end
      end

      subtitle_track_names = {}
      forced_subtitle_track = 0

      unless media.info[:subtitle].empty?
        track_order = media.info[:subtitle].keys
        stream = 0

        track_order.each do |track|
          subtitle_encoder = case media.info[:subtitle][track][:format]
          when 'Text'
            'text'
          when 'Bitmap'
            if media.info[:subtitle][track][:encoding] == 'VOBSUB'
              'dvdsub'
            else
              nil
            end
          else
            nil
          end

          unless subtitle_encoder.nil?
            map_options.concat([
              '-map', "0:#{media.info[:subtitle][track][:stream]}"
            ])
            copy_options.concat([
              "-c:s:#{stream}", subtitle_encoder
            ])
            stream += 1
            subtitle_track_names[stream] = media.info[:subtitle][track][:name]

            if forced_subtitle_track == 0 and media.info[:subtitle][track][:forced]
              forced_subtitle_track = stream
            end
          end
        end
      end

      convert media, output, [*map_options, *copy_options]
      mkv_media = Media.new(path: output, allow_directory: false)
      Console.debug mkv_media.info.inspect
      mkvpropedit_options = [
        '--tags', 'all:'
      ]

      mkv_media.info[:audio].each do |track, info|
        unless audio_track_names[track].nil?
          mkvpropedit_options.concat([
            '--edit', "track:a#{track}",
            '--set', "name=#{audio_track_names[track]}"
          ])
        end

        if track == 1 and not info[:default]
          default = 1
        elsif track != 1 and info[:default]
          default = 0
        else
          default = nil
        end

        unless default.nil?
          mkvpropedit_options.concat([
            '--edit', "track:a#{track}",
            '--set', "flag-default=#{default}"
          ])
        end
      end

      mkv_media.info[:subtitle].each do |track, info|
        unless subtitle_track_names[track].nil?
          mkvpropedit_options.concat([
            '--edit', "track:s#{track}",
            '--set', "name=#{subtitle_track_names[track]}"
          ])
        end
      end

      if forced_subtitle_track != 0
        mkvpropedit_options.concat([
          '--edit', "track:s#{forced_subtitle_track}",
          '--set', 'flag-forced=1'
        ])
      end

      Console.info 'Adjusting properties and tags with mkvpropedit...'

      begin
        IO.popen([
          MKVpropedit.command_name,
          *mkvpropedit_options,
          output,
        ], 'rb', :err=>[:child, :out]) do |io|
          io.each do |line|
            Console.debug line
          end
        end
      rescue SystemCallError => e
        raise "Adjusting properties and tags failed: #{e}"
      end

      fail "Adjusting properties and tags failed: #{output}" if $CHILD_STATUS.exitstatus == 2
    end

    def convert(media, output, ffmpeg_options)
      ffmpeg_command = [
        FFmpeg.command_name,
        '-hide_banner',
        '-nostdin',
        '-i', media.path,
        *ffmpeg_options,
        output
      ]
      Console.debug ffmpeg_command.inspect
      Console.info 'Converting with ffmpeg...'

      begin
        IO.popen(ffmpeg_command, 'rb', :err=>[:child, :out]) do |io|
          Signal.trap 'INT' do
            Process.kill 'INT', io.pid
          end

          buffer = ''

          io.each_char do |char|
            if (char.bytes[0] & 0x80) != 0
              buffer << char
            else
              if not buffer.empty?
                print buffer
                buffer = ''
              end

              print char
            end
          end

          if not buffer.empty?
            print buffer
          end
        end
      rescue SystemCallError => e
        raise "conversion failed: #{e}"
      end

      fail "conversion failed: #{media.path}" unless $CHILD_STATUS.exitstatus == 0
    end
  end
end

VideoTranscoding::Command.new.run
