import dis
import struct
import array
import types
import functools


try:
	_array_to_bytes = array.array.tobytes
except AttributeError:
	_array_to_bytes = array.array.tostring


class _Bytecode:
	def __init__(self):
	    code = (lambda: x if x else y).__code__.co_code
	    opcode, oparg = struct.unpack_from('BB', code, 2)

	    # Starting with Python 3.6, the bytecode format has changed, using
	    # 16-bit words (8-bit opcode + 8-bit argument) for each instruction,
	    # as opposed to previously 24 bit (8-bit opcode + 16-bit argument)
	    # for instructions that expect an argument and otherwise 8 bit.
	    # https://bugs.python.org/issue26647
	    if dis.opname[opcode] == 'POP_JUMP_IF_FALSE':
	        self.argument = struct.Struct('B')
	        self.have_argument = 0
	        # As of Python 3.6, jump targets are still addressed by their
	        # byte unit. This is matter to change, so that jump targets,
	        # in the future might refer to code units (address in bytes / 2).
	        # https://bugs.python.org/issue26647
	        self.jump_unit = 8 // oparg
	    else:
	        self.argument = struct.Struct('<H')
	        self.have_argument = dis.HAVE_ARGUMENT
	        self.jump_unit = 1

	@property
	def argument_bits(self):
	    return self.argument.size * 8


_BYTECODE = _Bytecode()


def _make_code(code, codestring):
	args = [
	    code.co_argcount,  code.co_nlocals,     code.co_stacksize,
	    code.co_flags,     codestring,          code.co_consts,
	    code.co_names,     code.co_varnames,    code.co_filename,
	    code.co_name,      code.co_firstlineno, code.co_lnotab,
	    code.co_freevars,  code.co_cellvars
	]

	try:
	    args.insert(1, code.co_kwonlyargcount)  # PY3
	except AttributeError:
	    pass
	
	try:
	    args.insert(1, code.co_posonlyargcount)  # PY3
	except AttributeError:
	    pass

	return types.CodeType(*args)


def _parse_instructions(code):
	extended_arg = 0
	extended_arg_offset = None
	pos = 0

	while pos < len(code):
	    offset = pos
	    if extended_arg_offset is not None:
	        offset = extended_arg_offset

	    opcode = struct.unpack_from('B', code, pos)[0]
	    pos += 1

	    oparg = None
	    if opcode >= _BYTECODE.have_argument:
	        oparg = extended_arg | _BYTECODE.argument.unpack_from(code, pos)[0]
	        pos += _BYTECODE.argument.size

	        if opcode == dis.EXTENDED_ARG:
	            extended_arg = oparg << _BYTECODE.argument_bits
	            extended_arg_offset = offset
	            continue

	    extended_arg = 0
	    extended_arg_offset = None
	    yield (dis.opname[opcode], oparg, offset)


def _get_instruction_size(opname, oparg=0):
	size = 1

	extended_arg = oparg >> _BYTECODE.argument_bits
	if extended_arg != 0:
	    size += _get_instruction_size('EXTENDED_ARG', extended_arg)
	    oparg &= (1 << _BYTECODE.argument_bits) - 1

	opcode = dis.opmap[opname]
	if opcode >= _BYTECODE.have_argument:
	    size += _BYTECODE.argument.size

	return size


def _get_instructions_size(ops):
	size = 0
	for op in ops:
	    if isinstance(op, str):
	        size += _get_instruction_size(op)
	    else:
	        size += _get_instruction_size(*op)
	return size


def _write_instruction(buf, pos, opname, oparg=0):
	extended_arg = oparg >> _BYTECODE.argument_bits
	if extended_arg != 0:
	    pos = _write_instruction(buf, pos, 'EXTENDED_ARG', extended_arg)
	    oparg &= (1 << _BYTECODE.argument_bits) - 1

	opcode = dis.opmap[opname]
	buf[pos] = opcode
	pos += 1

	if opcode >= _BYTECODE.have_argument:
	    _BYTECODE.argument.pack_into(buf, pos, oparg)
	    pos += _BYTECODE.argument.size

	return pos


