import urllib2
import urlparse
import httplib
import socket
import time

import errors

from client import __ver__

class Request(urllib2.Request):
	def __init__(self, url, data=None, headers={},
		origin_req_host=None, unverifiable=False, head = False):
		urllib2.Request.__init__(self, url, data, headers, origin_req_host, unverifiable)
		
		self.add_header('User-Agent', 'MwClient-' + __ver__)
		self.head = head
	def get_method(self):
		if self.head: return 'HEAD'
		return urllib2.Request.get_method(self)

		
class CookieJar(dict):
	def __init__(self):
		dict.__init__(self, ())
	def extract_cookies(self, response):
		for cookie in response.msg.getallmatchingheaders('Set-Cookie'):
			self.parse_cookie(cookie.strip())
		if response.getheader('set-cookie2', None):
			# ...
			raise RuntimeError, 'Set-Cookie2', value
	def parse_cookie(self, cookie):
		if not cookie: return
		value, attrs = cookie.split(': ', 1)[1].split(';', 1)
		i = value.strip().split('=')
		if len(i) == 1 and i[0] in self: 
			del self[i[0]]
		else:
			self[i[0]] = i[1]
		
	def get_cookie_header(self):
		return '; '.join(('%s=%s' % i for i in self.iteritems()))
	def __iter__(self):
		for k, v in self.iteritems():
			yield Cookie(k, v)
		
class Cookie(object):
	def __init__(self, name, value):
		self.name = name
		self.value = value
		
class HTTPPersistentConnection(object):
	def __init__(self, host, pool = None):
		self.cookies = {}
		self.pool = pool
		if pool: self.cookies = pool.cookies
		self._conn = httplib.HTTPConnection(host)
		self._conn.connect()
		self.last_request = time.time()
		
	def request(self, method, host, path, headers, data, stream_iter = None,
			raise_on_not_ok = True, auto_redirect = True):		
			
		# Dirty hack...
		if (time.time() - self.last_request) > 60:
			self._conn.close()
			self._conn.connect()			
		
		_headers = headers
		headers = {}
		if not data: data = ''
		headers['Connection'] = 'Keep-Alive'
		headers['User-Agent'] = 'MwClient/' + __ver__
		headers['Host'] = host
		if host in self.cookies: 
			headers['Cookie'] = self.cookies[host].get_cookie_header()
		if _headers: headers.update(_headers)
		
		try:
			self._conn.request(method, path, data, headers)
			if stream_iter:
				for s_data in stream_iter:
					self._conn.send(s_data)
			
			self.last_request = time.time()
			try:
				res = self._conn.getresponse()
			except httplib.BadStatusLine:
				self._conn.close()
				self._conn.connect()
				self._conn.request(method, path, data, headers)
				res = self._conn.getresponse()
		except socket.error, e:
			self._conn.close()
			raise errors.HTTPError, e
		except Exception, e:
			raise errors.HTTPError, e
				
		if not host in self.cookies: self.cookies[host] = CookieJar()
		self.cookies[host].extract_cookies(res)
		
		if res.status >= 300 and res.status <= 399 and auto_redirect:
			res.read()
			
			location = urlparse.urlparse(res.getheader('Location'))
			if res.status in (302, 303):
				if 'Content-Type' in headers:
					del headers['Content-Type']
				if 'Content-Length' in headers:
					del headers['Content-Length']
				method = 'GET'
				data = ''
			path = location[2]
			if location[4]: path = path + '?' + location[4]
			
			if self.pool is None:
				if location[1] != host: 
					raise errors.HTTPError, ('Redirecting to different hosts not supported', 
						res.getheader('Location'))
				return self.request(method, host, path, headers, data)
			else:
				return self.pool.request(method, location[1], path, 
					headers, data, stream_iter, raise_on_not_ok, auto_redirect)
			
		if res.status != 200 and raise_on_not_ok:
			try:
				raise errors.HTTPStatusError, (res.status, res)
			finally:
				res.close()
			
		return res
		
	def get(self, host, path, headers = None):
		return self.request('GET', host, path, headers, None)
	def post(self, host, path, headers = None, data = None, stream_iter = None):
		return self.request('POST', host, path, headers, data, stream_iter)
	def head(self, host, path, headers = None, auto_redirect = False):
		res = self.request('HEAD', host, path, headers, 
			data = None, raise_on_not_ok = False,
			auto_redirect = auto_redirect)
		res.read()
		return res.status, res.getheaders()
		
	def close(self):
		self._conn.close()

class HTTPConnection(HTTPPersistentConnection):
	def request(self, method, host, path, headers, data, stream_iter = None,
			raise_on_not_ok = True, auto_redirect = True):
		if not headers: headers = {}
		headers['Connection'] = 'Close'
		res = HTTPPersistentConnection.request(self, method, host, path, headers, data, 
			stream_iter, raise_on_not_ok, auto_redirect)
		return res

	
class HTTPPool(list):
	def __init__(self):
		list.__init__(self)
		self.cookies = {}
	def find_connection(self, host):
		for hosts, conn in self:
			if host in hosts: return conn
				
		for hosts, conn in self:
			status, headers = conn.head(host, '/')
			if status == 200:
				hosts.append(host)
				return conn
			if status >= 300 and status <= 399:
				headers = dict(headers)
				location = urlparse.urlparse(headers.get('location', ''))
				if location[1] == host:
					hosts.append(host)
					return conn
		conn = HTTPPersistentConnection(host, self)
		self.append(([host], conn))
		return conn
	def get(self, host, path, headers = None):
		return self.find_connection(host).get(host, 
			path, headers)
	def post(self, host, path, headers = None, data = None, stream_iter = None):
		return self.find_connection(host).post(host, 
			path, headers, data, stream_iter)
	def head(self, host, path, headers = None, auto_redirect = False):
		return self.find_connection(host).head(host, 
			path, headers, auto_redirect)
	def request(self, method, host, path, headers, data, stream_iter,
			raise_on_not_ok, auto_redirect):
		return self.find_connection(host).request(method, host, path,
			headers, data, stream_iter, raise_on_not_ok, auto_redirect)
	def close(self):
		for hosts, conn in self:
			conn.close()