# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-2-Clause

from collections import namedtuple, OrderedDict
import dis
import inspect

from types import CodeType, ModuleType

from numba_cuda_mlir.numba_cuda.core import errors
from numba_cuda_mlir.numba_cuda import serialize
from numba_cuda_mlir.numba_cuda import utils
from numba_cuda_mlir.numba_cuda.utils import PYVERSION
from numba_cuda_mlir._threading import _LockedCounter


if PYVERSION in ((3, 12), (3, 13), (3, 14)):
    from opcode import _inline_cache_entries

    # Instruction/opcode length in bytes
    INSTR_LEN = 2
elif PYVERSION in ((3, 9), (3, 10), (3, 11)):
    pass
else:
    raise NotImplementedError(PYVERSION)


opcode_info = namedtuple("opcode_info", ["argsize"])
_ExceptionTableEntry = namedtuple("_ExceptionTableEntry", "start end target depth lasti")

# The following offset is used as a hack to inject a NOP at the start of the
# bytecode. So that function starting with `while True` will not have block-0
# as a jump target. The Lowerer puts argument initialization at block-0.
_FIXED_OFFSET = 2


def get_function_object(obj):
    """
    Objects that wraps function should provide a "__numba__" magic attribute
    that contains a name of an attribute that contains the actual python
    function object.
    """
    attr = getattr(obj, "__numba__", None)
    if attr:
        return getattr(obj, attr)
    return obj


def get_code_object(obj):
    "Shamelessly borrowed from llpython"
    return getattr(obj, "__code__", getattr(obj, "func_code", None))


def _as_opcodes(seq):
    lst = []
    for s in seq:
        c = dis.opmap.get(s)
        if c is not None:
            lst.append(c)
    return lst


JREL_OPS = frozenset(dis.hasjrel)
JABS_OPS = frozenset(dis.hasjabs)
JUMP_OPS = JREL_OPS | JABS_OPS
TERM_OPS = frozenset(_as_opcodes(["RETURN_VALUE", "RAISE_VARARGS"]))
EXTENDED_ARG = dis.EXTENDED_ARG
HAVE_ARGUMENT = dis.HAVE_ARGUMENT


class ByteCodeInst:
    """
    Attributes
    ----------
    - offset:
        byte offset of opcode
    - opcode:
        opcode integer value
    - arg:
        instruction arg
    - lineno:
        -1 means unknown
    """

    __slots__ = "offset", "next", "opcode", "opname", "arg", "lineno"

    def __init__(self, offset, opcode, arg, nextoffset):
        self.offset = offset
        self.next = nextoffset
        self.opcode = opcode
        self.opname = dis.opname[opcode]
        self.arg = arg
        self.lineno = -1  # unknown line number

    @property
    def is_jump(self):
        return self.opcode in JUMP_OPS

    @property
    def is_terminator(self):
        return self.opcode in TERM_OPS

    def get_jump_target(self):
        # With Python 3.10 the addressing of "bytecode" instructions has
        # changed from using bytes to using 16-bit words instead. As a
        # consequence the code to determine where a jump will lead had to be
        # adapted.
        # See also:
        # https://bugs.python.org/issue26647
        # https://bugs.python.org/issue27129
        # https://github.com/python/cpython/pull/25069
        assert self.is_jump
        if PYVERSION in ((3, 13), (3, 14)):
            if self.opcode in (
                dis.opmap[k] for k in ["JUMP_BACKWARD", "JUMP_BACKWARD_NO_INTERRUPT"]
            ):
                return self.next - (self.arg * 2)
        elif PYVERSION in ((3, 12),):
            if self.opcode in (dis.opmap[k] for k in ["JUMP_BACKWARD"]):
                return self.offset - (self.arg - 1) * 2
        elif PYVERSION in ((3, 11),):
            if self.opcode in (
                dis.opmap[k]
                for k in (
                    "JUMP_BACKWARD",
                    "POP_JUMP_BACKWARD_IF_TRUE",
                    "POP_JUMP_BACKWARD_IF_FALSE",
                    "POP_JUMP_BACKWARD_IF_NONE",
                    "POP_JUMP_BACKWARD_IF_NOT_NONE",
                )
            ):
                return self.offset - (self.arg - 1) * 2
        elif PYVERSION in (
            (3, 9),
            (3, 10),
        ):
            pass
        else:
            raise NotImplementedError(PYVERSION)

        if PYVERSION in ((3, 10), (3, 11), (3, 12), (3, 13), (3, 14)):
            if self.opcode in JREL_OPS:
                return self.next + self.arg * 2
            else:
                assert self.opcode in JABS_OPS
                return self.arg * 2 - 2
        elif PYVERSION in ((3, 9),):
            if self.opcode in JREL_OPS:
                return self.next + self.arg
            else:
                assert self.opcode in JABS_OPS
                return self.arg
        else:
            raise NotImplementedError(PYVERSION)

    def __repr__(self):
        return "%s(arg=%s, lineno=%d)" % (self.opname, self.arg, self.lineno)

    @property
    def block_effect(self):
        """Effect of the block stack
        Returns +1 (push), 0 (none) or -1 (pop)
        """
        if self.opname.startswith("SETUP_"):
            return 1
        elif self.opname == "POP_BLOCK":
            return -1
        else:
            return 0


