require 'socket'
require 'yaml'
require 'thread'
require 'cipher/des'
require 'net/vnc/version'

module Net
	#
	# The VNC class provides for simple rfb-protocol based control of
	# a VNC server. This can be used, eg, to automate applications.
	#
	# Sample usage:
	#
	#	  # launch xclock on localhost. note that there is an xterm in the top-left
	#   Net::VNC.open 'localhost:0', :shared => true, :password => 'mypass' do |vnc|
	#     vnc.pointer_move 10, 10
	#     vnc.type 'xclock'
	#     vnc.key_press :return
	#   end
	#
	# = TODO
	#
	# * The server read loop seems a bit iffy. Not sure how best to do it.
	# * Should probably be changed to be more of a lower-level protocol wrapping thing, with the
	#   actual VNCClient sitting on top of that. all it should do is read/write the packets over
	#   the socket.
	#
	class VNC
		class PointerState
			attr_reader :x, :y, :button

			def initialize vnc
				@x = @y = @button = 0
				@vnc = vnc
			end

			# could have the same for x=, and y=
			def button= button
				@button = button
				refresh
			end

			def update x, y, button=@button
				@x, @y, @button = x, y, button
				refresh
			end

			def refresh
				packet = 0.chr * 6
				packet[0] = 5.chr
				packet[1] = button.chr
				packet[2, 2] = [x].pack 'n'
				packet[4, 2] = [y].pack 'n'
				@vnc.socket.write packet
			end
		end

		BASE_PORT = 5900
		CHALLENGE_SIZE = 16
		DEFAULT_OPTIONS = {
			:shared => false,
			:wait => 0.1
		}

		keys_file = File.dirname(__FILE__) + '/../../data/keys.yaml'
		KEY_MAP = YAML.load_file(keys_file).inject({}) { |h, (k, v)| h.update k.to_sym => v }
		def KEY_MAP.[] key
			super or raise ArgumentError.new('Invalid key name - %s' % key)
		end

		attr_reader :server, :display, :options, :socket, :pointer

		def initialize display=':0', options={}
			@server = 'localhost'
			if display =~ /^(.*)(:\d+)$/
				@server, display = $1, $2
			end
			@display = display[1..-1].to_i
			@options = DEFAULT_OPTIONS.merge options
			@clipboard = nil
			@pointer = PointerState.new self
			@mutex = Mutex.new
			connect
			@packet_reading_state = nil
			@packet_reading_thread = Thread.new { packet_reading_thread }
		end

		def self.open display=':0', options={}
			vnc = new display, options
			if block_given?
				begin
					yield vnc
				ensure
					vnc.close
				end
			else
				vnc
			end
		end

		def port
			BASE_PORT + @display
		end

		def connect
			@socket = TCPSocket.open server, port
			unless socket.read(12) =~ /^RFB (\d{3}.\d{3})\n$/
				raise 'invalid server response'
			end
			@server_version = $1
			socket.write "RFB 003.003\n"
			data = socket.read(4)
			auth = data.to_s.unpack('N')[0]
			case auth
			when 0, nil
				raise 'connection failed'
			when 1
				# ok...
			when 2
				password = @options[:password] or raise 'Need to authenticate but no password given'
				challenge = socket.read CHALLENGE_SIZE
				response = Cipher::DES.encrypt password, challenge
				socket.write response
				ok = socket.read(4).to_s.unpack('N')[0]
				raise 'Unable to authenticate - %p' % ok unless ok == 0
			else
				raise 'Unknown authentication scheme - %d' % auth
			end

			# ClientInitialisation
			socket.write((options[:shared] ? 1 : 0).chr)

			# ServerInitialisation
			# TODO: parse this.
			socket.read(20)
			data = socket.read(4)
			# read this many bytes in chunks of 20
			size = data.to_s.unpack('N')[0]
			while size > 0
				len = [20, size].min
				# this is the hostname, and other stuff i think...
				socket.read(len)
				size -= len
			end
		end

		# this types +text+ on the server
		def type text, options={}
			packet = 0.chr * 8
			packet[0] = 4.chr
			text.split(//).each do |char|
				packet[7] = char[0]
				packet[1] = 1.chr
				socket.write packet
				packet[1] = 0.chr
				socket.write packet
			end
			wait options
		end

		# this takes an array of keys, and successively holds each down then lifts them up in
		# reverse order.
		# FIXME: should wait. can't recurse in that case.
		def key_press(*args)
			options = Hash === args.last ? args.pop : {}
			keys = args
			raise ArgumentError, 'Must have at least one key argument' if keys.empty?
			begin
				key_down keys.first
				if keys.length == 1
					yield if block_given?
				else
					key_press(*(keys[1..-1] + [options]))
				end
			ensure
				key_up keys.first
			end
		end

		def get_key_code which
			if String === which
				if which.length != 1
					raise ArgumentError, 'can only get key_code of single character strings'
				end
				which[0]
			else
				KEY_MAP[which]
			end
		end
		private :get_key_code

		def key_down which, options={}
			packet = 0.chr * 8
			packet[0] = 4.chr
			key_code = get_key_code which
			packet[4, 4] = [key_code].pack('N')
			packet[1] = 1.chr
			socket.write packet
			wait options
		end

		def key_up which, options={}
			packet = 0.chr * 8
			packet[0] = 4.chr
			key_code = get_key_code which
			packet[4, 4] = [key_code].pack('N')
			packet[1] = 0.chr
			socket.write packet
			wait options
		end

		def pointer_move x, y, options={}
			# options[:relative]
			pointer.update x, y
			wait options
		end

		BUTTON_MAP = {
			:left => 0
		}

		def button_press button=:left, options={}
			begin
				button_down button, options
				yield if block_given?
			ensure
				button_up button, options
			end
		end

		def button_down which=:left, options={}
			button = BUTTON_MAP[which] || which
			raise ArgumentError, 'Invalid button - %p' % which unless (0..2) === button
			pointer.button |= 1 << button
			wait options
		end

		def button_up which=:left, options={}
			button = BUTTON_MAP[which] || which
			raise ArgumentError, 'Invalid button - %p' % which unless (0..2) === button
			pointer.button &= ~(1 << button)
			wait options
		end

		def wait options={}
			sleep options[:wait] || @options[:wait]
		end

		def close
			# destroy packet reading thread
			if @packet_reading_state == :loop
				@packet_reading_state = :stop
				while @packet_reading_state
					# do nothing
				end
			end
			socket.close
		end

		def clipboard
			if block_given?
				@clipboard = nil
				yield
				60.times do
					clipboard = @mutex.synchronize { @clipboard }
					return clipboard if clipboard
					sleep 0.5
				end
				warn 'clipboard still empty after 30s'
				nil
			else
				@mutex.synchronize { @clipboard }
			end
		end

		private

		def read_packet type
			case type
			when 3 # ServerCutText
				socket.read 3 # discard padding bytes
				len = socket.read(4).unpack('N')[0]
				@mutex.synchronize { @clipboard = socket.read len }
			else
				raise NotImplementedError, 'unhandled server packet type - %d' % type
			end
		end

		def packet_reading_thread
			@packet_reading_state = :loop
			loop do
				begin
					break if @packet_reading_state != :loop
					next unless IO.select [socket], nil, nil, 2
					type = socket.read(1)[0]
					read_packet type
				rescue
					warn "exception in packet_reading_thread: #{$!.class}:#{$!}"
					break
				end
			end
			@packet_reading_state = nil
		end
	end
end

