#! /usr/bin/env ruby
# vim: ai ts=2 sts=2 sw=2 expandtab
#
# ruutu-dl - Download video from ruutu.fi and jimtv.fi
# Copyright (C) 2013  Tommi Saviranta <wnd@iki.fi>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# 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.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


require 'net/http'
require 'time'
require 'pty'
require 'getoptlong'

APP_NAME = 'ruutu-dl'
APP_VERSION = '0.2.3'

# TODO partial_file



class RuutuDownloader
  def self.prerequisites?
    if !has_bin?('flvstreamer')
      raise RuntimeError, 'flvstreamer'
    end
  end

  def self.get(uri, opts = {})
    dl = RuutuDownloader.new(uri, opts)
    dl.get
  end

  def initialize(uri, opts = {})
    @web_uri = uri
    @out_file = opts[:out_file]
    @retry_count = opts[:retry_count] || 0
  end

  def get
    rtmp = find_rtmp
    generate_output_filename(web_data, rtmp) if @out_file.nil?
    download(rtmp, @out_file, @retry_count)
  end

  private

  def self.has_bin?(exe)
    !ENV['PATH'].split(':').find{|dir| File.executable?("#{dir}/#{exe}")}.nil?
  end

  def download(rtmp, out_file, retry_count)
    KludgyRTMPDumper.get(rtmp, out_file, :retry_count => retry_count)
  end

  def web_data
    @web_data ||= Net::HTTP.get(URI(@web_uri))
  end

  def find_rtmp
    config_uri = find_config_uri
    config_data = Net::HTTP.get(URI(config_uri))
    find_rtmp_uri(config_data)
  end

  def find_config_uri
    uris = find_media_configurator + find_media_xml_cache
    if uris.size == 0
      raise RuntimeError, 'Cannot find config uri'
    elsif uris.size > 1
      raise RuntimeError, 'Multiple configurations found'
    else
      uris.first
    end
  end

  def find_media_configurator
    web_data.each_line.select{|line| line.include?('/media_configurator')}.
      map{|line| line.sub(/.*http/, 'http').sub(/".*/, '').chomp}
  end

  def find_media_xml_cache
    matcher = Regexp.new(/(http:.*?\/media-xml-cache\?id=\d+)/)
    web_data.each_line.select{|line| matcher.match(line)}.
      map{|line| matcher.match(line)[1].sub(/.*http/, 'http')}
  end

  def find_rtmp_uri(data)
    re = Regexp.new(/SourceFile/)
    uris = data.each_line.select{|line| re.match(line)}
    if uris.count == 0
      raise RuntimeError, 'Cannot find SourceFile'
    elsif uris.count > 1
      raise RuntimeError, 'Multiple SourceFiles'
    else
      uris.first.sub(/.*<SourceFile>/, '').sub(/<\/SourceFile>/, '').chomp
    end
  end

  def generate_output_filename(html, rtmp_uri)
    re = Regexp.new('<title>')
    titles = web_data.each_line.select{|line| re.match(line)}
    if titles.count == 0
      raise RuntimeError, 'Cannot find title'
    elsif titles.count > 1
      raise RuntimeError, 'Multiple titles'
    else
      src = titles.first.sub(/.*<title>/, '').
        sub(/<.*/, '').sub(/ \| .*/, '').chomp
      suffix = rtmp_uri.sub(/.*\./, '')
      @out_file = "#{src}.#{suffix}"
    end
  end


  class FileState < Struct.new(:size, :tail, :tail_size, :n)
  end


  class KludgyRTMPDumper
    class LikelyError < RuntimeError
    end

    MAX_TIMEOUT = 10.0
    MIN_TIMEOUT = 0.5
    MIN_ADDITIONAL_DATA = 256 * 1024
    TAIL_SIZE = 256 * 1024

    attr_accessor :retry_count, :retry
    attr_accessor :min_timeout, :max_timeout


    def self.get(uri, out_file, opts = {})
      dumper = KludgyRTMPDumper.new(uri, out_file)
      dumper.retry_count = opts[:retry_count]
      dumper.retry = true if dumper.retry_count && dumper.retry_count > 0
      dumper.min_timeout = [opts[:min_timeout], MIN_TIMEOUT].compact.max
      dumper.max_timeout = [opts[:max_timeout], MAX_TIMEOUT].compact.min
      dumper.download
    end

    def initialize(uri, out_file)
      @uri = uri
      @out_file = out_file
      @cmd = "flvstreamer -e -r #{@uri} -o '#{@out_file}' 2>&1"
    end

    def init_known_states
      @known_states = Array.new
      if File.exists?(@out_file)
        size = File.size(@out_file)
        if tail = out_file_tail
          state = FileState.new(size, tail, TAIL_SIZE, @n)
          @known_states.push(state)
        end
      else
        state = FileState.new(0, nil, 0)
        @known_states.push(state)
      end
    end

    def download
      timeout = @min_timeout * 2.0
      init_known_states

      done = false
      until done
        old_size = File.size?(@out_file)
        begin
          done = run_streamer(timeout)

          new_size = File.size(@out_file)
          if old_size && new_size > old_size
            timeout = [timeout / 1.5, @min_timeout].max
            if new_size - old_size > MIN_ADDITIONAL_DATA
              tail = out_file_tail
              state = FileState.new(new_size, tail, TAIL_SIZE, @n)
              @known_states.push(state)
              while @known_states.count > 10
                @known_states.shift
              end
            end
          else
            timeout = [timeout * 1.5, @max_timeout].min
          end
        rescue LikelyError => e
          if @retry
            new_size = File.size(@out_file)
            check_tail_and_maybe_restore
          else
            $stderr.puts <<EOF
Child exited with error code.
Check output and consider deleting target file and retrying.
Also consider using --retry switch.

Download failed.
EOF
          end
          if @retry_count > 0
            $stderr.puts "Retrying. #{@retry_count} attempt(s) left."
            @retry_count -= 1
            done = false
          else
            $stderr.puts <<EOF
Out of retries.

Download failed.
EOF
            return false
          end
        end
      end
    end

    private

    def out_file_tail
      len = File.size(@out_file)
      return nil if len < TAIL_SIZE
      File.open(@out_file, "rb") do |f|
        f.seek(len - TAIL_SIZE, IO::SEEK_SET)
        f.read(TAIL_SIZE)
      end
#      IO.binread(@out_file, TAIL_SIZE, len - TAIL_SIZE)
    end

    def check_tail_and_maybe_restore
      if !@known_states.empty?
        restore_file_from_known
      end
    end

    def restore_file_from_known
      state = @known_states.pop
      restore_file(state)
    end

    def restore_file(state)
      $stderr.puts "Restoring file to #{state.size} bytes"
      if state.size == 0
        File.delete(@out_file)
      else
        File.open(@out_file, "ab") do |f|
          f.truncate(state.size)
          pos = state.size - state.tail_size
          f.seek(pos)
          f.write(state.tail.bytes.to_a[pos .. -1])
        end
      end
    end


    def run_streamer(timeout)
      File.delete(@out_file) if File.zero?(@out_file)

      puts
      begin
        stdin, stdout, pid = PTY.spawn(@cmd)
      end
      timeout_at = Time.now + timeout
      loop do
        begin
          data = stdin.read_nonblock(1024)
          print data
          timeout_at = Time.now + timeout
        rescue Errno::EAGAIN
          sleep 0.1
        rescue EOFError
          $stderr.puts "\nprocess EOF?"
          Process.wait(pid)
          raise
        rescue Errno::EIO
          Process.wait(pid)
          if $?.exitstatus != 0
            raise LikelyError
          else
            $stderr.puts "Child exited. All done?"
            return true
          end
        end
        if Time.now >= timeout_at
          Process.kill('HUP', pid)
          Process.wait(pid)
          # FIXME: magic sleep
          # Without this magic delay the file size is wrong, even though we
          # have waited for the child. This makes no sense.
          sleep 0.1 # FIXME
          return false
        end
      end
    end
  end
end


def usage
  $stderr.puts <<EOF
#{APP_NAME} #{APP_VERSION} - Download video from ruutu.fi and jimtv.fi
Copyright (C) 2013  Tommi Saviranta <wnd@iki.fi>
http://wnd.katei.fi/ruutu-dl/

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

Usage: #{APP_NAME} [OPTIONS] PROGRAM_WEB_URI [OUTPUT_FILE]

Where OPTIONS is
  --no-retry    Do not retry downloading if streaming fails.
  --retry=N     Retry downloading at most N times (default 100).
EOF
end



begin
  RuutuDownloader.prerequisites?
rescue RuntimeError => e
  $stderr.puts "#{e} required"
  exit 1
end


get_opts = GetoptLong.new(
  [ '--help',     '-h', GetoptLong::NO_ARGUMENT ],
  [ '--no-retry', '-n', GetoptLong::NO_ARGUMENT ],
  [ '--out',      '-o', GetoptLong::REQUIRED_ARGUMENT ],
  [ '--retry',    '-r', GetoptLong::OPTIONAL_ARGUMENT ]
)

opts = {:retry_count => 100}
begin
  get_opts.each do |opt, arg|
    case opt
    when '--retry'
      opts[:retry_count] = arg == '' ? 100 : arg.to_i
    when '--no-retry'
      opts[:retry_count] = 0
    when '--out'
      opts[:out_file] = arg
    else
      usage
      exit 1
    end
  end
rescue GetoptLong::InvalidOption
  usage
  exit 1
end

if ![1, 2].include?(ARGV.count)
  usage
  exit 1
end
opts[:out_file] ||= ARGV[1]

RuutuDownloader.get(ARGV[0], opts)