#     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.
#
""" Reformulation of contraction expressions.

Consult the developer manual for information. TODO: Add ability to sync
source code comments with developer manual sections.

"""

from nuitka.__past__ import intern  # pylint: disable=I0021,redefined-builtin
from nuitka.nodes.AssignNodes import (
    StatementAssignmentVariable,
    StatementReleaseVariable
)
from nuitka.nodes.AsyncgenNodes import (
    ExpressionAsyncgenObjectBody,
    ExpressionMakeAsyncgenObject
)
from nuitka.nodes.BuiltinIteratorNodes import (
    ExpressionAsyncIter,
    ExpressionAsyncNext,
    ExpressionBuiltinIter1
)
from nuitka.nodes.BuiltinNextNodes import ExpressionBuiltinNext1
from nuitka.nodes.CodeObjectSpecs import CodeObjectSpec
from nuitka.nodes.ConditionalNodes import makeStatementConditional
from nuitka.nodes.ConstantRefNodes import makeConstantRefNode
from nuitka.nodes.ContainerOperationNodes import (
    StatementListOperationAppend,
    StatementSetOperationAdd
)
from nuitka.nodes.DictionaryNodes import StatementDictOperationSet
from nuitka.nodes.FrameNodes import (
    StatementsFrameFunction,
    StatementsFrameGenerator
)
from nuitka.nodes.FunctionNodes import ExpressionFunctionRef
from nuitka.nodes.GeneratorNodes import (
    ExpressionGeneratorObjectBody,
    ExpressionMakeGeneratorObject,
    StatementGeneratorReturnNone
)
from nuitka.nodes.LoopNodes import StatementLoop, StatementLoopBreak
from nuitka.nodes.NodeMakingHelpers import makeVariableRefNode
from nuitka.nodes.OutlineNodes import (
    ExpressionOutlineBody,
    ExpressionOutlineFunction
)
from nuitka.nodes.ReturnNodes import StatementReturn
from nuitka.nodes.StatementNodes import (
    StatementExpressionOnly,
    StatementsSequence
)
from nuitka.nodes.VariableRefNodes import ExpressionTempVariableRef
from nuitka.nodes.YieldNodes import (
    ExpressionYield,
    ExpressionYieldFromWaitable
)
from nuitka.PythonVersions import python_version

from .ReformulationAssignmentStatements import buildAssignmentStatements
from .ReformulationBooleanExpressions import buildAndNode
from .ReformulationTryExceptStatements import makeTryExceptSingleHandlerNode
from .ReformulationTryFinallyStatements import makeTryFinallyStatement
from .TreeHelpers import (
    buildNode,
    buildNodeList,
    detectFunctionBodyKind,
    getKind,
    makeStatementsSequenceFromStatement,
    makeStatementsSequenceFromStatements,
    mergeStatements
)


def _makeIteratorCreation(provider, qual, for_asyncgen, source_ref):
    if getattr(qual, "is_async", 0):
        result = ExpressionAsyncIter(
            value      = buildNode(
                provider   = provider,
                node       = qual.iter,
                source_ref = source_ref
            ),
            source_ref = source_ref
        )

        if not for_asyncgen or python_version < 370:
            result = ExpressionYieldFromWaitable(
                expression = result,
                source_ref = source_ref
            )

        return result
    else:
        return ExpressionBuiltinIter1(
            value      = buildNode(
                provider   = provider,
                node       = qual.iter,
                source_ref = source_ref
            ),
            source_ref = source_ref
        )


def _makeIteratorNext(qual, iterator_ref, source_ref):
    if getattr(qual, "is_async", 0):
        return ExpressionYieldFromWaitable(
            expression = ExpressionAsyncNext(
                value      = iterator_ref,
                source_ref = source_ref
            ),
            source_ref = source_ref
        )
    else:
        return ExpressionBuiltinNext1(
            value      = iterator_ref,
            source_ref = source_ref
        )


def _getStopIterationName(qual):
    if getattr(qual, "is_async", 0):
        return "StopAsyncIteration"
    else:
        return "StopIteration"


