#!/usr/bin/env python

# from __future__ import absolute_import

import pickle
import signal
import sys
import gzip

import profile_pb2

MAX_STACK_DEPTH = 64

# Don't profile too quickly, each pause is for about 2ms
SAMPLE_HZ = 10
sample_interval = 1.0/SAMPLE_HZ

TEST_DURATION_S = 10

def main():
	signal.signal(signal.SIGINT, interrupt)
	signal.signal(signal.SIGALRM, interrupt)

	signal.signal(signal.SIGPROF, handle_signal)
	signal.setitimer(signal.ITIMER_PROF, sample_interval, sample_interval)
	signal.alarm(TEST_DURATION_S)

	pause()

def pause():
	val = range(10000)
	while True:
		a, b, c = lambda v: pickle.dumps(v), lambda v: pickle.dumps(v), lambda v: dumps(v)
		a(val)
		b(val)
		c(val)
		d = lambda v: pickle.dumps(v)
		d(val)
		pickle.dumps(val)

def dumps(val):
	pickle.dumps(val)

in_handler = False
stacks = {}
code_by_loc = {}

def interrupt(signum, frame):
	signal.setitimer(signal.ITIMER_PROF, 0)

	# write_text(stacks, code_by_loc)
	write_proto(stacks, code_by_loc)

	# print dir(sys.modules["pickle"].dumps.__code__)
	sys.exit(0)
	# for (k, v) in sys.modules.iteritems():
	# 	if v is None:
	# 		continue
	# 	print k, v.__dict__.get("__file__", None)
	# sys.exit(0)

def write_proto(stacks, code_by_loc):
	pr = profile_pb2.Profile()

	# Profile.StringTable
	strs = {}
	def add_string(s):
		if s not in strs:
			i = len(strs)
			strs[s] = i
			pr.string_table.append(s)
		return strs[s]

	add_string("")

	pr.period = long(sample_interval*1e9)
	pr.duration_nanos = long(TEST_DURATION_S*1e9)
	# pr.time_nanos
	# pr.x_drop_frames

	st = pr.sample_type.add()
	st.x_type = add_string("samples")
	st.x_unit = add_string("count")
	st = pr.sample_type.add()
	st.x_type = add_string("cpu")
	st.x_unit = add_string("nanoseconds")

	pr.period_type.x_type = add_string("cpu")
	pr.period_type.x_unit = add_string("nanoseconds")

	# Profile.Mapping
	m = pr.mapping.add()
	m.id = len(pr.mapping)
	m.x_file = add_string("python")
	m.has_functions = True
	m.has_filenames = True
	m.has_line_numbers = True

	# Profile.Function
	fns = {}
	fn_names = {}
	for (loc, code) in code_by_loc.iteritems():
		filename, name = code[0].co_filename, code[0].co_name
		key = (filename, name,)
		if key not in fns:
			i = len(fns)+1
			fns[key] = i

			unique_name = name
			if name not in fn_names:
				fn_names[name] = 1
			else:
				fn_names[name] += 1
				unique_name = "{}${}".format(name, fn_names[name])

			f = pr.function.add()
			f.id = i
			f.x_name = add_string(unique_name)
			f.x_system_name = add_string(unique_name)
			f.x_filename = add_string(filename)

	# Profile.Location
	locs = {}
	for (loc, code) in code_by_loc.iteritems():
		start, offset = loc[0], loc[1]
		key = (start, offset,)
		if key not in locs:
			i = len(locs)+1
			locs[key] = i
			l = pr.location.add()
			l.id = i
			l.mapping_idx = 1
			l.address = start+offset

			filename, name = code[0].co_filename, code[0].co_name
			fn_key = (filename, name,)
			ln = l.line.add()
			ln.function_idx = fns[fn_key]
			ln.line = code[1]

	# Profile.Sample
	for (stack, n) in stacks.iteritems():
		pass
		sample = pr.sample.add()
		for loc in stack:
			start, offset = loc[0], loc[1]
			l_key = (start, offset,)
			sample.location_idx.append(locs[l_key])
		sample.value.append(n)
		sample.value.append(n*pr.period)

	s = pr.SerializeToString()

	with gzip.open("/tmp/py.pb.gz", "wb") as f:
		f.write(s)

def write_text(stacks, code_by_loc):
	nn = 0
	for (_, n) in stacks.iteritems():
		nn += n
	print "cpu profile: total {}".format(nn)

	for (stack, n) in stacks.iteritems():
		print "{} @ {}".format(n, " ".join([hex(loc[0]+loc[1]) for loc in stack]))
		for loc in stack:
			code = code_by_loc[loc]
			print "#\t{}\t{}+{}\t{}:{}".format(
				hex(loc[0]+loc[1]),
				code[0].co_name, hex(loc[1]),
				code[0].co_filename, code[1])
		print

def handle_signal(signum, frame):
	global in_handler
	if in_handler:
		return

	try:
		in_handler = True
		# print "signal {}".format(signum)
		stack = []
		while frame is not None:
			if len(stack) >= MAX_STACK_DEPTH:
				stack = stack[1:]
			code = frame.f_code
			code_addr, code_offset = id(code.co_code), frame.f_lasti+1
			loc = (code_addr, code_offset,)
			if loc not in code_by_loc:
				code_by_loc[loc] = (code, frame.f_lineno,)
			stack.append(loc)

			# stack.append((frame.f_code.co_filename, frame.f_lineno, frame.f_code.co_name, frame.f_lasti,))
			# if frame.f_code is sys.modules["pickle"].dumps.__code__:
			# 	print "pickle.dumps!!"
			# frame_some(signum, frame)

			frame = frame.f_back
		stack = tuple(stack)
		stacks[stack] = stacks.get(stack, 0) + 1
	finally:
		in_handler = False

# def frame_some(signum, frame):
# 	print "{}:{} / {} / {}+{}".format(
# 		frame.f_code.co_filename, frame.f_lineno,
# 		frame.f_code.co_name,
# 		hex(id(frame.f_code.co_code)), hex(frame.f_lasti),
# 	)

# def frame_deets(signum, frame):
# 	print "signal={}".format(signum)
# 	print dir(frame)
# 	print "f_lasti={}".format(frame.f_lasti)
# 	print "f_lineno={}".format(frame.f_lineno)
# 	print "f_back={}".format(frame.f_back)
# 	# print "f_builtins={}".format(frame.f_builtins)
# 	print "f_code={}".format(frame.f_code)
# 	print "f_exc_traceback={}".format(frame.f_exc_traceback)
# 	print "f_exc_type={}".format(frame.f_exc_type)
# 	print "f_exc_value={}".format(frame.f_exc_value)
# 	print "f_globals={}".format(frame.f_globals)
# 	print "f_locals={}".format(frame.f_locals)
# 	print "f_restricted={}".format(frame.f_restricted)
# 	print "f_trace={}".format(frame.f_trace)
# 	print
# 	print dir(frame.f_code)
# 	print "co_filename={}".format(frame.f_code.co_filename)
# 	print "co_firstlineno={}".format(frame.f_code.co_firstlineno)
# 	print "co_name={}".format(frame.f_code.co_name)
# 	print "co_names={}".format(frame.f_code.co_names)

if __name__ == '__main__':
	main()
