#     Copyright 2018, 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:
        emit(
            "SAVE_%s_EXCEPTION( %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 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
    emit(
        "%s = %s;" % (
            to_name,
            yield_return_name
        )
    )

    if preserve_exception:
        emit(
            "RESTORE_%s_EXCEPTION( %s );" % (
                context.getContextObjectName().upper(),
                context.getContextObjectName()
            )
        )


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 == "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_yieldfrom = %(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
        )

        context.addCleanupTempName(result_name)


def generateYieldFromWaitableCode(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_yieldfrom = %(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
        )

        # TODO: Seems to be redundant with and _getYieldPreserveCode doing
        # it and could be removed
        getErrorExitCode(
            check_name = result_name,
            emit       = emit,
            context    = context
        )

        context.addCleanupTempName(result_name)


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