def _buildPython2ListContraction(provider, node, source_ref):
    # The contraction nodes are reformulated to function bodies, with loops as
    # described in the developer manual. They use a lot of temporary names,
    # nested blocks, etc. and so a lot of variable names.

    # Note: The assign_provider is only to cover Python2 list contractions,
    # assigning one of the loop variables to the outside scope.
    function_body = ExpressionOutlineBody(
        provider   = provider,
        name       = "list_contraction",
        source_ref = source_ref
    )

    iter_tmp = function_body.allocateTempVariable(
        temp_scope = None,
        name       = ".0"
    )

    container_tmp = function_body.allocateTempVariable(
        temp_scope = None,
        name       = "contraction_result"
    )

    statements, release_statements = _buildContractionBodyNode(
        provider        = provider,
        node            = node,
        emit_class      = StatementListOperationAppend,
        iter_tmp        = iter_tmp,
        temp_scope      = None,
        start_value     = [],
        container_tmp   = container_tmp,
        function_body   = function_body,
        assign_provider = True,
        for_asyncgen    = False,
        source_ref      = source_ref,
    )

    statements.append(
        StatementReturn(
            expression = ExpressionTempVariableRef(
                variable   = container_tmp,
                source_ref = source_ref
            ),
            source_ref = source_ref
        )
    )

    statement = makeTryFinallyStatement(
        provider   = function_body,
        tried      = statements,
        final      = release_statements,
        source_ref = source_ref.atInternal()
    )

    function_body.setBody(
        makeStatementsSequenceFromStatement(
            statement = statement
        )
    )

    return function_body


def buildListContractionNode(provider, node, source_ref):
    # List contractions are dealt with by general code.
    if python_version < 300:
        return _buildPython2ListContraction(
            provider   = provider,
            node       = node,
            source_ref = source_ref
        )

    return _buildContractionNode(
        provider    = provider,
        node        = node,
        name        = "<listcomp>",
        emit_class  = StatementListOperationAppend,
        start_value = [],
        source_ref  = source_ref
    )


def buildSetContractionNode(provider, node, source_ref):
    # Set contractions are dealt with by general code.

    return _buildContractionNode(
        provider    = provider,
        node        = node,
        name        = "<setcontraction>",
        emit_class  = StatementSetOperationAdd,
        start_value = set(),
        source_ref  = source_ref
    )


def buildDictContractionNode(provider, node, source_ref):
    # Dict contractions are dealt with by general code.

    return _buildContractionNode(
        provider    = provider,
        node        = node,
        name        = "<dictcontraction>",
        emit_class  = StatementDictOperationSet,
        start_value = {},
        source_ref  = source_ref
    )


def buildGeneratorExpressionNode(provider, node, source_ref):
    # Generator expressions are dealt with by general code.

    assert getKind(node) == "GeneratorExp"

    function_body = ExpressionOutlineBody(
        provider   = provider,
        name       = "genexpr",
        source_ref = source_ref
    )

    iter_tmp = function_body.allocateTempVariable(
        temp_scope = None,
        name       = ".0"
    )

    parent_module = provider.getParentModule()

    code_object = CodeObjectSpec(
        co_name           = "<genexpr>",
        co_kind           = "Generator",
        co_varnames       = (".0",),
        co_argcount       = 1,
        co_kwonlyargcount = 0,
        co_has_starlist   = False,
        co_has_stardict   = False,
        co_filename       = parent_module.getRunTimeFilename(),
        co_lineno         = source_ref.getLineNumber(),
        future_spec       = parent_module.getFutureSpec()
    )

    if python_version < 370:
        is_async = any(getattr(qual, "is_async", 0) for qual in node.generators)
    else:
        is_async = detectFunctionBodyKind(nodes = [node])[0] in ("Asyncgen", "Coroutine")

    if is_async:
        code_body = ExpressionAsyncgenObjectBody(
            provider    = provider,
            name        = "<genexpr>",
            code_object = code_object,
            flags       = set(),
            source_ref  = source_ref
        )

        maker_class = ExpressionMakeAsyncgenObject
    else:
        code_body = ExpressionGeneratorObjectBody(
            provider    = provider,
            name        = "<genexpr>",
            code_object = code_object,
            flags       = set(),
            source_ref  = source_ref
        )

        maker_class = ExpressionMakeGeneratorObject


    function_body.setBody(
        makeStatementsSequenceFromStatements(
            StatementAssignmentVariable(
                variable   = iter_tmp,
                source     = _makeIteratorCreation(
                    provider     = provider,
                    qual         = node.generators[0],
                    for_asyncgen = is_async,
                    source_ref   = source_ref
                ),
                source_ref = source_ref
            ),
            makeTryFinallyStatement(
                provider   = function_body,
                tried      = StatementReturn(
                    expression = maker_class(
                        ExpressionFunctionRef(
                            function_body = code_body,
                            source_ref    = source_ref
                        ),
                        source_ref = source_ref
                    ),
                    source_ref = source_ref
                ),
                final      = StatementReleaseVariable(
                    variable   = iter_tmp,
                    source_ref = source_ref
                ),
                source_ref = source_ref
            )
        )
    )

    statements, release_statements = _buildContractionBodyNode(
        provider        = provider,
        node            = node,
        emit_class      = ExpressionYield,
        iter_tmp        = iter_tmp,
        temp_scope      = None,
        start_value     = None,
        container_tmp   = None,
        function_body   = code_body,
        assign_provider = False,
        for_asyncgen    = is_async,
        source_ref      = source_ref,
    )

    if is_async:
        statements.append(
            StatementGeneratorReturnNone(
                source_ref = source_ref
            )
        )

    statements = (
        makeTryFinallyStatement(
            provider   = function_body,
            tried      = statements,
            final      = release_statements,
            source_ref = source_ref.atInternal()
        ),
    )

    code_body.setBody(
        makeStatementsSequenceFromStatement(
            statement = StatementsFrameGenerator(
                statements  = mergeStatements(statements, False),
                code_object = code_object,
                source_ref  = source_ref
            )
        )
    )

    return function_body


