#!/usr/local/bin/python

class base:
    @staticmethod
    def init_class(c, arg):
	if isinstance(arg, dict):
	    try:
		return c(**arg)
	    except:
		print "%s" % arg
		raise
	elif isinstance(arg, c):
	    return arg
	else:
	    return None

    @staticmethod
    def make_list(a):
	if isinstance(a, basestring) or isinstance(a, dict): return [ a ]
	elif getattr(a, '__iter__', None): return a
	else: return [ a ]

    def get_attribute(self, key):
	rv = None
	attrs = getattr(self, 'attribute', None)
	if attrs:
	    for a in attrs:
		if a.attribute == key:
		    rv = a.value
		    break
	return rv



class ConsistencyError(RuntimeError): pass

class Attribute(base):
    def __init__(self, attribute, value):
	self.attribute = attribute
	self.value = value

    def clone(self):
	return Attribute(attribute=self.attribute, value=self.value)

    def to_dict(self):
	return { 'attribute': self.attribute, 'value': self.value }

class Capacity(base):
    def __init__(self, rate, kind):
	self.rate = float(rate)
	self.kind = kind

    def clone(self):
	return Capacity(rate=self.rate, kind=self.kind)

    def to_dict(self):
	return { 'rate': float(self.rate), 'kind': self.kind }

class Latency(base):
    def __init__(self, time, kind):
	self.time = float(time)
	self.kind = kind

    def clone(self):
	return Latency(time=self.time, kind=self.kind)

    def to_dict(self):
	return { 'time': float(self.time), 'kind': self.kind }

class Substrate(base):
    def __init__(self, name, capacity=None, latency=None, attribute=[]):
	self.name = name
	self.capacity = self.init_class(Capacity, capacity)
	self.latency = self.init_class(Latency, latency)
	self.attribute = [ self.init_class(Attribute, a) \
		for a in self.make_list(attribute) ]
	self.interfaces = [ ]

    def clone(self):
	if self.capacity: c = self.capacity.clone()
	else: c = None

	if self.latency: l = self.latency.clone()
	else: l = None

	return Substrate(name=self.name,
		capacity=c,
		latency=l,
		attribute = [a.clone() for a in self.attribute])

    def to_dict(self):
	rv = { 'name': self.name }
	if self.capacity:
	    rv['capacity'] = self.capacity.to_dict()
	if self.latency:
	    rv['latency'] = self.latency.to_dict()
	if self.attribute:
	    rv['attribute'] = [ a.to_dict() for a in self.attribute ]
	return rv

class CPU(base):
    def __init__(self, type, attribute=[]):
	self.type = type
	self.attribute = [ self.init_class(Attribute, a) for a in \
		self.make_list(attribute) ]

    def clone(self):
	return CPU(type=self.type,
		attribute = [a.clone() for a in self.attribute])

    def to_dict(self):
	rv = { 'type': self.type}
	if self.attribute:
	    rv['attribute'] = [ a.to_dict() for a in self.attribute ]
	return rv

class Storage(base):
    def __init__(self, amount, persistence, attribute=[]):
	self.amount = float(amount)
	self.presistence = persistence
	self.attribute = [ self.init_class(Attribute, a) \
		for a in self.make_list(attribute) ]

    def clone(self):
	return Storage(amount=self.amount, persistence=self.persistence, 
		attribute = [a.clone() for a in self.attribute])

    def to_dict(self):
	rv = { 'amount': float(self.amount), 'persistence': self.persistence }
	if self.attribute:
	    rv['attribute'] = [ a.to_dict() for a in self.attribute ]
	return rv

class OperatingSystem(base):
    def __init__(self, name=None, version=None, distribution=None,
	    distributionversion=None, attribute=[]):
	self.name = name
	self.version = version
	self.distribution = distribution
	self.distributionversion = distributionversion
	self.attribute = [ self.init_class(Attribute, a) \
		for a in self.make_list(attribute) ]

    def clone(self):
	return OperatingSystem(name=self.name,
		version=self.version,
		distribution=self.distribution,
		distributionversion=self.distributionversion,
		attribute = [ a.clone() for a in self.attribute])

    def to_dict(self):
	rv = { }
	if self.name: rv['name'] = self.name
	if self.version: rv['version'] = self.version
	if self.distribution: rv['version'] = self.distribution
	if self.distributionversion: rv['version'] = self.distributionversion
	if self.attribute:
	    rv['attribute'] = [ a.to_dict() for a in self.attribute ]
	return rv

