#     Copyright 2023, Kay Hayen, mailto:kay.hayen@gmail.com
#
#     Part of "Nuitka", an optimizing Python compiler that is compatible and
#     integrates with CPython, but also works on its own.
#
#     Licensed under the Apache License, Version 2.0 (the "License");
#     you may not use this file except in compliance with the License.
#     You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#     Unless required by applicable law or agreed to in writing, software
#     distributed under the License is distributed on an "AS IS" BASIS,
#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#     See the License for the specific language governing permissions and
#     limitations under the License.
#
""" Frame codes

This is about frame stacks and their management. There are different kinds
of frames for different uses.
"""

from nuitka.PythonVersions import python_version
from nuitka.utils.Jinja2 import renderTemplateFromString

from .CodeHelpers import _generateStatementSequenceCode
from .Emission import SourceCodeCollector
from .ErrorCodes import getFrameVariableTypeDescriptionCode
from .ExceptionCodes import getTracebackMakingIdentifier
from .Indentation import indented
from .LabelCodes import getGotoCode, getLabelCode
from .ModuleCodes import getModuleAccessCode
from .templates.CodeTemplatesFrames import (
    template_frame_attach_locals,
    template_frame_guard_generator_exception_handler,
    template_frame_guard_generator_return_handler,
    template_frame_guard_normal_exception_handler,
    template_frame_guard_normal_main_block,
    template_frame_guard_normal_return_handler,
)


def getFrameLocalsStorageSize(type_descriptions):
    candidates = set()

    for type_description in type_descriptions:
        candidate = "+".join(
            getTypeSizeOf(type_indicator) for type_indicator in sorted(type_description)
        )

        candidates.add(candidate)

    if not candidates:
        return "0"

    candidates = list(sorted(candidates))
    result = candidates.pop()

    while candidates:
        # assert False, (type_descriptions, context.frame_variables_stack[-1])

        result = "MAX(%s, %s)" % (result, candidates.pop())

    return result


def _searchLocalVariableByName(local_variables, variable_name):
    for local_variable in local_variables:
        if local_variable.getName() == variable_name:
            return local_variable

    return None