def _buildContractionBodyNode(provider, node, emit_class, start_value,
                              container_tmp, iter_tmp, temp_scope,
                              assign_provider, function_body, for_asyncgen,
                              source_ref):

    # This uses lots of variables and branches. There is no good way
    # around that, and we deal with many cases, due to having generator
    # expressions sharing this code, pylint: disable=too-many-branches,too-many-locals,too-many-statements

    # Note: The assign_provider is only to cover Python2 list contractions,
    # assigning one of the loop variables to the outside scope.

    tmp_variables = []
    if emit_class is not ExpressionYield:
        tmp_variables.append(iter_tmp)

    if container_tmp is not None:
        tmp_variables.append(container_tmp)

    statements = []

    # First assign the iterator if we are an outline.
    if assign_provider:
        statements.append(
            StatementAssignmentVariable(
                variable   = iter_tmp,
                source     = _makeIteratorCreation(
                    provider     = provider,
                    qual         = node.generators[0],
                    for_asyncgen = False,
                    source_ref   = source_ref
                ),
                source_ref = source_ref.atInternal()
            )
        )

    if for_asyncgen and python_version >= 370 and node.generators[0].is_async:
        statements.append(
            StatementAssignmentVariable(
                variable   = iter_tmp,
                source     = ExpressionTempVariableRef(
                    variable   = iter_tmp,
                    source_ref = source_ref
                ),
                source_ref = source_ref
            )
        )

    if start_value is not None:
        statements.append(
            StatementAssignmentVariable(
                variable   = container_tmp,
                source     = makeConstantRefNode(
                    constant   = start_value,
                    source_ref = source_ref
                ),
                source_ref = source_ref.atInternal()
            )
        )

    if hasattr(node, "elt"):
        if start_value is not None:
            current_body = emit_class(
                ExpressionTempVariableRef(
                    variable   = container_tmp,
                    source_ref = source_ref
                ),
                buildNode(
                    provider   = function_body if not assign_provider else provider,
                    node       = node.elt,
                    source_ref = source_ref
                ),
                source_ref = source_ref
            )
        else:
            assert emit_class is ExpressionYield

            current_body = emit_class(
                buildNode(
                    provider   = function_body,
                    node       = node.elt,
                    source_ref = source_ref
                ),
                source_ref = source_ref
            )
    else:
        assert emit_class is StatementDictOperationSet

        current_body = StatementDictOperationSet(
            dict_arg   = ExpressionTempVariableRef(
                variable   = container_tmp,
                source_ref = source_ref
            ),
            key        = buildNode(
                provider   = function_body if not assign_provider else provider,
                node       = node.key,
                source_ref = source_ref,
            ),
            value      = buildNode(
                provider   = function_body if not assign_provider else provider,
                node       = node.value,
                source_ref = source_ref,
            ),
            source_ref = source_ref
        )

    if current_body.isExpression():
        current_body = StatementExpressionOnly(
            expression = current_body,
            source_ref = source_ref
        )

    for count, qual in enumerate(reversed(node.generators)):
        tmp_value_variable = function_body.allocateTempVariable(
            temp_scope = temp_scope,
            name       = "iter_value_%d" % count
        )

        tmp_variables.append(tmp_value_variable)

        # The first iterated value is to be calculated outside of the function
        # and will be given as a parameter "_iterated", the others are built
        # inside the function.

        if qual is node.generators[0]:
            iterator_ref = makeVariableRefNode(
                variable   = iter_tmp,
                source_ref = source_ref
            )

            if for_asyncgen and python_version >= 370:
                iterator_ref = ExpressionYieldFromWaitable(
                    expression = iterator_ref,
                    source_ref = source_ref
                )

            tmp_iter_variable = None

            nested_statements = []
        else:
            # First create the iterator and store it, next should be loop body
            value_iterator = _makeIteratorCreation(
                provider     = provider if assign_provider else function_body,
                qual         = qual,
                for_asyncgen = False,
                source_ref   = source_ref
            )

            tmp_iter_variable = function_body.allocateTempVariable(
                temp_scope = temp_scope,
                name       = "contraction_iter_%d" % count
            )

            tmp_variables.append(tmp_iter_variable)

            nested_statements = [
                StatementAssignmentVariable(
                    variable   = tmp_iter_variable,
                    source     = value_iterator,
                    source_ref = source_ref
                )
            ]

            iterator_ref = ExpressionTempVariableRef(
                variable   = tmp_iter_variable,
                source_ref = source_ref
            )

        loop_statements = [
            makeTryExceptSingleHandlerNode(
                tried          = StatementAssignmentVariable(
                    variable   = tmp_value_variable,
                    source     = _makeIteratorNext(
                        iterator_ref = iterator_ref,
                        qual         = qual,
                        source_ref   = source_ref
                    ),
                    source_ref = source_ref
                ),
                exception_name = _getStopIterationName(qual),
                handler_body   = StatementLoopBreak(
                    source_ref = source_ref
                ),
                source_ref     = source_ref
            ),
            buildAssignmentStatements(
                provider      = provider if assign_provider else function_body,
                temp_provider = function_body,
                node          = qual.target,
                source        = ExpressionTempVariableRef(
                    variable   = tmp_value_variable,
                    source_ref = source_ref
                ),
                source_ref    = source_ref
            )
        ]

        conditions = buildNodeList(
            provider   = provider if assign_provider else function_body,
            nodes      = qual.ifs,
            source_ref = source_ref
        )

        if len(conditions) >= 1:
            loop_statements.append(
                makeStatementConditional(
                    condition  = buildAndNode(
                        values     = conditions,
                        source_ref = source_ref
                    ),
                    yes_branch = current_body,
                    no_branch  = None,
                    source_ref = source_ref
                )
            )
        else:
            loop_statements.append(current_body)

        nested_statements.append(
            StatementLoop(
                body       = StatementsSequence(
                    statements = mergeStatements(loop_statements),
                    source_ref = source_ref
                ),
                source_ref = source_ref
            )
        )

        if tmp_iter_variable is not None:
            nested_statements.append(
                StatementReleaseVariable(
                    variable   = tmp_iter_variable,
                    source_ref = source_ref
                )
            )

        current_body = StatementsSequence(
            statements = mergeStatements(nested_statements, False),
            source_ref = source_ref
        )

    statements.append(current_body)
    statements = mergeStatements(statements)

    release_statements = [
        StatementReleaseVariable(
            variable   = tmp_variable,
            source_ref = source_ref
        )
        for tmp_variable in
        tmp_variables
    ]

    return statements, release_statements


