#!/bin/env python3
import pygraphviz

import sys
import os
import os.path
import subprocess as ps
import shutil
import argparse

import random

DBG = 1
PERMISSIIVE = False
DEAD_COUNT = 0 #number of files that failed, only relevant if PERMISSIVE is true

'''
Author: github.com/nilput

dependencies: 	
	pygraphviz 	(pip install)

	graphviz 	(command line tool: dot)
	llvm 		(command line tool: opt)
	clang 		(command line tool: clang)
'''

def find_c_sources(dir):
	files = []
	with os.scandir(dir) as it:
		for entry in it:
			if entry.is_dir():
				l = find_c_sources(entry.path)
				files.extend(l)
			elif entry.is_file() and entry.path.endswith('.c'):
				files.append(entry.path)
	return files

def sources_to_cmdlines(sources, cflags):
	assert(isinstance(cflags, str))
	cmdargs = tuple(cflags.split(' '))
	cmdlines = []
	for source in sources:
		cmdlines.append( (*cmdargs, source) )
	return cmdlines

def isworthy(line):
	#attempt to ignore autoconf tests, this isn't very robust :)
	if (line.find('conftest') == -1) and (line.find('-o') != -1) and (line.find('.c') != -1):
		#print(line)
		return True
	return False


def adjust_cmdline(cmdargs):
	assert(isinstance(cmdargs, tuple) or isinstance(cmdargs, list))
	for i,arg in enumerate(cmdargs):
		if arg.startswith('-o'):
			if arg == '-o':
				return cmdargs[:i] + cmdargs[i+2:]
			else:
				return cmdargs[:i] + cmdargs[i+1:]

	return cmdargs


def invkcache_to_cmdlines(invkcache):
	assert(isinstance(invkcache, str))
	cmdlines = []
	for line in invkcache.split('\n'):
		if isworthy(line):
			assert(line.startswith('pretendcc'))
			cmdargs = list(line.split(' '))[1:]
			cmdlines.append(adjust_cmdline(cmdargs))
	return cmdlines

def outify(path, ext):
	if not os.path.exists('outtmp'):
		os.mkdir('outtmp')
	name = os.path.basename(path)
	return os.path.join('outtmp', name + ext)

def getname(cmdline):
	return 'f' + str(random.randint(1, 100000))

#returns a list of dot files
def process_cmdlines(cmdlines):
	global PERMISSIVE
	global DEAD_COUNT

	assert(iter(cmdlines))
	dots = []
	for cmdline in cmdlines:
		assert(iter(cmdline))
		try:
			name = getname(cmdline)
			outllvm = outify(name, '.lv')
			if DBG:
				print('processing: "{}"'.format('CC ' + (' '.join(cmdline))))
			completed = ps.run(args=['clang', '-S', '-emit-llvm', *cmdline, '-o', outllvm])
			completed.check_returncode() #raises
			completed = ps.run(args=['opt', '-analyze', '-std-link-opts', '-dot-callgraph', outllvm])
			completed.check_returncode() #raises
			assert(os.path.exists('callgraph.dot'))
			outdot = outify(name, '.dot')
			shutil.move('callgraph.dot', outdot)
			dots.append(outdot)
		except ps.CalledProcessError as e:
			if PERMISSIVE:
				DEAD_COUNT = DEAD_COUNT + 1
			else:
				print('\nnote: you can use -w to try generating the graph anyways.\n')
				raise
	return dots


class cfunc:
	def __init__(self, name, registry):
		self.name = name
		self.callees_set = set()
		self.called_by_set = set()
		self.registry = registry
	def called_by(self, other):
		self.called_by_set.add(other)
	def calls(self, other):
		self.callees_set.add(other)
		self.registry.get_or_insert(other).called_by(self.name)

class registry:
	def __init__(self):
		self.dict = {}
	def get_or_insert(self, name):
		if name not in self.dict:
			self.dict[name] = cfunc(name, self)
		return self.dict[name]
	def iter_func_to_callees(self):
		for func in self.dict.values():
			yield (func.name, tuple(func.callees_set))


def try_get_label(graph, node_name):
	node = graph.get_node(node_name)
	return node.attr['label']

#returns graph
def process_dots(dots):
	reg = registry()
	for dot_file in dots:
		grph = pygraphviz.AGraph()
		with open(dot_file,'r') as f:
			grph.from_string(f.read())
		for edge in grph.edges():
			src,dest = edge[0], edge[1]
			srcn, destn = try_get_label(grph, src), try_get_label(grph, dest)
			if (srcn and destn):
				reg.get_or_insert(srcn).calls(destn)
			elif srcn or destn:
				srcn = srcn if srcn else '?'
				destn = destn if destn else '?'
				print('warning, unknown name for function call {} -> {}'.format(srcn, destn))
	ngraph = pygraphviz.AGraph(directed=True)
	for func_name, calls in reg.iter_func_to_callees():
		if func_name.find('external node') != -1:
			continue #hardcoded :)
		ngraph.add_node(func_name, label=func_name)
		for callee in calls:
			ngraph.add_edge(func_name, callee)
	return ngraph


def main():
	global PERMISSIVE
	global DEAD_COUNT

	parser = argparse.ArgumentParser()
	parser.add_argument('-d', help='input directory for source files')
	parser.add_argument('-o', help='output dotfile')
	parser.add_argument('--cflags', help='cflags')
	parser.add_argument('-p', help='invkcache.txt input for source files (generated by pretendcc.py)')
	parser.add_argument('-w', help='permissive (allow for compilation to fail)', action="store_true")


	args = parser.parse_args()

	PERMISSIVE = args.w
	
	if not args.o:
		print('error no output file, see -h')
		return

	cmdlines = None
	if args.d:
		source_dir = args.d
		cflags = args.cflags if args.cflags else ''
		sources = find_c_sources(source_dir)
		cmdlines = sources_to_cmdlines(sources,  cflags)
	elif args.p:
		invkcache = None
		with open(args.p, 'r') as f:
			invkcache = f.read()
		cmdlines = invkcache_to_cmdlines(invkcache)
	else:
		print('error no input dir, see -h')
		return

	assert(cmdlines)
	dots = process_cmdlines(cmdlines)

	G = process_dots(dots)
	G.write(args.o)
	G.draw(args.o+'.png', prog='dot')
	print()

	if DEAD_COUNT > 0:
		assert(PERMISSIVE)
		print('compilation of {} files failed'.format(DEAD_COUNT))
	print('generated "{}" and "{}"'.format(args.o, args.o + '.png'))


if __name__ == '__main__':
	main()