CODE_LEN = 1
ARG_LEN = 1
NO_ARG_LEN = 1

OPCODE_NOP = dis.opname.index("NOP")


if PYVERSION in ((3, 13), (3, 14)):

    def _unpack_opargs(code):
        buf = [(start_offset, op, arg) for _, start_offset, op, arg in dis._unpack_opargs(code)]
        for i, (start_offset, op, arg) in enumerate(buf):
            if i + 1 < len(buf):
                next_offset = buf[i + 1][0]
            else:
                next_offset = len(code)
            yield (start_offset, op, arg, next_offset)

elif PYVERSION in ((3, 9), (3, 10), (3, 11), (3, 12)):
    # Adapted from Lib/dis.py
    def _unpack_opargs(code):
        """
        Returns a 4-int-tuple of
        (bytecode offset, opcode, argument, offset of next bytecode).
        """
        extended_arg = 0
        n = len(code)
        offset = i = 0
        while i < n:
            op = code[i]
            i += CODE_LEN
            if op >= HAVE_ARGUMENT:
                arg = code[i] | extended_arg
                for j in range(ARG_LEN):
                    arg |= code[i + j] << (8 * j)
                i += ARG_LEN
                if PYVERSION in ((3, 12),):
                    # Python 3.12 introduced cache slots. We need to account for
                    # cache slots when we determine the offset of the next
                    # opcode. The number of cache slots is specific to each
                    # opcode and can be looked up in the _inline_cache_entries
                    # dictionary.
                    i += _inline_cache_entries[op] * INSTR_LEN
                elif PYVERSION in ((3, 9), (3, 10), (3, 11)):
                    pass
                else:
                    raise NotImplementedError(PYVERSION)
                if op == EXTENDED_ARG:
                    # This is a deviation from what dis does...
                    # In python 3.11 it seems like EXTENDED_ARGs appear more
                    # often and are also used as jump targets. So as to not have
                    # to do "book keeping" for where EXTENDED_ARGs have been
                    # "skipped" they are replaced with NOPs so as to provide a
                    # legal jump target and also ensure that the bytecode
                    # offsets are correct.
                    yield (offset, OPCODE_NOP, arg, i)
                    extended_arg = arg << 8 * ARG_LEN
                    offset = i
                    continue
            else:
                arg = None
                i += NO_ARG_LEN
                if PYVERSION in ((3, 12),):
                    # Python 3.12 introduced cache slots. We need to account for
                    # cache slots when we determine the offset of the next
                    # opcode. The number of cache slots is specific to each
                    # opcode and can be looked up in the _inline_cache_entries
                    # dictionary.
                    i += _inline_cache_entries[op] * INSTR_LEN
                elif PYVERSION in ((3, 9), (3, 10), (3, 11)):
                    pass
                else:
                    raise NotImplementedError(PYVERSION)

            extended_arg = 0
            yield (offset, op, arg, i)
            offset = i  # Mark inst offset at first extended

else:
    raise NotImplementedError(PYVERSION)


def _patched_opargs(bc_stream):
    """Patch the bytecode stream.

    - Adds a NOP bytecode at the start to avoid jump target being at the entry.
    """
    # Injected NOP
    yield (0, OPCODE_NOP, None, _FIXED_OFFSET)
    # Adjust bytecode offset for the rest of the stream
    for offset, opcode, arg, nextoffset in bc_stream:
        # If the opcode has an absolute jump target, adjust it.
        if opcode in JABS_OPS:
            arg += _FIXED_OFFSET
        yield offset + _FIXED_OFFSET, opcode, arg, nextoffset + _FIXED_OFFSET


