# -*- coding: utf-8 -*-
import pytest
import aerospike
from aerospike import exception as e
from aerospike_helpers.operations import map_operations as map_ops
from .test_base_class import TestBaseClass


def get_map_result_from_operation(client, key, operations, res_bin):
    """
    Map operations return a 1 bin record.
    This operation extracts the dictionary from it.
    """
    _, _, result_bins = client.operate(key, operations)
    return result_bins[res_bin]


def maps_have_same_keys(map1, map2):
    """
    Return True iff the two maps have an identical set of keys
    """
    if len(map1) != len(map2):
        return False

    for key in map1:
        if key not in map2:
            return False
    return True


def maps_have_same_values(map1, map2):
    if len(map1) != len(map2):
        return False

    for key in map1:
        if key not in map2 or map1[key] != map2[key]:
            return False
    return True


def sort_map(client, test_key, test_bin):
    map_policy = {"map_write_flags": aerospike.MAP_WRITE_FLAGS_CREATE_ONLY, "map_order": aerospike.MAP_KEY_ORDERED}
    operations = [map_ops.map_set_policy(test_bin, map_policy)]
    client.operate(test_key, operations)


# Used in one test case
MAP_WITH_ONE_KEY_BIN_NAME = "map_w_one_key"


class TestNewListOperationsHelpers(object):
    @pytest.fixture(autouse=True)
    def setup(self, request, as_connection):
        """
        Setup Method
        """
        self.Server61 = True
        if TestBaseClass.major_ver < 6 or (TestBaseClass.major_ver == 6 and TestBaseClass.minor_ver == 0):
            self.Server61 = False

        self.keys = []
        # In a CDT, different types are also ordered from each other
        # boolean < integer
        self.test_map = {"a": 5, "b": 4, "c": 3, "d": 2, "e": 1, "f": True, "g": False}

        self.unsorted_map_bin = "unsorted_map"
        self.unsorted_map = {"a": 1, "c": 3, "b": 2}

        self.test_key = "test", "demo", "new_map_op"
        self.test_bin = "map"
        self.as_connection.put(self.test_key, {self.test_bin: self.test_map, self.unsorted_map_bin: self.unsorted_map})
        self.keys.append(self.test_key)

        yield

        for key in self.keys:
            try:
                self.as_connection.remove(key)
            except e.AerospikeError:
                pass

    def test_map_set_policy(self):
        """
        Test setting map policy with an operation
        """
        map_policy = {
            "map_write_flags": aerospike.MAP_WRITE_FLAGS_CREATE_ONLY,
            "map_order": aerospike.MAP_KEY_VALUE_ORDERED,
            "persist_index": True
        }
        operations = [map_ops.map_set_policy(self.test_bin, map_policy)]

        self.as_connection.operate(self.test_key, operations)

    def test_map_policy_invalid_persist_index(self):
        map_policy = {
            "persist_index": 1
        }
        operations = [map_ops.map_set_policy(self.test_bin, map_policy)]

        with pytest.raises(e.ParamError):
            self.as_connection.operate(self.test_key, operations)

    # Default persist index value should be tested automatically
    # from other tests that don't set the persist index option

    def test_map_put(self):
        operations = [map_ops.map_put(self.test_bin, "new", "map_put")]
        self.as_connection.operate(self.test_key, operations)
        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert res_map["new"] == "map_put"

    def test_map_put_items(self):
        new_items = {"new": 1, "new2": 3}
        expected_items = {"a": 5, "b": 4, "c": 3, "d": 2, "e": 1, "f": True, "g": False, "new": 1, "new2": 3}
        operations = [map_ops.map_put_items(self.test_bin, new_items)]
        self.as_connection.operate(self.test_key, operations)
        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]

        assert res_map == expected_items

    def test_map_increment(self):
        operations = [map_ops.map_increment(self.test_bin, "a", 100)]
        self.as_connection.operate(self.test_key, operations)
        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert res_map["a"] == self.test_map["a"] + 100

    def test_map_decrement(self):
        operations = [map_ops.map_decrement(self.test_bin, "a", 100)]
        self.as_connection.operate(self.test_key, operations)
        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert res_map["a"] == self.test_map["a"] - 100

    def test_map_size(self):
        operations = [map_ops.map_size(self.test_bin)]
        assert get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin) == len(
            self.test_map
        )

    def test_map_clear(self):
        operations = [map_ops.map_clear(self.test_bin)]
        self.as_connection.operate(self.test_key, operations)
        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert res_map == {}

    def test_map_remove_by_key(self):
        operations = [map_ops.map_remove_by_key(self.test_bin, "a", return_type=aerospike.MAP_RETURN_VALUE)]
        assert (
            get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
            == self.test_map["a"]
        )

        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert "a" not in res_map

    def test_map_remove_by_key_list(self):
        operations = [
            map_ops.map_remove_by_key_list(self.test_bin, ["a", "b", "c"], return_type=aerospike.MAP_RETURN_VALUE)
        ]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert set(ret_vals) == set([self.test_map["a"], self.test_map["b"], self.test_map["c"]])

        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert "a" not in res_map
        assert "b" not in res_map
        assert "c" not in res_map

    def test_map_remove_by_key_range(self):
        operations = [map_ops.map_remove_by_key_range(self.test_bin, "a", "d", return_type=aerospike.MAP_RETURN_VALUE)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert set(ret_vals) == set([self.test_map["a"], self.test_map["b"], self.test_map["c"]])

        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert "a" not in res_map
        assert "b" not in res_map
        assert "c" not in res_map

    def test_map_remove_by_value(self):
        operations = [
            map_ops.map_remove_by_value(self.test_bin, self.test_map["a"], return_type=aerospike.MAP_RETURN_KEY)
        ]
        assert get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin) == ["a"]

        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert "a" not in res_map

    def test_map_remove_by_value_list(self):
        operations = [
            map_ops.map_remove_by_value_list(
                self.test_bin,
                [self.test_map["a"], self.test_map["b"], self.test_map["c"]],
                return_type=aerospike.MAP_RETURN_KEY,
            )
        ]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert set(ret_vals) == set(["a", "b", "c"])

        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert "a" not in res_map
        assert "b" not in res_map
        assert "c" not in res_map

    def test_map_remove_by_value_range(self):
        operations = [map_ops.map_remove_by_value_range(self.test_bin, 1, 4, return_type=aerospike.MAP_RETURN_KEY)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert set(ret_vals) == set(["e", "d", "c"])

        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert "e" not in res_map
        assert "d" not in res_map
        assert "c" not in res_map

    def test_map_remove_by_index(self):
        sort_map(self.as_connection, self.test_key, self.test_bin)
        operations = [map_ops.map_remove_by_index(self.test_bin, 1, return_type=aerospike.MAP_RETURN_KEY)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        assert ret_vals == "b"
        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert "b" not in res_map

    # For the rest of the map_remove* operations, there are test cases that produce an empty map
    @pytest.mark.parametrize(
        "op",
        [
            map_ops.map_remove_by_key(
                bin_name=MAP_WITH_ONE_KEY_BIN_NAME,
                key="a",
                return_type=aerospike.MAP_RETURN_NONE
            ),
            map_ops.map_remove_by_index(
                bin_name=MAP_WITH_ONE_KEY_BIN_NAME,
                index=0,
                return_type=aerospike.MAP_RETURN_NONE
            ),
            map_ops.map_remove_by_rank(
                bin_name=MAP_WITH_ONE_KEY_BIN_NAME,
                rank=0,
                return_type=aerospike.MAP_RETURN_NONE
            ),
            map_ops.map_remove_by_value(
                bin_name=MAP_WITH_ONE_KEY_BIN_NAME,
                value=1,
                return_type=aerospike.MAP_RETURN_NONE
            )
        ]
    )
    def test_map_remove_ops_causing_empty_map(self, op):
        # Add a map bin with one key on the fly
        bins = {
            MAP_WITH_ONE_KEY_BIN_NAME: {"a": 1}
        }
        self.as_connection.put(self.test_key, bins)

        ops = [
            op
        ]
        self.as_connection.operate(self.test_key, ops)

        _, _, bins = self.as_connection.get(self.test_key)
        # "a" should've been deleted
        assert bins[MAP_WITH_ONE_KEY_BIN_NAME] == {}

    def test_map_remove_by_index_range(self):
        sort_map(self.as_connection, self.test_key, self.test_bin)
        operations = [map_ops.map_remove_by_index_range(self.test_bin, 1, 2, return_type=aerospike.MAP_RETURN_KEY)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        assert ret_vals == ["b", "c"]
        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert "b" not in res_map
        assert "c" not in res_map

    def test_map_remove_by_rank(self):
        operations = [map_ops.map_remove_by_rank(self.test_bin, 1, return_type=aerospike.MAP_RETURN_KEY)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        # True > False
        assert ret_vals == "f"
        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert "f" not in res_map

    def test_map_remove_by_rank_range(self):
        operations = [map_ops.map_remove_by_rank_range(self.test_bin, 1, 2, return_type=aerospike.MAP_RETURN_KEY)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        # f: True, e: 1
        # Removes 2 items from rank 1 and going up the rank
        assert set(ret_vals) == set(["e", "f"])
        res_map = self.as_connection.get(self.test_key)[2][self.test_bin]
        assert "e" not in res_map
        assert "f" not in res_map

    def test_map_get_by_key(self):
        operations = [map_ops.map_get_by_key(self.test_bin, "a", return_type=aerospike.MAP_RETURN_VALUE)]
        assert (
            get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
            == self.test_map["a"]
        )

    def test_map_get_by_key_list(self):
        operations = [
            map_ops.map_get_by_key_list(self.test_bin, ["a", "b", "c"], return_type=aerospike.MAP_RETURN_VALUE)
        ]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert set(ret_vals) == set([self.test_map["a"], self.test_map["b"], self.test_map["c"]])

    def test_map_get_by_key_range(self):
        operations = [map_ops.map_get_by_key_range(self.test_bin, "a", "d", return_type=aerospike.MAP_RETURN_VALUE)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert set(ret_vals) == set([self.test_map["a"], self.test_map["b"], self.test_map["c"]])

    def test_map_get_by_value(self):
        operations = [map_ops.map_get_by_value(self.test_bin, self.test_map["a"], return_type=aerospike.MAP_RETURN_KEY)]
        assert get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin) == ["a"]

    def test_map_get_by_value_list(self):
        operations = [
            map_ops.map_get_by_value_list(
                self.test_bin,
                [self.test_map["a"], self.test_map["b"], self.test_map["c"]],
                return_type=aerospike.MAP_RETURN_KEY,
            )
        ]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert set(ret_vals) == set(["a", "b", "c"])

    def test_map_get_by_value_range(self):
        operations = [map_ops.map_get_by_value_range(self.test_bin, 1, 4, return_type=aerospike.MAP_RETURN_KEY)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert set(ret_vals) == set(["e", "d", "c"])

    def test_map_get_by_index(self):
        sort_map(self.as_connection, self.test_key, self.test_bin)
        operations = [map_ops.map_get_by_index(self.test_bin, 1, return_type=aerospike.MAP_RETURN_KEY)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        assert ret_vals == "b"

    def test_map_get_by_index_range(self):
        sort_map(self.as_connection, self.test_key, self.test_bin)
        operations = [map_ops.map_get_by_index_range(self.test_bin, 1, 2, return_type=aerospike.MAP_RETURN_KEY)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        assert ret_vals == ["b", "c"]

    def test_map_get_by_rank(self):
        operations = [map_ops.map_get_by_rank(self.test_bin, 1, return_type=aerospike.MAP_RETURN_KEY)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        assert ret_vals == "f"

    def test_map_get_by_rank_range(self):
        sort_map(self.as_connection, self.test_key, self.test_bin)
        operations = [map_ops.map_get_by_rank_range(self.test_bin, 1, 2, return_type=aerospike.MAP_RETURN_KEY)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        assert ret_vals == ["f", "e"]

    @pytest.mark.parametrize(
            "map_return_type",
            [
                aerospike.MAP_RETURN_ORDERED_MAP,
                aerospike.MAP_RETURN_UNORDERED_MAP
            ]
    )
    def test_map_return_types(self, map_return_type):
        if (TestBaseClass.major_ver, TestBaseClass.minor_ver) < (6, 3):
            pytest.skip("It only applies to >= 6.3 enterprise edition")

        operations = [
            map_ops.map_get_by_key_list(self.unsorted_map_bin, ["a", "b", "c"], map_return_type)
        ]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.unsorted_map_bin)

        # The C client will always reorder maps using their keys
        # so it doesn't matter whether the user requests an unordered or ordered map
        # Both will return the same results
        assert list(ret_vals.items()) == list({"a": 1, "b": 2, "c": 3}.items())

    def test_map_get_exists_by_key_list(self):
        if not self.Server61:
            pytest.skip("It only applies to >= 6.1 enterprise edition")
        operations = [
            map_ops.map_get_by_key_list(self.test_bin, ["a", "b", "c"], aerospike.MAP_RETURN_EXISTS)
        ]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert ret_vals is True

    def test_map_get_exists_by_key_range(self):
        if not self.Server61:
            pytest.skip("It only applies to >= 6.1 enterprise edition")
        operations = [map_ops.map_get_by_key_range(self.test_bin, "a", "d", return_type=aerospike.MAP_RETURN_EXISTS)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert ret_vals is True

    def test_map_get_exists_by_value(self):
        if not self.Server61:
            pytest.skip("It only applies to >= 6.1 enterprise edition")
        operations = [
            map_ops.map_get_by_value(self.test_bin, self.test_map["a"], return_type=aerospike.MAP_RETURN_EXISTS)
        ]
        assert get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin) is True

    def test_map_get_exists_by_value_list(self):
        if not self.Server61:
            pytest.skip("It only applies to >= 6.1 enterprise edition")
        operations = [
            map_ops.map_get_by_value_list(
                self.test_bin,
                [self.test_map["a"], self.test_map["b"], self.test_map["c"]],
                return_type=aerospike.MAP_RETURN_EXISTS,
            )
        ]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert ret_vals is True

    def test_map_get_exists_by_value_range(self):
        if not self.Server61:
            pytest.skip("It only applies to >= 6.1 enterprise edition")
        operations = [map_ops.map_get_by_value_range(self.test_bin, 1, 4, return_type=aerospike.MAP_RETURN_EXISTS)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)

        assert ret_vals is True

    def test_map_get_exists_by_index(self):
        if not self.Server61:
            pytest.skip("It only applies to >= 6.1 enterprise edition")
        sort_map(self.as_connection, self.test_key, self.test_bin)
        operations = [map_ops.map_get_by_index(self.test_bin, 1, return_type=aerospike.MAP_RETURN_EXISTS)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        assert ret_vals is True

    def test_map_get_exists_by_index_range(self):
        if not self.Server61:
            pytest.skip("It only applies to >= 6.1 enterprise edition")
        sort_map(self.as_connection, self.test_key, self.test_bin)
        operations = [map_ops.map_get_by_index_range(self.test_bin, 1, 2, return_type=aerospike.MAP_RETURN_EXISTS)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        assert ret_vals is True

    def test_map_get_exists_by_rank(self):
        if not self.Server61:
            pytest.skip("It only applies to >= 6.1 enterprise edition")
        operations = [map_ops.map_get_by_rank(self.test_bin, 1, return_type=aerospike.MAP_RETURN_EXISTS)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        assert ret_vals is True

    def test_map_get_exists_by_rank_range(self):
        if not self.Server61:
            pytest.skip("It only applies to >= 6.1 enterprise edition")
        sort_map(self.as_connection, self.test_key, self.test_bin)
        operations = [map_ops.map_get_by_rank_range(self.test_bin, 1, 2, return_type=aerospike.MAP_RETURN_EXISTS)]
        ret_vals = get_map_result_from_operation(self.as_connection, self.test_key, operations, self.test_bin)
        assert ret_vals is True

    def test_map_create(self):
        # This should create an empty dictionary
        # map_create only works if a map does not already exist at the given bin and context path
        operations = [
            map_ops.map_create(bin_name="new_map", map_order=aerospike.MAP_KEY_ORDERED, persist_index=False, ctx=None)
        ]
        get_map_result_from_operation(self.as_connection, self.test_key, operations, "new_map")

        res_map = self.as_connection.get(self.test_key)[2]["new_map"]
        assert res_map == {}
