#! /usr/bin/env ruby

# This code is in the public domain
# Written by Robbert Haarman

require 'net/http'
require 'uri'

## Defaults
$loglevel = 10

def log priority, message
	$stderr.puts message if priority < $loglevel
end

def human_readable n
	if n < 0x400
		"#{n} bytes"
	elsif n < 0x100000
		"#{(n * 10 / 0x400).to_i / 10.0} KB"
	elsif n < 0x40000000
		"#{(n * 10 / 0x100000).to_i / 10.0} MB"
	else
		"#{(n * 10 / 0x40000000).to_i / 10.0} GB"
	end
end

def log_fetched n, total = nil
	if total
		$stderr.print "\r\x1b[KFetched #{human_readable n} of " +
			"#{human_readable total}\r"
	else
		$stderr.print "\r\x1b[KFetched #{human_readable n}\r"
	end
end

## Parse command line
uri = ARGV[0]

## Get video id
match = uri.match(/watch\?v=([^&]*)/)

unless match
	log 0, "Could not extract video id from URI"
	exit 1
end

video_id = match[1]
log 4, "Video ID: #{video_id}"

## Get video page
log 1, "Fetching video page: #{uri}"
page = Net::HTTP::get_response URI.parse(uri)

unless page.instance_of? Net::HTTPOK
	log 0, "Error fetching #{uri}"
	exit 1
end

## Extract params
match = page.body.match(/player2\.swf\?([^"]*)/)

unless match
	log 0, "Could not extract video parameters"
	exit 1
end

params = match[1]
log 5, "Video Parameters: #{params.inspect}"

## Compute video URI
video_uri = "http://www.youtube.com/get_video?#{params}"

## Follow redirects
while true
	log 6, "Fetching headers for #{video_uri}"

	uri = URI.parse video_uri
	resp = Net::HTTP::start(uri.host, uri.port) do |http|
		http.request_head "#{uri.path}?#{uri.query}"
	end

	case resp
	when Net::HTTPSuccess
		break
	when Net::HTTPRedirection
		video_uri = resp['location']
		log 5, "Redirected to #{video_uri}"
	else
		log 0, "Unexpected response: #{resp}"
		exit 1
	end
end

log 4, "Real video URI: #{video_uri}"

## Compute filename
filename = "#{video_id}.flv"
log 2, "Saving video to #{filename}"

## Save video to file
File.open(filename, 'wb') do |file|
	Net::HTTP::get_response(URI.parse(video_uri)) do |resp|
		begin
			size = resp['content-length'].to_i
		rescue
			size = nil
		end

		fetched = 0
		msgtime = 0.0
		resp.read_body do |chunk|
			file.print chunk
			fetched = fetched + chunk.size
			if $loglevel > 0 && Time.now.to_f - msgtime >= 1.0
				log_fetched fetched, size
				msgtime = Time.now.to_f
			end
		end

		$stderr.print "\r\x1b[KDone" if $loglevel > 0
	end
end