import markdown

class UrlRebaseExtension(markdown.Extension):
	def __init__ (self, configs):
		# set extension defaults
		self.config = {
			'rebase' : [None, 'Callable to rebase elementes with custom logic.'],
			'base_url' : ['/', 'String to append to beginning or URL.'],
		}

		# Override defaults with user settings
		for key, value in configs :
			self.setConfig(key, value)		

	def extendMarkdown(self, md, md_globals):
		""" Add pieces to Markdown. """
		md.registerExtension(self)
		md.treeprocessors.add("urlrebase", UrlRebaseTreeprocessor(self.getConfigs()), "_end")
		
class UrlRebaseTreeprocessor(markdown.treeprocessors.Treeprocessor):
	def __init__ (self, config):
		self.config = config
	
	def run(self, root):
		rebase = self.rebase_element
		if "rebase" in self.config:
			rebase = self.config["rebase"]
			
		for a in root.findall(".//a"):
			rebase(self, a, 'href')
		for i in root.findall(".//img"):
			rebase(self, i, 'src')
		return root

	def rebase_element(self, also_self, e, attr):
		uri = e.get(attr, '')
		if '://' in uri or uri.startswith('/'):
			return
		e.set(attr, self.rebase(uri))

	def rebase(self, base, uri):
		if base[0] and not base.endswith('/'):
			base = base + '/'
		return base + uri

def makeExtension(configs=None) :
	return UrlRebaseExtension(configs=configs)