# Website
#
# Douglas Thrift
#
# $Id$

from BeautifulSoup import MinimalSoup
import copy
import mechanize
import re

# XXX: hack to make Wells Fargo http-equiv redirects actually work
mechanize._http.AbstractHeadParser.head_elems = tuple(list(mechanize._http.AbstractHeadParser.head_elems) + ['body'])
mechanize._http.HTTPRedirectHandler.max_redirections = 20

class DuckSoup(MinimalSoup):
	MARKUP_MASSAGE = copy.copy(MinimalSoup.MARKUP_MASSAGE)

	MARKUP_MASSAGE.append((re.compile(r'''</?[a-z]+['"]\+['"][a-z]+'''), lambda match: ''))

class Factory(mechanize.DefaultFactory):
	def __init__(self, form_parser):
		mechanize.DefaultFactory.__init__(self)

		self._forms_factory = mechanize.FormsFactory(form_parser_class = form_parser)

class Website(object):
	Soup = DuckSoup

	def __init__(self, debug, *args, **kwargs):
		self.browser = mechanize.Browser(*args, **kwargs)
		self.browser.addheaders = [('User-agent', 'Mozilla/5.0 (Windows; U; Windows NT 5.1; en-US) AppleWebKit/525.19 (KHTML, like Gecko) Chrome/0.4.154.25 Safari/525.19')]

		if debug:
			self.browser.set_debug_http(True)

		self.browser.set_handle_equiv(True)
		self.browser.set_handle_redirect(True)
		self.browser.set_handle_refresh(True)
		self.browser.set_handle_robots(False)

		self._count = 0

	def _back(self):
		if self._count:
			result = self.browser.back(self._count)

			self._count = 0

			return result

	def _follow_link(self, *args, **kwargs):
		result = self.browser.follow_link(*args, **kwargs)

		self._count += 1

		return result

	def _submit(self, *args, **kwargs):
		result = self.browser.submit(*args, **kwargs)

		self._count += 1

		return result