#
# Copyright 2019 GridGain Systems, Inc. and Contributors.
#
# Licensed under the GridGain Community Edition License (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.gridgain.com/products/software/community-edition/gridgain-community-edition-license
#
# 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.
#
from typing import Iterable, Union

from pygridgain.connection import AioConnection, Connection
from pygridgain.datatypes import Bool, Int, Long, UUIDObject
from pygridgain.datatypes.internal import StructArray, Conditional, Struct
from pygridgain.queries import Query, query_perform
from pygridgain.queries.op_codes import OP_CACHE_PARTITIONS
from pygridgain.utils import is_iterable
from .result import APIResult


cache_ids = StructArray([
    ('cache_id', Int),
])

cache_config = StructArray([
    ('key_type_id', Int),
    ('affinity_key_field_id', Int),
])

node_partitions = StructArray([
    ('partition_id', Int),
])

node_mapping = StructArray([
    ('node_uuid', UUIDObject),
    ('node_partitions', node_partitions)
])

cache_mapping = StructArray([
    ('cache_id', Int),
    ('cache_config', cache_config),
])

empty_cache_mapping = StructArray([
    ('cache_id', Int)
])

empty_node_mapping = Struct([])

partition_mapping = StructArray([
    ('is_applicable', Bool),

    ('cache_mapping', Conditional(['is_applicable'],
                                  lambda ctx: ctx['is_applicable'] and ctx['is_applicable'].value == 1,
                                  lambda ctx: ctx['is_applicable'],
                                  cache_mapping, empty_cache_mapping)),

    ('node_mapping', Conditional(['is_applicable'],
                                 lambda ctx: ctx['is_applicable'] and ctx['is_applicable'].value == 1,
                                 lambda ctx: ctx['is_applicable'],
                                 node_mapping, empty_node_mapping)),
])


def cache_get_node_partitions(conn: 'Connection', caches: Union[int, Iterable[int]], query_id: int = None) -> APIResult:
    """
    Gets partition mapping for a GridGain cache or a number of caches. See
    “IEP-23: Best Effort Affinity for thin clients”.

    :param conn: connection to GridGain server,
    :param caches: cache ID(s) the mapping is provided for,
    :param query_id: (optional) a value generated by client and returned as-is
     in response.query_id. When the parameter is omitted, a random value
     is generated,
    :return: API result data object.
    """
    return __cache_get_node_partitions(conn, caches, query_id)


async def cache_get_node_partitions_async(conn: 'AioConnection', caches: Union[int, Iterable[int]],
                                          query_id: int = None) -> APIResult:
    """
    Async version of cache_get_node_partitions.
    """
    return await __cache_get_node_partitions(conn, caches, query_id)


def __post_process_partitions(result):
    if result.status == 0:
        # tidying up the result
        value = {
            'version': (
                result.value['version_major'],
                result.value['version_minor']
            ),
            'partition_mapping': {},
        }
        for partition_map in result.value['partition_mapping']:
            is_applicable = partition_map['is_applicable']

            node_mapping = None
            if is_applicable:
                node_mapping = {
                    p['node_uuid']: set(x['partition_id'] for x in p['node_partitions'])
                    for p in partition_map['node_mapping']
                }

            for cache_info in partition_map['cache_mapping']:
                cache_id = cache_info['cache_id']

                cache_partition_mapping = {
                    'is_applicable': is_applicable,
                }

                parts = 0
                if is_applicable:
                    cache_partition_mapping['cache_config'] = {
                        a['key_type_id']: a['affinity_key_field_id']
                        for a in cache_info['cache_config']
                    }
                    cache_partition_mapping['node_mapping'] = node_mapping

                    parts = sum(len(p) for p in cache_partition_mapping['node_mapping'].values())

                cache_partition_mapping['number_of_partitions'] = parts

                value['partition_mapping'][cache_id] = cache_partition_mapping
        result.value = value
    return result


def __cache_get_node_partitions(conn, caches, query_id):
    query_struct = Query(
        OP_CACHE_PARTITIONS,
        [
            ('cache_ids', cache_ids),
        ],
        query_id=query_id
    )
    if not is_iterable(caches):
        caches = [caches]

    return query_perform(
        query_struct,
        conn,
        query_params={
            'cache_ids': [{'cache_id': cache} for cache in caches],
        },
        response_config=[
            ('version_major', Long),
            ('version_minor', Int),
            ('partition_mapping', partition_mapping),
        ],
        post_process_fun=__post_process_partitions
    )
