import sys
import beatbox
import xmltramp
import datetime

sf = beatbox._tPartnerNS
svc = beatbox.Client()

def nofilter(f):
	return True

def idFilter(f):
	return str(f[sf.name]) == "Id"

def fkFilter(f):
	return idFilter(f) or str(f[sf.type]) == "reference"

def idAndNameFilter(f):
	return idFilter(f) or str(f[sf.nameField]) == "true"

def fkOrNameFilter(f):
	return fkFilter(f) or idAndNameFilter(f)
		
class GvWriter:
	def __init__(self):
		self.refs = set()
		self.describes = dict()
		
	def start(self, primaryName, depth):
		self.primaryName = primaryName
		self.depth = depth
		print "digraph g {"
		print "graph [rankdir=LR overlap=scale];"
	
	def describe(self, name):
		k = name.lower()
		if k in self.describes:
			return self.describes[k]
		d = svc.describeSObjects(name)
		self.describes[k] = d
		return d
		
	def writeSObjectNode(self, color, obj, filter):
		print str(obj[sf.name]) + " [shape=Mrecord style=filled fillcolor=" + color + "  label=\"<f0> " + str(obj[sf.name]),
		fIdx = 1
		for field in obj[sf.fields:]:
			if filter(field):
				print " | <f" + str(fIdx) + "> " + str(field[sf.name]),
				fIdx = fIdx + 1
		print "\"];"

	def writeFkConnections(self, obj, filter):
		desc = self.describe(obj)
		fidx = 1;
		for f in desc[sf.fields:]:
			if filter(f):
				for refTo in f[sf.referenceTo:]:
					if str(refTo) in self.refs:
						print str(desc[sf.name]) + ":f" + str(fidx) + " -> " + str(refTo) + ":f1;"
				fidx = fidx + 1
						
	# returns a set of objects related to the current things in refs that we don't already know about
	def findNewRefs(self):
		found = set()
		for o in self.refs:
			d = self.describe(o)
			# find all the types we have a reference to
			for field in d[sf.fields:]:
				for foreignType in field[sf.referenceTo:]:
					if str(foreignType) not in self.refs:
						found.add(str(foreignType))

			# find all the types that have a reference to us, unless we're User (which connects to everything)
			if self.primaryName == "User" or str(d[sf.name]) != "User":
				for cr in d[sf.childRelationships:]:
					t = str(cr[sf.childSObject])
					if t not in self.refs:
						found.add(t)
		
		return found
			
	def generate(self):
		svc.serverUrl = "https://login.salesforce.com/services/Soap/u/18.0"
		svc.login(sys.argv[1], sys.argv[2])
		primary = self.describe(sys.argv[3])
	
		d = 1
		if len(sys.argv) == 5: d = int(sys.argv[4])
		self.start(str(primary[sf.name]), d)
		
		self.writeSObjectNode("coral", primary, nofilter)
		self.refs.add(self.primaryName)

		# spider the object mode from the primary, for the specified depth
		currentDepth = self.depth
		while (currentDepth > 0):
			newRefs = self.findNewRefs()
			for f in newRefs:
				self.refs.add(f)
			currentDepth = currentDepth -1
				
		# write out all the object nodes.
		for rt in self.refs:
			if rt != self.primaryName:
				desc = self.describe(rt)
				self.writeSObjectNode("cornflowerblue", desc, fkOrNameFilter)

		# write out all the connections from the primary
		self.writeFkConnections(self.primaryName, nofilter)
		
		# write out all the other connections
		for rt in self.refs:
			if rt != self.primaryName:
				self.writeFkConnections(rt, fkOrNameFilter)	# have to use the same filter as we defined the node with

		print "}"

if __name__ == "__main__":
	if len(sys.argv) < 4 or len(sys.argv) > 5:
		print "usage is graphviz.py <username> <password> <primarySObject> [depth]"
	else:
		g = GvWriter()
		g.generate()