def _write_instructions(buf, pos, ops):
	for op in ops:
	    if isinstance(op, str):
	        pos = _write_instruction(buf, pos, op)
	    else:
	        pos = _write_instruction(buf, pos, *op)
	return pos


def _find_labels_and_gotos(code):
	labels = {}
	gotos = []

	block_stack = []
	block_counter = 0

	opname1 = oparg1 = offset1 = None
	opname2 = oparg2 = offset2 = None
	opname3 = oparg3 = offset3 = None

	for opname4, oparg4, offset4 in _parse_instructions(code.co_code):
	    if opname1 in ('LOAD_GLOBAL', 'LOAD_NAME'):
	        if opname2 == 'LOAD_ATTR' and opname3 == 'POP_TOP':
	            name = code.co_names[oparg1]
	            if name == 'label':
	                if oparg2 in labels:
	                    raise SyntaxError('Ambiguous label {0!r}'.format(
	                        code.co_names[oparg2]
	                    ))
	                labels[oparg2] = (offset1,
	                                  offset4,
	                                  tuple(block_stack))
	            elif name == 'goto':
	                gotos.append((offset1,
	                              offset4,
	                              oparg2,
	                              tuple(block_stack)))
	    elif opname1 in ('SETUP_LOOP',
	                     'SETUP_EXCEPT', 'SETUP_FINALLY',
	                     'SETUP_WITH', 'SETUP_ASYNC_WITH'):
	        block_counter += 1
	        block_stack.append(block_counter)
	    elif opname1 == 'POP_BLOCK' and block_stack:
	        block_stack.pop()

	    opname1, oparg1, offset1 = opname2, oparg2, offset2
	    opname2, oparg2, offset2 = opname3, oparg3, offset3
	    opname3, oparg3, offset3 = opname4, oparg4, offset4

	return labels, gotos


def _inject_nop_sled(buf, pos, end):
	while pos < end:
	    pos = _write_instruction(buf, pos, 'NOP')


def _patch_code(code):
	labels, gotos = _find_labels_and_gotos(code)
	buf = array.array('B', code.co_code)

	for pos, end, _ in labels.values():
	    _inject_nop_sled(buf, pos, end)

	for pos, end, label, origin_stack in gotos:
	    try:
	        _, target, target_stack = labels[label]
	    except KeyError:
	        raise SyntaxError('Unknown label {0!r}'.format(
	            code.co_names[label]
	        ))

	    target_depth = len(target_stack)
	    if origin_stack[:target_depth] != target_stack:
	        raise SyntaxError('Jump into different block')

	    ops = []
	    for i in range(len(origin_stack) - target_depth):
	        ops.append('POP_BLOCK')
	    ops.append(('JUMP_ABSOLUTE', target // _BYTECODE.jump_unit))

	    if pos + _get_instructions_size(ops) > end:
	        # not enough space, add code at buffer end and jump there
	        buf_end = len(buf)

	        go_to_end_ops = [('JUMP_ABSOLUTE', buf_end // _BYTECODE.jump_unit)]

	        if pos + _get_instructions_size(go_to_end_ops) > end:
	            # not sure if reachable
	            raise SyntaxError('Goto in an incredibly huge function')

	        pos = _write_instructions(buf, pos, go_to_end_ops)
	        _inject_nop_sled(buf, pos, end)

	        buf.extend([0] * _get_instructions_size(ops))
	        _write_instructions(buf, buf_end, ops)
	    else:
	        pos = _write_instructions(buf, pos, ops)
	        _inject_nop_sled(buf, pos, end)

	return _make_code(code, _array_to_bytes(buf))


def with_goto(func_or_code):
	if isinstance(func_or_code, types.CodeType):
	    return _patch_code(func_or_code)

	return functools.update_wrapper(
	    types.FunctionType(
	        _patch_code(func_or_code.__code__),
	        func_or_code.__globals__,
	        func_or_code.__name__,
	        func_or_code.__defaults__,
	        func_or_code.__closure__,
	    ),
	    func_or_code
	)
