# -*- coding: utf-8 -*-

"""
A simple library for managing the availability of ports.
"""

from __future__ import print_function

import time
import socket
import datetime
import argparse
import sys
import functools
import itertools

from jaraco import timing
from py26compat import total_seconds


def client_host(server_host):
	"""Return the host on which a client can connect to the given listener."""
	if server_host == '0.0.0.0':
		# 0.0.0.0 is INADDR_ANY, which should answer on localhost.
		return '127.0.0.1'
	if server_host in ('::', '::0', '::0.0.0.0'):
		# :: is IN6ADDR_ANY, which should answer on localhost.
		# ::0 and ::0.0.0.0 are non-canonical but common
		# ways to write IN6ADDR_ANY.
		return '::1'
	return server_host


def _getaddrinfo(host, port, *args, **kwargs):
	"""
	Provide a fallback when getaddrinfo fails.

	TODO: why is this needed?
	"""
	try:
		return socket.getaddrinfo(host, port, *args, **kwargs)
	except socket.gaierror:
		host = client_host(host)
		if ':' in host:
			family = socket.AF_INET6
			addr = host, port, 0, 0
		else:
			family = socket.AF_INET
			addr = host, port
		item = family, socket.SOCK_STREAM, 0, "", addr
		return [item]


def _check_port(host, port, timeout=1.0):
	"""
	Raise an error if the given port is not free on the given host
	(i.e. all attempts to connect fail within the timeout).

	>>> free_port = find_available_local_port()
	>>> _check_port('localhost', free_port)
	>>> _check_port('127.0.0.1', free_port)
	>>> _check_port('::1', free_port)
	"""
	info = _getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
	_timeout_connect = functools.partial(_check_free_info, timeout=timeout)
	list(itertools.starmap(_timeout_connect, info))


def _check_free_info(af, socktype, proto, canonname, sa, timeout):
	s = socket.socket(af, socktype, proto)
	try:
		# fail fast with a small timeout
		s.settimeout(timeout)
		s.connect(sa)
	except socket.error:
		return
	finally:
		s.close()

	# the connect succeeded, so the port isn't free
	port, host = sa[:2]
	tmpl = "Port {port} is in use on {host}."
	raise IOError(tmpl.format(**locals()))


class Timeout(IOError):
	pass


def free(host, port, timeout=float('Inf')):
	"""
	Wait for the specified port to become free (dropping or rejecting
	requests). Return when the port is free or raise a Timeout if timeout has
	elapsed.

	Timeout may be specified in seconds or as a timedelta.
	If timeout is None or ∞, the routine will run indefinitely.

	>>> free('localhost', find_available_local_port())
	"""
	if not host:
		raise ValueError("Host values of '' or None are not allowed.")

	if isinstance(timeout, datetime.timedelta):
		timeout = timeout.total_seconds()

	if timeout is None:
		# treat None as infinite timeout
		timeout = float('Inf')

	watch = timing.Stopwatch()

	while total_seconds(watch.split()) < timeout:
		try:
			# Expect a free port, so use a small timeout
			_check_port(host, port, timeout=0.1)
			return
		except IOError:
			# Politely wait.
			time.sleep(0.1)

	raise Timeout("Port {port} not free on {host}.".format(**locals()))
wait_for_free_port = free


def occupied(host, port, timeout=float('Inf')):
	"""
	Wait for the specified port to become occupied (accepting requests).
	Return when the port is occupied or raise a Timeout if timeout has
	elapsed.

	Timeout may be specified in seconds or as a timedelta.
	If timeout is None or ∞, the routine will run indefinitely.

	>>> occupied('localhost', find_available_local_port(), .1) # doctest: +IGNORE_EXCEPTION_DETAIL
	Traceback (most recent call last):
	    ...
	Timeout: Port ... not bound on localhost.
	"""
	if not host:
		raise ValueError("Host values of '' or None are not allowed.")

	if isinstance(timeout, datetime.timedelta):
		timeout = timeout.total_seconds()

	if timeout is None:
		# treat None as infinite timeout
		timeout = float('Inf')

	watch = timing.Stopwatch()

	while total_seconds(watch.split()) < timeout:
		try:
			_check_port(host, port, timeout=.5)
			# Politely wait
			time.sleep(0.1)
		except IOError:
			# port is occupied
			return

	raise Timeout("Port {port} not bound on {host}.".format(**locals()))
wait_for_occupied_port = occupied


def find_available_local_port():
	"""
	Find a free port on localhost.

	>>> 0 < find_available_local_port() < 65536
	True
	"""
	sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
	addr = '', 0
	sock.bind(addr)
	addr, port = sock.getsockname()[:2]
	sock.close()
	return port


class HostPort(str):
	"""
	A simple representation of a host/port pair as a string

	>>> hp = HostPort('localhost:32768')

	>>> hp.host
	'localhost'

	>>> hp.port
	32768

	>>> len(hp)
	15
	"""

	@property
	def host(self):
		host, sep, port = self.partition(':')
		return host

	@property
	def port(self):
		host, sep, port = self.partition(':')
		return int(port)


def _main():
	parser = argparse.ArgumentParser()
	global_lookup = lambda key: globals()[key]
	parser.add_argument('target', metavar='host:port', type=HostPort)
	parser.add_argument('func', metavar='state', type=global_lookup)
	parser.add_argument('-t', '--timeout', default=None, type=float)
	args = parser.parse_args()
	try:
		args.func(args.target.host, args.target.port, timeout=args.timeout)
	except Timeout as timeout:
		print(timeout, file=sys.stderr)
		raise SystemExit(1)


if __name__ == '__main__':
	_main()

@@ -225,16 +225,34 @@
 tertools
+%0Aimport contextlib
 %0A%0Afrom j
@@ -1829,22 +1829,54 @@
  proto)%0A
+%0A%09with contextlib.closing(s):%0A%09
 %09try:%0A
+%09
 %09%09# fail
@@ -1902,16 +1902,17 @@
 timeout%0A
+%09
 %09%09s.sett
@@ -1929,16 +1929,17 @@
 eout)%0A%09%09
+%09
 s.connec
@@ -1944,16 +1944,17 @@
 ect(sa)%0A
+%09
 %09except 
@@ -1973,36 +1973,15 @@
 :%0A%09%09
+%09
 return
-%0A%09finally:%0A%09%09s.close()
 %0A%0A%09#

Use contextlib.closing to unconditionally close the socket.