# Copyright (c) 2013, Albert Zeyer, www.az2000.de
# All rights reserved.
# Code under 2-clause BSD licence.

import dis
import types

def _modified_jumps(codestr, jumprel=None, jumpaddrmap=None, start=None, end=None):
	if jumprel is not None:
		assert jumpaddrmap is None, "specify only one of jumprel and jumpaddrmap"
		jumpaddrmap = lambda addr, op, n: n + jumprel
	assert jumpaddrmap is not None, "one of jumprel and jumpaddrmap must be specified"

	codestr = bytearray(codestr)
	if start is None: start = 0
	if end is None: end = len(codestr)
	i = start
	while i < end:
		op = codestr[i]
		i += 1

		if op >= dis.HAVE_ARGUMENT:
			b1 = codestr[i]
			b2 = codestr[i+1]
			num = b2 * 256 + b1
			del b1,b2
			i += 2
		else:
			num = 0

		if op in dis.hasjabs:
			assert op >= dis.HAVE_ARGUMENT
			num = jumpaddrmap(i, op, num)
			codestr[i-2] = chr(num & 255)
			codestr[i-1] = chr(num >> 8)

		if op in dis.hasjrel and jumprel is None:
			assert op >= dis.HAVE_ARGUMENT
			num += i # because it is a relative jump
			num = jumpaddrmap(i, op, num) # map
			num -= i # convert back to relative jump
			codestr[i-2] = chr(num & 255)
			codestr[i-1] = chr(num >> 8)

	return str(codestr)

def _codestr_without_final_return(codestr):
	assert len(codestr) >= 4
	assert codestr[-4] == dis.opmap["LOAD_CONST"]
	assert codestr[-1] == dis.opmap["RETURN_VALUE"]
	return codestr[:-4]

def _prefix_codestr(codestr1, codestr2):
	# see dis.findlinestarts() about co_firstlineno and co_lnotab
	codestr = codestr1 + codestr2
	codestr = _modified_jumps(codestr, start=len(codestr1), end=len(codestr), jumprel=len(codestr1))
	return codestr

def _modified_code(c, **kwargs):
	CodeArgs = [
		"argcount", "nlocals", "stacksize", "flags", "code",
		"consts", "names", "varnames", "filename", "name",
		"firstlineno", "lnotab", "freevars", "cellvars"]
	c_dict = dict([(arg, getattr(c, "co_" + arg)) for arg in CodeArgs])

	for key,value in kwargs.items():
		assert key in c_dict
		c_dict[key] = value

	c = types.CodeType(*[c_dict[arg] for arg in CodeArgs])
	return c

def _find_setup_blocks(codestr, start, end):
	"Yields (op, absolute target instraddr, loop-type)"
	i = start
	SETUPS = [dis.opmap[opname] for opname in ["SETUP_LOOP", "SETUP_EXCEPT", "SETUP_FINALLY"]]
	POP_BLOCK = dis.opmap["POP_BLOCK"]
	FOR_ITER = dis.opmap["FOR_ITER"]
	blockstack = []
	while i < end:
		op = ord(codestr[i])
		i += 1

		if op >= dis.HAVE_ARGUMENT:
			b1 = ord(codestr[i])
			b2 = ord(codestr[i+1])
			num = b2 * 256 + b1
			del b1,b2
			i += 2
		else:
			num = 0

		if op in SETUPS:
			assert op >= dis.HAVE_ARGUMENT
			blockstack += [[op, i + num, None]]

		elif op == FOR_ITER:
			assert len(blockstack) > 0
			blockstack[-1][2] = op

		elif op == POP_BLOCK:
			assert len(blockstack) > 0
			blockstack.pop(len(blockstack) - 1)

	return blockstack


def replace_code(codeobj, instaddr, removelen=0, addcodestr=""):
	assert isinstance(codeobj, types.CodeType)
	assert removelen >= 0
	if removelen == 0 and len(addcodestr) == 0: return codeobj

	# Search right place in lnotab.
	lnotab = codeobj.co_lnotab
	assert len(lnotab) % 2 == 0
	lnotab_instaddr = 0
	lnotab_idx = 0
	lnotab_len = len(lnotab)
	while lnotab_idx < lnotab_len:
		if lnotab_instaddr >= instaddr: break
		addrincr, lineincr = map(ord, lnotab[lnotab_idx:lnotab_idx+2])
		lnotab_instaddr += addrincr
		lnotab_idx += 2
	assert lnotab_instaddr >= instaddr, "instaddr %i not in lnotab" % instaddr

	# If we skipped it, insert a dummy entry to lnotab.
	if lnotab_instaddr > instaddr:
		# Insert. addrincr, lineincr are from the last entry.
		lnotab = \
			lnotab[:lnotab_idx-2] + \
			chr(addrincr - (lnotab_instaddr - instaddr)) + chr(0) + \
			chr(lnotab_instaddr - instaddr) + chr(lineincr) + \
			lnotab[lnotab_idx:]
		lnotab_instaddr = instaddr
	del addrincr, lineincr
	assert lnotab_instaddr == instaddr
	# And lnotab_idx is right where the upcoming lnotab-data starts.

	# Check whether instaddr is sane.
	codestr = codeobj.co_code
	codeidx = 0
	codelen = len(codestr)
	while codeidx < codelen:
		if codeidx >= instaddr: break
		op = ord(codestr[codeidx])
		codeidx += 1