#!/usr/bin/env python
# -*- coding: utf-8 -*-

__author__ = "Yoan Blanc <yoan at dosimple dot ch>"
__revision__ = "20080113"
__license__ = "MIT"

# Twisted Imports
from twisted.words.protocols.jabber import client, jid
from twisted.words.xish import domish, xmlstream
from twisted.internet import reactor
import logging

class Roster(object):
	def __init__(self):
		self.buddies = {}
	
	def addBuddy(self, jid, buddy):
		logging.info("%s %s" % (jid, repr(buddy)))
		if not jid in self.buddies:
			self.buddies[jid] = {"name": jid}
			if "name" in buddy:
				self.buddies[jid]["name"] = buddy["name"]
			self.buddies[jid]["connected"] = False
		
		return self.buddies[jid]
	
	def removeBuddy(self, jid):
		if jid in self.buddies:
			return self.buddies.pop(jid)
		else:
			return None
	
	def hasBuddy(self, jid):
		if jid in self.buddies:
			return True
		else:
			logging.warn("%s not in roster" % jid)
			return False
	
	def setOnline(self, jid, online):
		if self.hasBuddy(jid):
			self.buddies[jid]["connected"] = online
			return self.buddies[jid]
		
	def setStatus(self, jid, status, show):
		if self.hasBuddy(jid):
			self.buddies[jid]["status"] = status
			self.buddies[jid]["show"] = show
			return self.buddies[jid]
	
	def isOnline(self, jid):
		logging.info("isOnline " + jid)
		return (self.hasBuddy(jid) \
			and self.buddies[jid]["connected"] \
			and self.buddies[jid]["status"] != "dnd")
	
	def __str__(self):
		return u"<Roster [%s]" % u",".join([k for k, v in self.buddies.iteritems()])