class ByteCodeIter:
    def __init__(self, code):
        self.code = code
        self.iter = iter(_patched_opargs(_unpack_opargs(self.code.co_code)))

    def __iter__(self):
        return self

    def _fetch_opcode(self):
        return next(self.iter)

    def next(self):
        offset, opcode, arg, nextoffset = self._fetch_opcode()
        return offset, ByteCodeInst(offset=offset, opcode=opcode, arg=arg, nextoffset=nextoffset)

    __next__ = next

    def read_arg(self, size):
        buf = 0
        for i in range(size):
            _offset, byte = next(self.iter)
            buf |= byte << (8 * i)
        return buf


class _ByteCode:
    """
    The decoded bytecode of a function, and related information.
    """

    __slots__ = (
        "func_id",
        "co_names",
        "co_varnames",
        "co_consts",
        "co_cellvars",
        "co_freevars",
        "exception_entries",
        "table",
        "labels",
    )

    def __init__(self, func_id):
        code = func_id.code

        labels = set(x + _FIXED_OFFSET for x in dis.findlabels(code.co_code))
        labels.add(0)

        # A map of {offset: ByteCodeInst}
        table = OrderedDict(ByteCodeIter(code))
        self._compute_lineno(table, code)

        self.func_id = func_id
        self.co_names = code.co_names
        self.co_varnames = code.co_varnames
        self.co_consts = code.co_consts
        self.co_cellvars = code.co_cellvars
        self.co_freevars = code.co_freevars

        self.table = table
        self.labels = sorted(labels)

    @classmethod
    def _compute_lineno(cls, table, code):
        """
        Compute the line numbers for all bytecode instructions.
        """
        for offset, lineno in dis.findlinestarts(code):
            adj_offset = offset + _FIXED_OFFSET
            if adj_offset in table:
                table[adj_offset].lineno = lineno
        # Assign unfilled lineno
        # Start with first bytecode's lineno
        known = code.co_firstlineno
        for inst in table.values():
            if inst.lineno is not None and inst.lineno >= 0:
                known = inst.lineno
            else:
                inst.lineno = known
        return table

    def __iter__(self):
        return iter(self.table.values())

    def __getitem__(self, offset):
        return self.table[offset]

    def __contains__(self, offset):
        return offset in self.table

    def dump(self):
        def label_marker(i):
            if i[1].offset in self.labels:
                return ">"
            else:
                return " "

        return "\n".join(
            "%s %10s\t%s" % ((label_marker(i),) + i)
            for i in self.table.items()
            if i[1].opname != "CACHE"
        )

    @classmethod
    def _compute_used_globals(cls, func, table, co_consts, co_names):
        """
        Compute the globals used by the function with the given
        bytecode table.
        """
        d = {}
        globs = func.__globals__
        builtins = globs.get("__builtins__", utils.builtins)
        if isinstance(builtins, ModuleType):
            builtins = builtins.__dict__
        # Look for LOAD_GLOBALs in the bytecode
        for inst in table.values():
            if inst.opname == "LOAD_GLOBAL":
                name = co_names[_fix_LOAD_GLOBAL_arg(inst.arg)]
                if name not in d:
                    try:
                        value = globs[name]
                    except KeyError:
                        value = builtins[name]
                    d[name] = value
        # Add globals used by any nested code object
        for co in co_consts:
            if isinstance(co, CodeType):
                subtable = OrderedDict(ByteCodeIter(co))
                d.update(cls._compute_used_globals(func, subtable, co.co_consts, co.co_names))
        return d

    def get_used_globals(self):
        """
        Get a {name: value} map of the globals used by this code
        object and any nested code objects.
        """
        return self._compute_used_globals(
            self.func_id.func, self.table, self.co_consts, self.co_names
        )


def _fix_LOAD_GLOBAL_arg(arg):
    if PYVERSION in ((3, 11), (3, 12), (3, 13), (3, 14)):
        return arg >> 1
    elif PYVERSION in (
        (3, 9),
        (3, 10),
    ):
        return arg
    else:
        raise NotImplementedError(PYVERSION)


class ByteCodePy311(_ByteCode):
    def __init__(self, func_id):
        super().__init__(func_id)
        entries = dis.Bytecode(func_id.code).exception_entries
        self.exception_entries = tuple(map(self.fixup_eh, entries))

    @staticmethod
    def fixup_eh(ent):
        # Patch up the exception table offset
        # because we add a NOP in _patched_opargs
        out = dis._ExceptionTableEntry(
            start=ent.start + _FIXED_OFFSET,
            end=ent.end + _FIXED_OFFSET,
            target=ent.target + _FIXED_OFFSET,
            depth=ent.depth,
            lasti=ent.lasti,
        )
        return out

    def find_exception_entry(self, offset):
        """
        Returns the exception entry for the given instruction offset
        """
        candidates = [
            (ent.depth, ent) for ent in self.exception_entries if ent.start <= offset < ent.end
        ]
        if candidates:
            _, ent = max(candidates)
            return ent


