from .structure import *
from functools import partial

class Namespace:
	def __init__(self):
		self.counts = {}
		self.sigs = {}
	
	def GetName(self, sig):
		try:
			n = self.sigs[sig]
			if n:
				return sig.name + "_" + str(n)
			else:
				return sig.name
		except KeyError:
			try:
				n = self.counts[sig.name]
			except KeyError:
				n = 0
			self.sigs[sig] = n
			self.counts[sig.name] = n + 1
			if n:
				return sig.name + "_" + str(n)
			else:
				return sig.name

def ListSignals(node):
	if isinstance(node, Constant):
		return set()
	elif isinstance(node, Signal):
		return {node}
	elif isinstance(node, Operator):
		l = list(map(ListSignals, node.operands))
		return set().union(*l)
	elif isinstance(node, Slice):
		return ListSignals(node.value)
	elif isinstance(node, Cat):
		l = list(map(ListSignals, node.l))
		return set().union(*l)
	elif isinstance(node, Assign):
		return ListSignals(node.l) | ListSignals(node.r)
	elif isinstance(node, StatementList):
		l = list(map(ListSignals, node.l))
		return set().union(*l)
	elif isinstance(node, If):
		return ListSignals(node.cond) | ListSignals(node.t) | ListSignals(node.f)
	elif isinstance(node, Fragment):
		return ListSignals(node.comb) | ListSignals(node.sync)
	else:
		raise TypeError

def Convert(f, ins, outs, name="top"):
	ns = Namespace()
	
	clks = Signal(name="sys_clk")
	rsts = Signal(name="sys_rst")
	clk = ns.GetName(clks)
	rst = ns.GetName(rsts)
	
	def printsig(s):
		if s.bv.signed:
			n = "signed "
		else:
			n = ""
		if s.bv.width > 1:
			n += "[" + str(s.bv.width-1) + ":0] "
		n += ns.GetName(s)
		return n
	
	def printnode(level, node):
		if isinstance(node, Constant):
			return str(node)
		elif isinstance(node, Signal):
			return ns.GetName(node)
		elif isinstance(node, Operator):
			arity = len(node.operands)
			if arity == 1:
				r = self.op + str(node.operands[0])
			elif arity == 2:
				r = printnode(level, node.operands[0]) + " " + node.op + " " + printnode(level, node.operands[1])
			else:
				raise TypeError
			return "(" + r + ")"
		elif isinstance(node, Slice):
			if node.start + 1 == node.stop:
				sr = "[" + str(node.start) + "]"
			else:
				sr = "[" + str(node.stop-1) + ":" + str(node.start) + "]"
			return str(node.value) + sr
		elif isinstance(node, Cat):
			l = list(map(partial(printnode, level), node.l))
			l.reverse()
			return "{" + ", ".join(l) + "}"
		elif isinstance(node, Assign):
			# TODO: variables
			return "\t"*level + printnode(level, node.l) + " <= " + printnode(level, node.r) + ";\n"
		elif isinstance(node, StatementList):
			return "".join(list(map(partial(printnode, level), node.l)))
		elif isinstance(node, If):
			r = "\t"*level + "if " + printnode(level, node.cond) + " begin\n"
			r += printnode(level + 1, node.t)
			if node.f.l:
				r += "\t"*level + "end else begin\n"
				r += printnode(level + 1, node.f)
			r += "\t"*level + "end\n"
			return r
		else:
			raise TypeError
		
	r = "/* Autogenerated by Migen */\n"
	r += "module " + name + "(\n"
	r += "\tinput " + clk + ",\n"
	r += "\tinput " + rst
	if ins:
		r += ",\n\tinput " + ",\n\tinput ".join(map(printsig, ins)) 
	if outs:
		r += ",\n\toutput reg " + ",\n\toutput reg ".join(map(printsig, outs)) 
	r += "\n);\n\n"
	sigs = ListSignals(f).difference(ins, outs)
	for sig in sigs:
		r += "reg " + printsig(sig) + ";\n"
	r += "\n"
	
	if f.comb.l:
		r += "always @(*) begin\n"
		r += printnode(1, f.comb)
		r += "end\n\n"
	
	if f.sync.l:
		r += "always @(posedge " + clk + ") begin\n"
		r += printnode(1, f.sync)
		r += "end\n\n"
	
	r += "endmodule\n"
	
	return r