#     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.
#
""" Shared definitions of what comparison operation helpers are available.

These are functions to work with helper names, as well as sets of functions to
generate specialized code variants for.

Note: These are ordered, so we can define the order they are created in the
code generation of specialized helpers, as this set is used for input there
too.

"""

from nuitka.containers.OrderedSets import buildOrderedSet

# Mapping of rich comparison Python level name, to C level name of helpers.
rich_comparison_codes = {
    "Lt": "LT",
    "LtE": "LE",
    "Eq": "EQ",
    "NotEq": "NE",
    "Gt": "GT",
    "GtE": "GE",
}

# Subset of comparisons which will be used for identical types.
rich_comparison_subset_codes = {
    "Lt": "LT",
    "LtE": "LE",
    "Eq": "EQ",
}


def _makeDefaultOps():
    for comparator in rich_comparison_codes.values():
        yield "RICH_COMPARE_%s_OBJECT_OBJECT_OBJECT" % comparator
        yield "RICH_COMPARE_%s_NBOOL_OBJECT_OBJECT" % comparator


def _makeTypeOps(type_name, may_raise_same_type, shortcut=False):
    for result_part in "OBJECT", "CBOOL", "NBOOL":
        for comparator in rich_comparison_codes.values():
            if result_part == "CBOOL":
                continue

            yield "RICH_COMPARE_%s_%s_OBJECT_%s" % (comparator, result_part, type_name)
            yield "RICH_COMPARE_%s_%s_%s_OBJECT" % (comparator, result_part, type_name)

        if may_raise_same_type and result_part == "CBOOL":
            continue
        if not may_raise_same_type and result_part == "NBOOL":
            continue

        for comparator in (
            rich_comparison_codes.values()
            if not shortcut
            else rich_comparison_subset_codes.values()
        ):
            yield "RICH_COMPARE_%s_%s_%s_%s" % (
                comparator,
                result_part,
                type_name,
                type_name,
            )


def _makeFriendOps(type_name1, type_name2, may_raise):
    assert type_name1 != type_name2

    for result_part in "OBJECT", "CBOOL", "NBOOL":
        if not may_raise:
            if result_part == "NBOOL":
                continue

        for comparator in rich_comparison_codes.values():
            yield "RICH_COMPARE_%s_%s_%s_%s" % (
                comparator,
                result_part,
                type_name1,
                type_name2,
            )


specialized_cmp_helpers_set = buildOrderedSet(
    # Default implementation.
    _makeDefaultOps(),
    # Pure types:
    _makeTypeOps("STR", may_raise_same_type=False, shortcut=True),
    _makeTypeOps("UNICODE", may_raise_same_type=False, shortcut=True),
    _makeTypeOps("BYTES", may_raise_same_type=False, shortcut=True),
    _makeTypeOps("INT", may_raise_same_type=False, shortcut=True),
    _makeTypeOps("LONG", may_raise_same_type=False, shortcut=True),
    _makeTypeOps("FLOAT", may_raise_same_type=False, shortcut=True),
    # TODO: What would shortcut mean, how do tuples compare their elements then?
    _makeTypeOps("TUPLE", may_raise_same_type=True),
    _makeTypeOps("LIST", may_raise_same_type=True),
    # Mixed Python types:
    _makeFriendOps("LONG", "INT", may_raise=False),
    # Partial Python with C types
    _makeFriendOps("INT", "CLONG", may_raise=False),
    _makeFriendOps("LONG", "DIGIT", may_raise=False),
    _makeFriendOps("FLOAT", "CFLOAT", may_raise=False),
    # TODO: Add "CLONG_CLONG" type ops once we use that for local variables too.
)

_non_specialized_cmp_helpers_set = set()


def getSpecializedComparisonOperations():
    return specialized_cmp_helpers_set


def getNonSpecializedComparisonOperations():
    return _non_specialized_cmp_helpers_set