class ByteCodePy312(ByteCodePy311):
    def __init__(self, func_id):
        super().__init__(func_id)

        # initialize lazy property
        self._ordered_offsets = None

        # Fixup offsets for all exception entries.
        entries = [self.fixup_eh(e) for e in dis.Bytecode(func_id.code).exception_entries]

        # Remove exceptions, innermost ones first
        # Can be done by using a stack
        entries = self.remove_build_list_swap_pattern(entries)

        # If this is a generator, we need to skip any exception table entries
        # that point to the exception handler with the highest offset.
        if func_id.is_generator:
            # Get the exception handler with the highest offset.
            max_exception_target = max([e.target for e in entries])
            # Remove any exception table entries that point to that exception
            # handler.
            entries = [e for e in entries if e.target != max_exception_target]

        self.exception_entries = tuple(entries)

    @property
    def ordered_offsets(self):
        if not self._ordered_offsets:
            # Get an ordered list of offsets.
            self._ordered_offsets = [o for o in self.table]
        return self._ordered_offsets

    def remove_build_list_swap_pattern(self, entries):
        """Find the following bytecode pattern:

        BUILD_{LIST, MAP, SET}
        SWAP(2)
        FOR_ITER
        ...
        END_FOR
        SWAP(2)

        This pattern indicates that a list/dict/set comprehension has
        been inlined. In this case we can skip the exception blocks
        entirely along with the dead exceptions that it points to.
        A pair of exception that sandwiches these exception will
        also be merged into a single exception.

        Update for Python 3.13, the ending of the pattern has a extra
        POP_TOP:

        ...
        END_FOR
        POP_TOP
        SWAP(2)

        Update for Python 3.13.1, there's now a GET_ITER before FOR_ITER.
        This patch the GET_ITER to NOP to minimize changes downstream
        (e.g. array-comprehension).
        """

        def pop_and_merge_exceptions(entries: list, entry_to_remove: _ExceptionTableEntry):
            lower_entry_idx = entries.index(entry_to_remove) - 1
            upper_entry_idx = entries.index(entry_to_remove) + 1

            # Merge the upper and lower exceptions if possible.
            if lower_entry_idx >= 0 and upper_entry_idx < len(entries):
                lower_entry = entries[lower_entry_idx]
                upper_entry = entries[upper_entry_idx]
                if lower_entry.target == upper_entry.target:
                    entries[lower_entry_idx] = _ExceptionTableEntry(
                        lower_entry.start,
                        upper_entry.end,
                        lower_entry.target,
                        lower_entry.depth,
                        upper_entry.lasti,
                    )
                    entries.remove(upper_entry)

            # Remove the exception entry.
            entries.remove(entry_to_remove)
            # Remove dead exceptions, if any, that the entry above may point to.
            entries = [e for e in entries if not e.start == entry_to_remove.target]
            return entries

        change_to_nop = set()
        work_remaining = True
        while work_remaining:
            # Temporarily set work_remaining to False, if we find a pattern
            # then work is not complete, hence we set it again to True.
            work_remaining = False
            current_nop_fixes = set()
            for entry in entries.copy():
                # Check start of pattern, three instructions.
                # Work out the index of the instruction.
                index = self.ordered_offsets.index(entry.start)
                # If there is a BUILD_{LIST, MAP, SET} instruction at this
                # location.
                curr_inst = self.table[self.ordered_offsets[index]]
                if curr_inst.opname not in (
                    "BUILD_LIST",
                    "BUILD_MAP",
                    "BUILD_SET",
                ):
                    continue
                # Check if the BUILD_{LIST, MAP, SET} instruction is followed
                # by a SWAP(2).
                next_inst = self.table[self.ordered_offsets[index + 1]]
                if not next_inst.opname == "SWAP" and next_inst.arg == 2:
                    continue
                next_inst = self.table[self.ordered_offsets[index + 2]]
                # Check if the SWAP is followed by a FOR_ITER
                # BUT Python3.13.1 introduced an extra GET_ITER.
                # If we see a GET_ITER here, check if the next thing is a
                # FOR_ITER.
                if next_inst.opname == "GET_ITER":
                    # Add the inst to potentially be replaced to NOP
                    current_nop_fixes.add(next_inst)
                    # Loop up next instruction.
                    next_inst = self.table[self.ordered_offsets[index + 3]]

                if not next_inst.opname == "FOR_ITER":
                    continue

                if PYVERSION in ((3, 13), (3, 14)):
                    # Check end of pattern, two instructions.
                    # Check for the corresponding END_FOR, exception table end
                    # is non-inclusive, so subtract one.
                    index = self.ordered_offsets.index(entry.end)
                    curr_inst = self.table[self.ordered_offsets[index - 2]]
                    if not curr_inst.opname == "END_FOR":
                        continue
                    next_inst = self.table[self.ordered_offsets[index - 1]]
                    if PYVERSION in ((3, 13),):
                        if not next_inst.opname == "POP_TOP":
                            continue
                    elif PYVERSION in ((3, 14),):
                        if not next_inst.opname == "POP_ITER":
                            continue
                    else:
                        raise NotImplementedError(PYVERSION)
                    # END_FOR must be followed by SWAP(2)
                    next_inst = self.table[self.ordered_offsets[index]]
                    if not next_inst.opname == "SWAP" and next_inst.arg == 2:
                        continue
                elif PYVERSION in ((3, 9), (3, 10), (3, 11), (3, 12)):
                    # Check end of pattern, two instructions.
                    # Check for the corresponding END_FOR, exception table end
                    # is non-inclusive, so subtract one.
                    index = self.ordered_offsets.index(entry.end)
                    curr_inst = self.table[self.ordered_offsets[index - 1]]
                    if not curr_inst.opname == "END_FOR":
                        continue
                    # END_FOR must be followed by SWAP(2)
                    next_inst = self.table[self.ordered_offsets[index]]
                    if not next_inst.opname == "SWAP" and next_inst.arg == 2:
                        continue
                else:
                    raise NotImplementedError(PYVERSION)
                # If all conditions are met that means this exception entry
                # is for a list/dict/set comprehension and can be removed.
                # Also if there exist exception entries above and below this
                # entry pointing to the same target. those can be merged into
                # a single bigger exception block.
                entries = pop_and_merge_exceptions(entries, entry)
                work_remaining = True

                # Commit NOP fixes since we confirmed the suspects belong to
                # a comprehension code.
                change_to_nop |= current_nop_fixes

        # Complete fixes to NOPs
        for inst in change_to_nop:
            self.table[inst.offset] = ByteCodeInst(inst.offset, dis.opmap["NOP"], None, inst.next)
        return entries