class Software(base):
    def __init__(self, location, install=None, attribute=[]):
	self.location = location
	self.install = install
	self.attribute = [ self.init_class(Attribute, a)\
		for a in self.make_list(attribute) ]

    def clone(self):
	return Software(location=self.location, install=self.install, 
		attribute=[a.clone() for a in self.attribute])

    def to_dict(self):
	rv = { 'location': self.location }
	if self.install: rv['install'] = self.install
	if self.attribute:
	    rv['attribute'] = [ a.to_dict() for a in self.attribute ]
	return rv

class Interface(base):
    def __init__(self, substrate, capacity=None, latency=None, attribute=[],
	    element=None):
	self.substrate = self.make_list(substrate)
	self.capacity = self.init_class(Capacity, capacity)
	self.latency = self.init_class(Latency, latency)
	self.attribute = [ self.init_class(Attribute, a) \
		for a in self.make_list(attribute) ]
	self.element = element 
	self.subs = [ ]

    def clone(self):
	if self.capacity: c = self.capacity.clone()
	else: c = None

	if self.latency: l = self.latency.clone()
	else: l = None

	return Interface(substrate=self.substrate,
		capacity=c, latency=l,
		attribute = [ a.clone() for a in self.attribute])

    def to_dict(self):
	rv = { 'substrate': self.substrate }
	if self.capacity:
	    rv['capacity'] = self.capacity.to_dict()
	if self.latency:
	    rv['latency'] = self.latency.to_dict()
	if self.attribute:
	    rv['attribute'] = [ a.to_dict() for a in self.attribute ]
	return rv

class Computer(base):
    def __init__(self, name=[], cpu=[], os=[], software=[], storage=[],
	    interface=[], attribute=[]):
	def assign_element(i):
	    i.element = self

	self.name = self.make_list(name)
	self.cpu = [ self.init_class(CPU, c)  for c in self.make_list(cpu) ]
	self.os = [ self.init_class(OperatingSystem, c) \
		for c in self.make_list(os) ]
	self.software = [ self.init_class(Software, c) \
		for c in self.make_list(software) ]
	self.storage = [ self.init_class(Storage, c) \
		for c in self.make_list(storage) ]
	self.interface = [ self.init_class(Interface, c) \
		for c in self.make_list(interface) ]
	self.attribute = [ self.init_class(Attribute, a) \
		for a in self.make_list(attribute) ]
	map(assign_element, self.interface)

    def clone(self):
	return Computer(name=self.name,
		cpu=[x.clone() for x in self.cpu],
		os=[x.clone() for x in self.os],
		software=[x.clone() for x in self.software],
		storage=[x.clone() for x in self.storage],
		interface=[x.clone() for x in self.interface],
		attribute=[x.clone() for x in self.attribute])

    def to_dict(self):
	rv = { }
	if self.name:
	    rv['name'] = self.name
	if self.cpu:
	    rv['cpu'] = [ c.to_dict() for  c in self.cpu ]
	if self.os:
	    rv['os'] = [ o.to_dict() for o in self.os ]
	if self.software:
	    rv['software'] = [ s.to_dict() for s in self.software ]
	if self.storage:
	    rv['storage'] = [ s.to_dict for s in self.storage ]
	if self.interface:
	    rv['interface'] = [ i.to_dict() for i in self.interface ]
	if self.attribute:
	    rv['attribute'] = [ i.to_dict() for i in self.attribute ]
	return { 'computer': rv }

class Other(base):
    def __init__(self, interface=[], attribute=[]):
	self.interface = [ self.init_class(Interface, c) \
		for c in self.make_list(interface) ]
	self.attribute = [ self.init_class(Attribute, c) \
		for c in self.make_list(attribute) ]

    def clone(self):
	return Other(interface=[i.clone() for i in self.interface], 
		attribute=[a.clone() for a in attribute])

    def to_dict(self):
	if self.interface:
	    rv['interface'] = [ i.to_dict() for i in self.interface ]
	if self.attribute:
	    rv['attribute'] = [ a.to_dict() for a in self.attribute ]


