#     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.
#
""" Yield related codes.

The normal "yield", and the Python 3.3 or higher "yield from" variant.
"""

from .CodeHelpers import (
    generateChildExpressionsCode,
    withObjectCodeTemporaryAssignment,
)
from .ErrorCodes import getErrorExitCode
from .PythonAPICodes import getReferenceExportCode
from .VariableDeclarations import VariableDeclaration


def _getYieldPreserveCode(
    to_name, value_name, preserve_exception, yield_code, resume_code, emit, context
):
    yield_return_label = context.allocateLabel("yield_return")
    yield_return_index = yield_return_label.split("_")[-1]

    locals_preserved = context.variable_storage.getLocalPreservationDeclarations()

    # Need not preserve it, if we are not going to use it for the purpose
    # of releasing it.
    if type(value_name) is tuple:
        value_names = value_name
    else:
        value_names = (value_name,)

    for name in value_names:
        if not context.needsCleanup(name):
            locals_preserved.remove(name)

    # Target name is not assigned, no need to preserve it.
    if to_name in locals_preserved:
        locals_preserved.remove(to_name)

    if locals_preserved:
        yield_tmp_storage = context.variable_storage.getVariableDeclarationTop(
            "yield_tmps"
        )

        if yield_tmp_storage is None:
            yield_tmp_storage = context.variable_storage.addVariableDeclarationTop(
                "char[1024]", "yield_tmps", None
            )

        emit(
            "Nuitka_PreserveHeap(%s, %s, NULL);"
            % (
                yield_tmp_storage,
                ", ".join(
                    "&%s, sizeof(%s)" % (local_preserved, local_preserved.c_type)
                    for local_preserved in locals_preserved
                ),
            )
        )

    if preserve_exception:
        # Expands to SAVE_GENERATOR_EXCEPTION,SAVE_COROUTINE_EXCEPTION,SAVE_ASYNCGEN_EXCEPTION
        emit(
            "SAVE_%s_EXCEPTION(tstate, %s);"
            % (context.getContextObjectName().upper(), context.getContextObjectName())
        )

    emit(
        """\
%(context_object_name)s->m_yield_return_index = %(yield_return_index)s;"""
        % {
            "context_object_name": context.getContextObjectName(),
            "yield_return_index": yield_return_index,
        }
    )

    emit(yield_code)

    emit("%(yield_return_label)s:" % {"yield_return_label": yield_return_label})

    if preserve_exception:
        # Expands to RESTORE_GENERATOR_EXCEPTION,RESTORE_COROUTINE_EXCEPTION,RESTORE_ASYNCGEN_EXCEPTION
        emit(
            "RESTORE_%s_EXCEPTION(tstate, %s);"
            % (context.getContextObjectName().upper(), context.getContextObjectName())
        )

    if locals_preserved:
        emit(
            "Nuitka_RestoreHeap(%s, %s, NULL);"
            % (
                yield_tmp_storage,
                ", ".join(
                    "&%s, sizeof(%s)" % (local_preserved, local_preserved.c_type)
                    for local_preserved in locals_preserved
                ),
            )
        )

    if resume_code:
        emit(resume_code)

    yield_return_name = VariableDeclaration(
        "PyObject *", "yield_return_value", None, None
    )

    getErrorExitCode(check_name=yield_return_name, emit=emit, context=context)

    # Called with object, so we can simply do this.
    emit("%s = %s;" % (to_name, yield_return_name))
    context.addCleanupTempName(to_name)


def generateYieldCode(to_name, expression, emit, context):
    (value_name,) = generateChildExpressionsCode(
        expression=expression, emit=emit, context=context
    )

    # In handlers, we must preserve/restore the exception.
    preserve_exception = expression.isExceptionPreserving()

    getReferenceExportCode(value_name, emit, context)
    if context.needsCleanup(value_name):
        context.removeCleanupTempName(value_name)

    yield_code = "return %(yielded_value)s;" % {"yielded_value": value_name}

    with withObjectCodeTemporaryAssignment(
        to_name, "yield_result", expression, emit, context
    ) as result_name:
        _getYieldPreserveCode(
            to_name=result_name,
            value_name=value_name,
            yield_code=yield_code,
            resume_code=None,
            preserve_exception=preserve_exception,
            emit=emit,
            context=context,
        )

        # This conversion will not use it, and since it is borrowed, debug mode
        # would otherwise complain.
        if to_name.c_type == "nuitka_void":
            result_name.maybe_unused = True

        # Comes as only borrowed.
        # context.addCleanupTempName(result_name)


def generateYieldFromCode(to_name, expression, emit, context):
    (value_name,) = generateChildExpressionsCode(
        expression=expression, emit=emit, context=context
    )

    # In handlers, we must preserve/restore the exception.
    preserve_exception = expression.isExceptionPreserving()

    getReferenceExportCode(value_name, emit, context)
    if context.needsCleanup(value_name):
        context.removeCleanupTempName(value_name)

    yield_code = """\
generator->m_yield_from = %(yield_from)s;
return NULL;
""" % {
        "yield_from": value_name
    }

    with withObjectCodeTemporaryAssignment(
        to_name, "yieldfrom_result", expression, emit, context
    ) as result_name:
        _getYieldPreserveCode(
            to_name=result_name,
            value_name=value_name,
            yield_code=yield_code,
            resume_code=None,
            preserve_exception=preserve_exception,
            emit=emit,
            context=context,
        )


def generateYieldFromAwaitableCode(to_name, expression, emit, context):
    # In handlers, we must preserve/restore the exception.
    preserve_exception = expression.isExceptionPreserving()

    (awaited_name,) = generateChildExpressionsCode(
        expression=expression, emit=emit, context=context
    )

    yield_code = """\
%(object_name)s->m_yield_from = %(yield_from)s;
%(object_name)s->m_awaiting = true;
return NULL;
""" % {
        "object_name": context.getContextObjectName(),
        "yield_from": awaited_name,
    }

    resume_code = """\
%(object_name)s->m_awaiting = false;
""" % {
        "object_name": context.getContextObjectName()
    }

    getReferenceExportCode(awaited_name, emit, context)
    if context.needsCleanup(awaited_name):
        context.removeCleanupTempName(awaited_name)

    with withObjectCodeTemporaryAssignment(
        to_name, "await_result", expression, emit, context
    ) as result_name:
        _getYieldPreserveCode(
            to_name=result_name,
            value_name=awaited_name,
            yield_code=yield_code,
            resume_code=resume_code,
            preserve_exception=preserve_exception,
            emit=emit,
            context=context,
        )


def getYieldReturnDispatchCode(context):
    function_dispatch = [
        "case %(index)d: goto yield_return_%(index)d;" % {"index": yield_index}
        for yield_index in range(context.getLabelCount("yield_return"), 0, -1)
    ]

    if function_dispatch:
        function_dispatch.insert(
            0, "switch(%s->m_yield_return_index) {" % context.getContextObjectName()
        )
        function_dispatch.append("}")

    return function_dispatch