def generateStatementsFrameCode(statement_sequence, emit, context):
    # This is a wrapper that provides also handling of frames. The standard
    # and generator frame variety ought to be merged once generators are
    # possible to inline.
    # TODO: This will get simpler, pylint: disable=too-many-locals
    context.pushCleanupScope()

    guard_mode = statement_sequence.getGuardMode()

    code_object = statement_sequence.getCodeObject()
    code_identifier = context.getCodeObjectHandle(code_object=code_object)

    parent_exception_exit = context.getExceptionEscape()

    # For nested frames, make sure to restore set the type description.
    if context.getFrameHandle() is not None:
        real_parent_exception_exit = parent_exception_exit
        parent_exception_exit = context.allocateLabel("nested_frame_exit")

    # Allow stacking of frame handles.
    context.pushFrameHandle(code_identifier, statement_sequence.hasStructureMember())

    context.setExceptionEscape(context.allocateLabel("frame_exception_exit"))

    needs_preserve = statement_sequence.needsFrameExceptionPreserving()

    if statement_sequence.mayReturn():
        parent_return_exit = context.getReturnTarget()

        context.setReturnTarget(context.allocateLabel("frame_return_exit"))
    else:
        parent_return_exit = None

    # We need to define that all the variables needs to be pushed. We do not
    # have a flag that says "always NULL" for variables. With efficient NULL
    # passing however (not at all, TODO), that doesn't matter much.
    provider = statement_sequence.getParentVariableProvider()
    local_variables = provider.getLocalVariables()

    context.pushFrameVariables(
        tuple(
            _searchLocalVariableByName(local_variables, variable_name)
            for variable_name in code_object.getVarNames()
        )
    )

    # Now generate the statements code into a local buffer, to we can wrap
    # the frame stuff around it.
    local_emit = SourceCodeCollector()

    _generateStatementSequenceCode(
        statement_sequence=statement_sequence, emit=local_emit, context=context
    )

    if statement_sequence.mayRaiseException(BaseException):
        frame_exception_exit = context.getExceptionEscape()
    else:
        frame_exception_exit = None

    if parent_return_exit is not None:
        frame_return_exit = context.getReturnTarget()
    else:
        frame_return_exit = None

    if guard_mode == "generator":
        # TODO: This case should also care about "needs_preserve", as for
        # Python3 it is actually not a stub of empty code.

        getFrameGuardGeneratorCode(
            frame_node=statement_sequence,
            code_identifier=code_identifier,
            codes=local_emit.codes,
            parent_exception_exit=parent_exception_exit,
            frame_exception_exit=frame_exception_exit,
            parent_return_exit=parent_return_exit,
            frame_return_exit=frame_return_exit,
            emit=emit,
            context=context,
        )
    elif guard_mode in ("full", "once"):
        getFrameGuardHeavyCode(
            frame_node=statement_sequence,
            code_identifier=code_identifier,
            parent_exception_exit=parent_exception_exit,
            parent_return_exit=parent_return_exit,
            frame_exception_exit=frame_exception_exit,
            frame_return_exit=frame_return_exit,
            codes=local_emit.codes,
            needs_preserve=needs_preserve,
            emit=emit,
            context=context,
        )
    else:
        assert False, guard_mode

    context.popFrameVariables()
    context.popFrameHandle()

    # For nested frames, make sure to restore set the type description.
    if context.getFrameHandle() is not None:
        label = context.allocateLabel("skip_nested_handling")
        getGotoCode(label, emit)
        getLabelCode(parent_exception_exit, emit)
        emit(getFrameVariableTypeDescriptionCode(context))
        getGotoCode(real_parent_exception_exit, emit)
        getLabelCode(label, emit)

        parent_exception_exit = real_parent_exception_exit

    context.popCleanupScope()

    context.setExceptionEscape(parent_exception_exit)

    if frame_return_exit is not None:
        context.setReturnTarget(parent_return_exit)


def getTypeSizeOf(type_indicator):
    if type_indicator in ("O", "o", "N", "c"):
        return "sizeof(void *)"
    elif type_indicator == "b":
        return "sizeof(nuitka_bool)"
    elif type_indicator == "L":
        return "sizeof(nuitka_ilong)"
    else:
        assert False, type_indicator


def getFrameAttachLocalsCode(context, frame_identifier):
    frame_variable_codes = context.getFrameVariableCodeNames()
    frame_variable_codes = ",\n    ".join(frame_variable_codes)
    if frame_variable_codes:
        frame_variable_codes = ",\n    " + frame_variable_codes

    return template_frame_attach_locals % {
        "frame_identifier": frame_identifier,
        "type_description": context.getFrameTypeDescriptionDeclaration(),
        "frame_variable_refs": frame_variable_codes,
    }