if PYVERSION == (3, 11):
    ByteCode = ByteCodePy311
elif PYVERSION in (
    (3, 12),
    (3, 13),
    (3, 14),
):
    ByteCode = ByteCodePy312
elif PYVERSION < (3, 11):
    ByteCode = _ByteCode
else:
    raise NotImplementedError(PYVERSION)


class FunctionIdentity(serialize.ReduceMixin):
    """
    A function's identity and metadata.

    Note this typically represents a function whose bytecode is
    being compiled, not necessarily the top-level user function
    (the two might be distinct).
    """

    _unique_ids = _LockedCounter(1)

    @classmethod
    def from_function(cls, pyfunc):
        """
        Create the FunctionIdentity of the given function.
        """
        func = get_function_object(pyfunc)
        code = get_code_object(func)
        pysig = utils.pysignature(func)
        if not code:
            raise errors.ByteCodeSupportError("%s does not provide its bytecode" % func)

        try:
            func_qualname = func.__qualname__
        except AttributeError:
            func_qualname = func.__name__

        self = cls()
        self.func = func
        self.func_qualname = func_qualname
        self.func_name = func_qualname.split(".")[-1]
        self.code = code
        self.module = inspect.getmodule(func)
        self.modname = utils._dynamic_modname if self.module is None else self.module.__name__
        self.is_generator = inspect.isgeneratorfunction(func)
        self.pysig = pysig
        self.filename = code.co_filename
        self.firstlineno = code.co_firstlineno
        self.arg_count = len(pysig.parameters)
        self.arg_names = list(pysig.parameters)

        # Even the same function definition can be compiled into
        # several different function objects with distinct closure
        # variables, so we make sure to disambiguate using an unique id.
        uid = next(cls._unique_ids)
        self.unique_name = "{}${}".format(self.func_qualname, uid)
        self.unique_id = uid

        return self

    def derive(self):
        """Copy the object and increment the unique counter."""
        return self.from_function(self.func)

    def _reduce_states(self):
        """
        NOTE: part of ReduceMixin protocol
        """
        return dict(pyfunc=self.func)

    @classmethod
    def _rebuild(cls, pyfunc):
        """
        NOTE: part of ReduceMixin protocol
        """
        return cls.from_function(pyfunc)