class Topology(base):
    @staticmethod
    def init_element(e):
	"""
	e should be of the form { typename: args } where args is a dict full of
	the right parameters to initialize the element.  e should have only one
	key, but we walk e's keys in an arbitrary order and instantiate the
	first key we know how to.
	"""
	classmap = {
		'computer': Computer,
		'other': Other,
	    }

	if isinstance(e, dict):
	    for k in e.keys():
		cl = classmap.get(k, None)
		if cl: return cl(**e[k])
	else:
	    return e

    def __init__(self, substrates=[], elements=[]):
	self.substrates = [ self.init_class(Substrate, s) \
		for s in self.make_list(substrates) ]
	self.elements = [ self.init_element(e) \
		for e in self.make_list(elements) ]
	self.incorporate_elements()

    def incorporate_elements(self):

	# Could to this init in one gulp, but we want to look for duplicate
	# substrate names
	substrate_map = { }
	for s in self.substrates:
	    s.interfaces = [ ]
	    if not substrate_map.has_key(s.name):
		substrate_map[s.name] = s
	    else:
		raise ConsistencyError("Duplicate substrate name %s" % s.name)

	for e in self.elements:
	    for i in e.interface:
		i.element = e
		i.subs = [ ]
		for sn in i.substrate:
		    # NB, interfaces have substrate names in their substrate
		    # attribute.
		    if substrate_map.has_key(sn):
			sub = substrate_map[sn]
			i.subs.append(sub)
			sub.interfaces.append(i)
		    else:
			raise ConsistencyError("No such substrate for %s" % sn)

    def clone(self):
	return Topology(substrates=[s.clone() for s in self.substrates], 
		elements=[e.clone() for e in self.elements])


    def make_indices(self):
	sub_index = dict([(s.name, s) for s in self.substrates])
	elem_index = dict([(n, e) for e in self.elements for n in e.name])

    def to_dict(self):
	rv = { }
	if self.substrates:
	    rv['substrates'] = [ s.to_dict() for s in self.substrates ]
	if self.elements:
	    rv['elements'] = [ s.to_dict() for s in self.elements ]
	return rv

def topology_from_xml(string=None, file=None, filename=None, top="topology"):
    import xml.parsers.expat

    class parser:
	def __init__(self):
	    self.stack = [ ]
	    self.chars = ""
	    self.key = ""
	    self.have_chars = False
	    self.current = { }
	    self.in_cdata = False
	
	def start_element(self, name, attrs):
	    self.chars = ""
	    self.have_chars = False
	    self.key = str(name)
	    self.stack.append((self.current, self.key))
	    self.current = { }

	def end_element(self, name):
	    if self.have_chars:
		self.chars = self.chars.strip()
		if len(self.chars) >0:
		    addit = self.chars
		else:
		    addit = self.current
	    else:
		addit = self.current

	    parent, key = self.stack.pop()
	    if parent.has_key(key):
		if isinstance(parent[key], list):
		    parent[key].append(addit)
		else:
		    parent[key] = [parent[key], addit]
	    else:
		parent[key] = addit
	    self.current = parent
	    self.key = key

	    self.chars = ""
	    self.have_chars = False

	def char_data(self, data):
	    self.have_chars = True
	    self.chars += data

    p = parser()
    xp = xml.parsers.expat.ParserCreate()

    xp.StartElementHandler = p.start_element
    xp.EndElementHandler = p.end_element
    xp.CharacterDataHandler = p.char_data

    num_set = len([ x for x in (string, filename, file)\
	    if x is not None ])

    if num_set != 1:
	raise RuntimeError("Exactly one one of file, filename and string " + \
		"must be set")
    elif filename:
	f = open(filename, "r")
	xp.ParseFile(f)
	f.close()
    elif file:
	xp.ParseFile(file)
    elif string:
	xp.Parse(string, isfinal=True)
    else:
	return None

    return Topology(**p.current[top])

def topology_to_xml(t, top=None):
    """
    Print the topology as XML.  This is quick and dirty, but should work for
    many purposes.  Convert the topology to a dict and print it recursively.
    """
    from xml.sax.saxutils import escape

    def dict_to_xml(e, top=None):
	if top: rv = "<%s>" % top
	else: rv = ""

	for k in e.keys():
	    if isinstance(e[k], (basestring, int, float, long)):
		rv += "<%s>%s</%s>" % (k, escape(e[k]), k)
	    elif isinstance(e[k], dict):
		rv += "<%s>%s</%s>" % (k, dict_to_xml(e[k]), k)
	    elif getattr(e[k], '__iter__', None):
		for ee in e[k]:
		    if isinstance(ee, dict):
			rv += "<%s>%s</%s>" % (k, dict_to_xml(ee), k)
		    else:
			rv += "<%s>%s</%s>" % (k, escape(ee), k)
	    else:
		raise ConsistencyError("What is this?? %s %s" % (k, e[k]))
	if top: rv += "</%s>" % top
	return rv

    return dict_to_xml(t.to_dict(), top)