def getFrameGuardHeavyCode(
    frame_node,
    code_identifier,
    codes,
    parent_exception_exit,
    parent_return_exit,
    frame_exception_exit,
    frame_return_exit,
    needs_preserve,
    emit,
    context,
):
    # We really need this many parameters here and it gets very
    # detail rich, pylint: disable=too-many-locals

    no_exception_exit = context.allocateLabel("frame_no_exception")

    frame_identifier = context.getFrameHandle()
    if frame_node.getGuardMode() == "full":
        frame_cache_identifier = context.variable_storage.addFrameCacheDeclaration(
            frame_identifier.code_name
        )
    else:
        frame_cache_identifier = None

    frame_init_code = ""
    frame_exit_code = ""

    # Expose the locals dictionary with the frame locals if it exists.
    if frame_node.isStatementsFrameClass():
        attach_locals_code = getFrameAttachLocalsCode(context, frame_identifier)

        use_locals_dict = (
            frame_node.getLocalsScope().getCodeName() in context.getLocalsDictNames()
        )

        make_frame_code = (
            """MAKE_CLASS_FRAME(tstate, %(code_identifier)s, %(module_identifier)s, %(locals_dict_name)s, %(locals_size)s)"""
            % {
                "code_identifier": code_identifier,
                "module_identifier": getModuleAccessCode(context),
                "locals_size": getFrameLocalsStorageSize(
                    context.getFrameVariableTypeDescriptions()
                ),
                "locals_dict_name": frame_node.getLocalsScope().getCodeName()
                if use_locals_dict
                else "NULL",
            }
        )

        if use_locals_dict:
            # TODO: Have a function that clears an optional existing reference and assigns a new
            # value taking a reference.
            frame_init_code = """\
assert(%(frame_identifier)s->m_frame.f_locals == NULL);
%(frame_identifier)s->m_frame.f_locals = %(locals_dict)s;
Py_INCREF(%(frame_identifier)s->m_frame.f_locals);
""" % {
                "frame_identifier": frame_identifier,
                "locals_dict": frame_node.getLocalsScope().getCodeName(),
            }
            frame_exit_code = """
Py_CLEAR(%(frame_identifier)s->m_frame.f_locals);
""" % {
                "frame_identifier": frame_identifier,
            }
    elif frame_node.isStatementsFrameFunction():
        attach_locals_code = getFrameAttachLocalsCode(context, frame_identifier)

        make_frame_code = (
            """MAKE_FUNCTION_FRAME(tstate, %(code_identifier)s, %(module_identifier)s, %(locals_size)s)"""
            % {
                "code_identifier": code_identifier,
                "module_identifier": getModuleAccessCode(context),
                "locals_size": getFrameLocalsStorageSize(
                    context.getFrameVariableTypeDescriptions()
                ),
            }
        )
    elif frame_node.isStatementsFrameModule():
        attach_locals_code = ""
        make_frame_code = (
            """MAKE_MODULE_FRAME(%(code_identifier)s, %(module_identifier)s)"""
            % {
                "code_identifier": code_identifier,
                "module_identifier": getModuleAccessCode(context),
            }
        )
    else:
        assert False, frame_node

    context_identifier = frame_node.getStructureMember()

    emit(
        renderTemplateFromString(
            template_frame_guard_normal_main_block,
            frame_identifier=frame_identifier,
            frame_cache_identifier=frame_cache_identifier,
            codes=indented(codes, 0),
            no_exception_exit=no_exception_exit,
            needs_preserve=needs_preserve,
            make_frame_code=make_frame_code,
            frame_init_code=frame_init_code,
            frame_exit_code=frame_exit_code,
            context_identifier=context_identifier,
            is_python34_or_later=python_version >= 0x340,
            is_python3=python_version >= 0x300,
        )
    )

    if frame_return_exit is not None:
        emit(
            renderTemplateFromString(
                template_frame_guard_normal_return_handler,
                frame_identifier=frame_identifier,
                return_exit=parent_return_exit,
                frame_return_exit=frame_return_exit,
                needs_preserve=needs_preserve,
                frame_exit_code=frame_exit_code,
            )
        )

    if frame_exception_exit is not None:
        (
            _exception_type,
            _exception_value,
            exception_tb,
            exception_lineno,
        ) = context.variable_storage.getExceptionVariableDescriptions()

        emit(
            renderTemplateFromString(
                template_frame_guard_normal_exception_handler,
                frame_identifier=context.getFrameHandle(),
                frame_cache_identifier=frame_cache_identifier,
                tb_making_code=getTracebackMakingIdentifier(
                    context=context, lineno_name=exception_lineno
                ),
                attach_locals_code=attach_locals_code,
                parent_exception_exit=parent_exception_exit,
                frame_exception_exit=frame_exception_exit,
                frame_exit_code=frame_exit_code,
                needs_preserve=needs_preserve,
                exception_tb=exception_tb,
                exception_lineno=exception_lineno,
            )
        )

    getLabelCode(no_exception_exit, emit)