class JabberBot(object):
	"""Basic jabber bot"""
	
	def __init__(self, jid, password, reactor=reactor, port=5222, resource="JabberBot"):
		self.jabberid = jid
		self.password = password
		self.servername = jid[jid.find('@')+1:]
		self.port = port
		self.resource = resource
		self.roster = Roster()
		
		# internal values
		self._jid = None
		self._factory = None
		self._reactor = reactor
		self._resource = None
		self._xmlstream = None
		self.tryandregister = 1
	
	def run(self):
		self.__initFactory()
	
	def __repr__(self):
		return "<%s (%s)>" % (type(self).__name__, self.jabberid)
	
	def __initFactory(self):
		self._jid = jid.JID("%s/%s" % (self.jabberid, self.resource))
		self._factory = client.basicClientFactory(self._jid, self.password)
		
		self._factory.addBootstrap('//event/stream/authd', self._authd)
		
		self._reactor.connectTCP(self.servername, self.port, self._factory)
		self._reactor.run() 
	
	def _authd(self, xmlstream):
		if xmlstream:
			self._xmlstream = xmlstream
			
			# set it as online
			self._presence = domish.Element(('jabber:client', 'presence'))
			self._presence.addElement('status').addContent('Online')
			self._xmlstream.send(self._presence)

			self.__initOnline()
			self.__initRoster()
	
	def __initOnline(self):
		self._xmlstream.addObserver('/message', self._gotMessage)
		self._xmlstream.addObserver('/presence', self._gotPresence)
		self._xmlstream.addObserver('/iq/query[@xmlns="jabber:iq:roster"]', self._gotRosterIq)
		#self._xmlstream.addObserver('/iq/query[@xmlns="jabber:iq:version"]', self._gotVersionIq)
		#self._xmlstream.addObserver('/iq', self._gotIq)
		# debug
		#self._xmlstream.addObserver('/*', self._gotSomething)
	
	def __initRoster(self):
		roster = domish.Element((None, "iq"))
		roster.attributes["type"] = "get"
		roster.attributes["id"] = "getRoster"
		rosterQuery = roster.addElement('query')
		rosterQuery.attributes['xmlns'] = "jabber:iq:roster"
		
		self._xmlstream.send(roster)
	
	def __parseJID(self, jid):
		"""
		>>> self.__parseJID("foo@example.com/Bar")
		('foo@example.com', 'Bar')
		"""
		
		if not jid:
			return None, None
		
		try:
			jid, resource = jid.split("/", 1)
		except:
			resource = None
		
		return (jid, resource)
	
	def sendPresence(self, to, value):
		"""
		Send the presence to.
		"""
		
		if self._xmlstream:
			presence = domish.Element(('jabber:client', 'presence'))
			presence.attributes["to"] = jid.JID(to).full()
			presence.attributes["from"] = self._jid.full()
			presence.attributes["type"] = value
			
			self._xmlstream.send(presence)
	
	def add(self, to):
		"""
		Add the user to the roster
		"""
		
		if self._xmlstream:
			iq = domish.Element((None, "iq"))
			iq.attributes["type"] = "set"
			iqQuery = iq.addElement('query')
			iqQuery.attributes['xmlns'] = "jabber:iq:roster"
			iqItem = iqQuery.addElement('item')
			iqItem.attributes["jid"] = to
			iqItem.attributes["name"] = to.replace("@", " @ ")
			# can add it in some groups here.
			
			self._xmlstream.send(iq)
		
	def remove(self, to):
		"""
		Remove the user from the roster.
		"""
		
		if self._xmlstream:
			iq = domish.Element((None, "iq"))
			iq.attributes["type"] = "set"
			iqQuery = iq.addElement('query')
			iqQuery.attributes['xmlns'] = "jabber:iq:roster"
			iqItem = iqQuery.addElement('item')
			iqItem.attributes["jid"] = to
			iqItem.attributes["subscription"] = "remove"
			
			self._xmlstream.send(iq)
	
	def _gotSomething(self, el):
		"""
		Debug function
		"""
		logging.info("[something] %s" % el.toXml())
	
	def _gotPresence(self, el):
		"""
		Auto-subscribe to my new friends
		"""
		
		jid, resource = self.__parseJID(el["from"])
		
		show = None
		status = None
		priority = None
		
		for element in el.elements():
			if element.name == "status":
				status = element.__str__()
			elif element.name == "show":
				show = element.__str__()
			elif element.name == "priority":
				priority = element.__str__()
		
		if "type" in el.attributes:
			presence = el.attributes["type"]
			
			if presence == "subscribe":
				# asking for being a friend.
				self.sendPresence(jid, "subscribe")
			elif presence == "subscribed":
				self.sendPresence(jid, "subscribed")
			elif presence == "unsubscribe":
				self.sendPresence(jid, "unsubscribe")
			elif presence == "unsubscribed":
				# denied the subscribe asking
				self.sendPresence(jid, "unsubscribed")
			elif presence == "unavailable":
				# just logged out.
				self.roster.setOnline(jid, False)
			else:
				logging.warn("[presence] is %s" % el.attributes["type"])
		else:
			self.roster.setOnline(jid, True)
			self.roster.setStatus(jid, status, show)
			
			self.gotPresence(jid, resource, show, status)			
	
	def _gotMessage(self, el):
		"""
		Called when a message is received
		"""
		
		jid, resource = self.__parseJID(el["from"])
		
		body = None
		composing = False
		for e in el.elements():
			if e.name == "body":
				body = e.__str__()
			elif e.name == "composing":
				composing = True
		
		if el["type"] == "error":
			self.gotError(jid, resource, body)
		else:
			self.gotMessage(jid, resource, body, composing)
	
	def _gotRosterIq(self, el):
		
		action = el["type"]
		
		for e in el.elements():
			for item in e.elements():
				if "jid" in item.attributes and "subscription" in item.attributes:
					subscription = item.attributes["subscription"]
					from_jid = item.attributes["jid"]
					
					self.roster.addBuddy(from_jid, item.attributes)
					
					if subscription == "both":
						pass
					elif subscription == "from":
						self.add(from_jid)
					elif subscription == "to":	
						self.remove(from_jid)	
					elif subscription == "none":
						if "ask" in item.attributes:
							# inscription is pending a yes will get us a "to"
							# and a no to the unsubscribed presence info.
							pass
						else:
							self.remove(from_jid)
				else:
					logging.error("[Unknown] %s" % repr(item.attributes))
	
	def _gotIq(self, el):
		logging.info("[iq] %s" % el.toXml())
	
	def gotMessage(self, jid, resource, body, composing):
		raise NotImplementedError, u"Please use gotMessage to treat the incoming messages"
	
	def gotPresence(self, jid, resource, show, satus):
		raise NotImplementedError, u"Please use gotPresence to treat the incoming presences"
	
	def gotError(self, jid, ressource, body):
		pass
	
	def sendMessage(self, to, body):
		message = domish.Element(('jabber:client','message'))
		message["to"] = jid.JID(to).full()
		message["from"] = self._jid.full()
		message["type"] = "chat"
		message.addElement("body", "jabber:client", body)
		
		self._xmlstream.send(message)