def _buildContractionNode(provider, node, name, emit_class, start_value,
                          source_ref):
    # The contraction nodes are reformulated to function bodies, with loops as
    # described in the developer manual. They use a lot of temporary names,
    # nested blocks, etc. and so a lot of variable names.

    function_body = ExpressionOutlineFunction(
        provider   = provider,
        name       = intern(name[1:-1]),
        source_ref = source_ref
    )

    iter_tmp = function_body.allocateTempVariable(
        temp_scope = None,
        name       = ".0"
    )

    container_tmp = function_body.allocateTempVariable(
        temp_scope = None,
        name       = "contraction"
    )

    statements, release_statements = _buildContractionBodyNode(
        provider        = provider,
        node            = node,
        emit_class      = emit_class,
        iter_tmp        = iter_tmp,
        temp_scope      = None,
        start_value     = start_value,
        container_tmp   = container_tmp,
        function_body   = function_body,
        assign_provider = False,
        for_asyncgen    = False,
        source_ref      = source_ref,
    )

    assign_iter_statement = StatementAssignmentVariable(
        source     = _makeIteratorCreation(
            provider     = provider,
            qual         = node.generators[0],
            for_asyncgen = False,
            source_ref   = source_ref
        ),
        variable   = iter_tmp,
        source_ref = source_ref
    )

    statements.append(
        StatementReturn(
            expression = ExpressionTempVariableRef(
                variable   = container_tmp,
                source_ref = source_ref
            ),
            source_ref = source_ref
        )
    )

    statements = (
        makeTryFinallyStatement(
            provider   = function_body,
            tried      = statements,
            final      = release_statements,
            source_ref = source_ref.atInternal()
        ),
    )

    if python_version < 300:
        body = makeStatementsSequenceFromStatements(
            assign_iter_statement,
            statements
        )
    else:
        parent_module = provider.getParentModule()

        code_object = CodeObjectSpec(
            co_name           = name,
            co_kind           = "Function",
            co_varnames       = (),
            co_argcount       = 1,
            co_kwonlyargcount = 0,
            co_has_starlist   = False,
            co_has_stardict   = False,
            co_filename       = parent_module.getRunTimeFilename(),
            co_lineno         = source_ref.getLineNumber(),
            future_spec       = parent_module.getFutureSpec()
        )

        body = makeStatementsSequenceFromStatements(
            assign_iter_statement,
            StatementsFrameFunction(
                statements  = mergeStatements(statements, False),
                code_object = code_object,
                source_ref  = source_ref
            )
        )

    function_body.setBody(body)

    return function_body