def getFrameGuardGeneratorCode(
    frame_node,
    code_identifier,
    codes,
    parent_exception_exit,
    parent_return_exit,
    frame_exception_exit,
    frame_return_exit,
    emit,
    context,
):
    # We really need this many parameters here and it gets very
    # detail rich, pylint: disable=too-many-locals
    (
        exception_type,
        _exception_value,
        exception_tb,
        exception_lineno,
    ) = context.variable_storage.getExceptionVariableDescriptions()

    context_identifier = frame_node.getStructureMember()

    no_exception_exit = context.allocateLabel("frame_no_exception")

    frame_identifier = context.getFrameHandle()
    frame_cache_identifier = context.variable_storage.addFrameCacheDeclaration(
        frame_identifier.code_name
    )

    make_frame_code = (
        """MAKE_FUNCTION_FRAME(tstate, %(code_identifier)s, %(module_identifier)s, %(locals_size)s)"""
        % {
            "code_identifier": code_identifier,
            "module_identifier": getModuleAccessCode(context),
            "locals_size": getFrameLocalsStorageSize(
                context.getFrameVariableTypeDescriptions()
            ),
        }
    )
    is_generator = True

    frame_init_code = ""
    frame_exit_code = ""

    emit(
        renderTemplateFromString(
            template_frame_guard_normal_main_block,
            frame_identifier=frame_identifier,
            frame_cache_identifier=frame_cache_identifier,
            context_identifier=context_identifier,
            codes=indented(codes, 0),
            no_exception_exit=no_exception_exit,
            needs_preserve=False,  # TODO: Clears stuff
            make_frame_code=make_frame_code,
            frame_init_code=frame_init_code,
            frame_exit_code=frame_exit_code,
            is_generator=is_generator,
            is_python34_or_later=python_version >= 0x340,
            is_python3=python_version >= 0x300,
        )
    )

    if frame_return_exit is not None:
        emit(
            template_frame_guard_generator_return_handler
            % {
                "context_identifier": context_identifier,
                "return_exit": parent_return_exit,
                "frame_return_exit": frame_return_exit,
            }
        )

    if frame_exception_exit is not None:
        emit(
            template_frame_guard_generator_exception_handler
            % {
                "context_identifier": context_identifier,
                "frame_identifier": frame_identifier,
                "frame_cache_identifier": frame_cache_identifier,
                "exception_type": exception_type,
                "exception_tb": exception_tb,
                "exception_lineno": exception_lineno,
                "tb_making": getTracebackMakingIdentifier(
                    context=context, lineno_name=exception_lineno
                ),
                "attach_locals": indented(
                    getFrameAttachLocalsCode(context, frame_identifier)
                ),
                "frame_exception_exit": frame_exception_exit,
                "parent_exception_exit": parent_exception_exit,
            }
        )

    getLabelCode(no_exception_exit, emit)


def generateFramePreserveExceptionCode(statement, emit, context):
    if python_version < 0x300:
        emit("// Preserve existing published exception.")

        emit(
            "PRESERVE_FRAME_EXCEPTION(tstate, %(frame_identifier)s);"
            % {"frame_identifier": context.getFrameHandle()}
        )
    else:
        preserver_id = statement.getPreserverId()

        assert preserver_id != 0, statement
        exception_preserved = context.addExceptionPreserverVariables(preserver_id)

        emit(
            """\
// Preserve existing published exception id %(preserver_id)d.
%(exception_preserved)s = GET_CURRENT_EXCEPTION(tstate);
"""
            % {
                "exception_preserved": exception_preserved,
                "preserver_id": preserver_id,
            }
        )


def generateFrameRestoreExceptionCode(statement, emit, context):
    if python_version < 0x300:
        emit(
            """\
// Restore previous exception.
RESTORE_FRAME_EXCEPTION(%(frame_identifier)s);"""
            % {"frame_identifier": context.getFrameHandle()}
        )
    else:
        preserver_id = statement.getPreserverId()

        exception_preserved = context.addExceptionPreserverVariables(preserver_id)

        emit(
            """\
// Restore previous exception id %(preserver_id)d.
SET_CURRENT_EXCEPTION(tstate, &%(exception_preserved)s);
"""
            % {
                "exception_preserved": exception_preserved,
                "preserver_id": preserver_id,
            }
        )
