# -*- coding: utf8 -*-

info = {
	'author': 'Jon Bergli Heier',
	'title': 'Traffic data',
	'description': 'Traffic data from Statens Vegvesen',
}

import urllib2, urllib
from xml.etree import ElementTree as ET

counties = {
	'Akershus': 2,
	'Aust-Agder': 9,
	'Buskerud': 6,
	'Finnmark': 20,
	'Hedmark': 4,
	'Hordaland': 12,
	u'Møre og Romsdal': 15,
	u'Nord-Trøndelag': 17,
	'Nordland': 18,
	'Oppland': 5,
	'Oslo': 3,
	'Rogaland': 11,
	'Sogn og Fjordane': 14,
	u'Sør-Trøndelag': 16,
	'Telemark': 8,
	'Troms': 19,
	'Vest-Agder': 10,
	'Vestfold': 7,
	u'Østfold': 1,
}

class InvalidCountyException(Exception): pass
class AmbiguousCountyException(Exception): pass

class Module:
	def __init__(self, bot):
		self.irc = bot
		if self.irc:
			self.irc.register_keyword('!traffic', self)

	def traffic(self, args):
		if len(args) == 2:
			rn, c = args
		elif args[0][-1].isdigit():
			rn = args[0]
			c = None
		else:
			rn = None
			c = args[0]
		if c:
			c = c.decode('utf8').lower()
			matches = [s for s in counties if c in s.lower()]
			if not len(matches):
				raise InvalidCountyException()
			elif len(matches) > 1:
				raise AmbiguousCountyException(matches)
			else:
				c = matches[0]

		if rn:
			if rn.lower().startswith('e'):
				rt = 'Ev'
			elif rn.lower().startswith('r'):
				rt = 'Rv'
			else:
				rt = 'Alle'
		else:
			rt = None

		while rn and not rn[0].isdigit():
			rn = rn[1:]

		if rt and rn:
			url = 'http://www.vegvesen.no/trafikk/xml/search.xml?searchFocus.roadTypes=%s&searchFocus.roadNumber=%s' % (rt, rn)
		elif c:
			url = 'http://www.vegvesen.no/trafikk/xml/search.xml?searchFocus.counties=%d' % counties[c]
		u = urllib2.urlopen(url)

		xmldoc = ET.parse(u)
		xml = xmldoc.getroot()

		results = []
		for message in xml.findall('result-array/result/messages/message'):
			if (rt and c and c.lower() in (x.text.lower() for x in message.findall('actualCounties/string'))) or not c or not rt:
				if len(results) == 5:
					if rt and rn:
						results.append('More results: http://www.vegvesen.no/Trafikkinformasjon/Reiseinformasjon/Trafikkmeldinger?%s' % urllib.urlencode({
							'type': 'veg',
							'vegkategorier': rt,
							'vegnummer': '6',
							}))
					else:
						results.append('More results: http://www.vegvesen.no/Trafikkinformasjon/Reiseinformasjon/Trafikkmeldinger?%s' % urllib.urlencode({
							'type': 'fylke',
							'fylke': counties[c],
							}))
					break
				results.append('\002%s\002 - %s' % (message.find('heading').text, message.find('ingress').text))
		return results

	def keyword(self, nick, channel, kw, msg):
		if not len(msg):
			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], 'Usage: !traffic [ROADNO] [COUNTY]')
			return

		try:
			for line in self.traffic(msg.split()):
				self.irc.msg(nick.split('!')[0], line.encode('utf8'))
		except InvalidCountyException:
			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], 'Invalid county.')
		except AmbiguousCountyException as e:
			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], ('Ambiguous county: %s' % (', '.join(e.message))).encode('utf8'))

if __name__ == '__main__':
	import sys
	m = Module(None)
	print '\n'.join(m.traffic(sys.argv[1:]))