import opcode
import sys
from enum import IntFlag
from typing import Optional, Union

# alias to keep the 'bytecode' variable free
import bytecode as _bytecode


class CompilerFlags(IntFlag):
    """Possible values of the co_flags attribute of Code object.

    Note: We do not rely on inspect values here as some of them are missing and
    furthermore would be version dependent.

    """

    OPTIMIZED = 0x00001
    NEWLOCALS = 0x00002
    VARARGS = 0x00004
    VARKEYWORDS = 0x00008
    NESTED = 0x00010
    GENERATOR = 0x00020
    NOFREE = 0x00040
    # New in Python 3.5
    # Used for coroutines defined using async def ie native coroutine
    COROUTINE = 0x00080
    # Used for coroutines defined as a generator and then decorated using
    # types.coroutine
    ITERABLE_COROUTINE = 0x00100
    # New in Python 3.6
    # Generator defined in an async def function
    ASYNC_GENERATOR = 0x00200

    # __future__ flags
    # future flags changed in Python 3.9
    if sys.version_info < (3, 9):
        FUTURE_GENERATOR_STOP = 0x80000
        FUTURE_ANNOTATIONS = 0x100000
    else:
        FUTURE_GENERATOR_STOP = 0x800000
        FUTURE_ANNOTATIONS = 0x1000000


def infer_flags(
    bytecode: Union[
        "_bytecode.Bytecode", "_bytecode.ConcreteBytecode", "_bytecode.ControlFlowGraph"
    ],
    is_async: Optional[bool] = None,
):
    """Infer the proper flags for a bytecode based on the instructions.

    Because the bytecode does not have enough context to guess if a function
    is asynchronous the algorithm tries to be conservative and will never turn
    a previously async code into a sync one.

    Parameters
    ----------
    bytecode : Bytecode | ConcreteBytecode | ControlFlowGraph
        Bytecode for which to infer the proper flags
    is_async : bool | None, optional
        Force the code to be marked as asynchronous if True, prevent it from
        being marked as asynchronous if False and simply infer the best
        solution based on the opcode and the existing flag if None.

    """
    flags = CompilerFlags(0)
    if not isinstance(
        bytecode,
        (_bytecode.Bytecode, _bytecode.ConcreteBytecode, _bytecode.ControlFlowGraph),
    ):
        msg = (
            "Expected a Bytecode, ConcreteBytecode or ControlFlowGraph "
            "instance not %s"
        )
        raise ValueError(msg % bytecode)

    instructions = (
        bytecode._get_instructions()
        if isinstance(bytecode, _bytecode.ControlFlowGraph)
        else bytecode
    )
    instr_names = {
        i.name
        for i in instructions
        if not isinstance(
            i,
            (
                _bytecode.SetLineno,
                _bytecode.Label,
                _bytecode.TryBegin,
                _bytecode.TryEnd,
            ),
        )
    }

    # Identify optimized code
    if not (instr_names & {"STORE_NAME", "LOAD_NAME", "DELETE_NAME"}):
        flags |= CompilerFlags.OPTIMIZED

    # Check for free variables
    if not (instr_names & {opcode.opname[i] for i in opcode.hasfree}):
        flags |= CompilerFlags.NOFREE

    # Copy flags for which we cannot infer the right value
    flags |= bytecode.flags & (
        CompilerFlags.NEWLOCALS
        | CompilerFlags.VARARGS
        | CompilerFlags.VARKEYWORDS
        | CompilerFlags.NESTED
    )

    sure_generator = instr_names & {"YIELD_VALUE"}
    maybe_generator = instr_names & {"YIELD_VALUE", "YIELD_FROM"}

    sure_async = instr_names & {
        "GET_AWAITABLE",
        "GET_AITER",
        "GET_ANEXT",
        "BEFORE_ASYNC_WITH",
        "SETUP_ASYNC_WITH",
        "END_ASYNC_FOR",
        "ASYNC_GEN_WRAP",  # New in 3.11
    }

    # If performing inference or forcing an async behavior, first inspect
    # the flags since this is the only way to identify iterable coroutines
    if is_async in (None, True):
        if bytecode.flags & CompilerFlags.COROUTINE:
            if sure_generator:
                flags |= CompilerFlags.ASYNC_GENERATOR
            else:
                flags |= CompilerFlags.COROUTINE
        elif bytecode.flags & CompilerFlags.ITERABLE_COROUTINE:
            if sure_async:
                msg = (
                    "The ITERABLE_COROUTINE flag is set but bytecode that"
                    "can only be used in async functions have been "
                    "detected. Please unset that flag before performing "
                    "inference."
                )
                raise ValueError(msg)
            flags |= CompilerFlags.ITERABLE_COROUTINE
        elif bytecode.flags & CompilerFlags.ASYNC_GENERATOR:
            if not sure_generator:
                flags |= CompilerFlags.COROUTINE
            else:
                flags |= CompilerFlags.ASYNC_GENERATOR

        # If the code was not asynchronous before determine if it should now be
        # asynchronous based on the opcode and the is_async argument.
        else:
            if sure_async:
                # YIELD_FROM is not allowed in async generator
                if sure_generator:
                    flags |= CompilerFlags.ASYNC_GENERATOR
                else:
                    flags |= CompilerFlags.COROUTINE

            elif maybe_generator:
                if is_async:
                    if sure_generator:
                        flags |= CompilerFlags.ASYNC_GENERATOR
                    else:
                        flags |= CompilerFlags.COROUTINE
                else:
                    flags |= CompilerFlags.GENERATOR

            elif is_async:
                flags |= CompilerFlags.COROUTINE

    # If the code should not be asynchronous, check first it is possible and
    # next set the GENERATOR flag if relevant
    else:
        if sure_async:
            raise ValueError(
                "The is_async argument is False but bytecodes "
                "that can only be used in async functions have "
                "been detected."
            )

        if maybe_generator:
            flags |= CompilerFlags.GENERATOR

    flags |= bytecode.flags & CompilerFlags.FUTURE_GENERATOR_STOP

    return flags