def topology_to_vtopo(t):
    nodes = [ ]
    lans = [ ]

    for eidx, e in enumerate(t.elements):
	if e.name: name = e.name[0]
	else: name = "unnamed_node%d" % eidx
	
	ips = [ ]
	for idx, i in enumerate(e.interface):
	    ip = i.get_attribute('ip4_address')
	    ips.append(ip)
	    port = "%s:%d" % (name, idx)
	    for idx, s in enumerate(i.subs):
		bw = 100000
		delay = 0.0
		if s.capacity:
		    bw = s.capacity.rate
		if i.capacity:
		    bw = i.capacity.rate

		if s.latency:
		    delay = s.latency.time
		if i.latency:
		    bw = i.latency.time

		lans.append({
		    'member': port,
		    'vname': s.name,
		    'ip': ip,
		    'vnode': name,
		    'delay': delay,
		    'bandwidth': bw,
		    })
	nodes.append({
	    'ips': ":".join(ips),
	    'vname': name,
	    })

    return { 'node': nodes, 'lan': lans }

def topology_to_ns2(t):
    out = """
set ns [new Simulator]
source tb_compat.tcl

"""
    for e in t.elements:
	rpms = ""
	tarfiles = ""
	if isinstance(e, Computer):
	    name = e.name[0]
	    out += "set %s [$ns node]\n" % name
	    if e.os and len(e.os) == 1:
		osid = e.os[0].get_attribute('osid')
		if osid:
		    out += "tb-set-node-os $%s %s\n" % (name, osid)
	    for s in e.software:
		if s.install:
		    tarfiles += "%s %s " % (s.install, s.location)
		else:
		    rpms += "%s " % s.location
	    if rpms:
		out += "tb-set-node-rpms $%s %s\n" % (name, rpms)
	    if tarfiles:
		out += "tb-set-node-tarfiles $%s %s\n" % (name, tarfiles)
	    startcmd = e.get_attribute('startup')
	    if startcmd:
		out+= 'tb-set-node-startcmd $%s "%s"\n' % (name, startcmd)
	    out+= "\n"
    
    for idx, s in enumerate(t.substrates):
	loss = s.get_attribute('loss')
	if s.latency:
	    delay = s.latency.time
	else:
	    delay = 0
	name = s.name or "sub%d" % idx

	if len(s.interfaces) > 2:
	    # Lan
	    members = [ i.element.name[0] for i in s.interfaces]
	    out += 'set %s [$ns make-lan "%s" %f %fms ]\n' % \
		    (name, " ".join(members), s.capacity.rate, delay)
	    if loss:
		"tb-set-lan-loss $%s %f\n" % (name, float(loss))

	    for i in s.interfaces:
		e = i.element
		ip = e.get_attribute("ip4_address")
		if ip:
		    out += "tb-set-ip-lan $%s $%s %s\n" % (e.name, name, ip)
		if i.capacity and i.capacity.rate != s.capacity.rate:
		    out += "tb-set-node-lan-bandwidth $%s $%s %f\n" % \
			    (e.name[0], name, i.capacity.rate)
		if i.latency and i.latency.time != delay:
		    out += "tb-set-node-lan-delay $%s $%s %fms\n" % \
			    (e.name[0], name, i.latency.time)
		iloss = i.get_attribute('loss')
		if loss and iloss != loss :
		    out += "tb-set-node-lan-loss $%s $%s %f" % \
			    (e.name[0], name, float(loss))
	    out+= "\n"
	elif len(s.interfaces) == 2:
	    f = s.interfaces[0]
	    t = s.interfaces[1]

	    out += "set %s [$ns duplex-link $%s $%s %f %fms DropTail]\n" %\
		    (name, f.element.name[0], t.element.name[0],
			    s.capacity.rate, delay)
	    if loss:
		out += "tb-set-link-loss $%s %f\n" % (name, float(loss))

	    for i in s.interfaces:
		lloss = i.get_attribute("loss")
		cap_override = i.capacity and \
			i.capacity.rate != s.capacity.rate
		delay_override = i.latency and \
			i.latency.time != delay
		loss_override = lloss and lloss != loss
		if cap_override or delay_override or loss_override:
		    if i.capacity: cap = i.capacity.rate
		    else: cap = s.capacity.rate

		    if i.latency: delay = i.latency.time

		    if lloss: loss = lloss
		    else: loss = loss or 0.0

		    out += "tb-set-link-simplex-params $%s $%s %fms %f %f\n" % \
			    (name, i.element.name[0], delay, cap, loss)
	    out+= "\n"
    out+="""
$ns run
"""
    return out