# -*- coding: utf-8 -*-

import pytest
from .test_base_class import TestBaseClass
from .as_status_codes import AerospikeStatus
import aerospike
from aerospike import exception as e
from aerospike import predicates as p
from aerospike_helpers import expressions as exp
from aerospike_helpers import cdt_ctx
from threading import Lock
import time

from aerospike_helpers.expressions.arithmetic import Add
from aerospike_helpers.expressions.base import IntBin, GeoBin, ListBin, BlobBin

list_index = "list_index"
list_rank = "list_rank"
list_value = "list_value"
map_index = "map_index"
map_key = "map_key"
map_rank = "map_rank"
map_value = "map_value"

ctx_ops = {
    list_index: cdt_ctx.cdt_ctx_list_index,
    list_rank: cdt_ctx.cdt_ctx_list_rank,
    list_value: cdt_ctx.cdt_ctx_list_value,
    map_index: cdt_ctx.cdt_ctx_map_index,
    map_key: cdt_ctx.cdt_ctx_map_key,
    map_rank: cdt_ctx.cdt_ctx_map_rank,
    map_value: cdt_ctx.cdt_ctx_map_value,
}


def add_ctx_op(ctx_type, value):
    ctx_func = ctx_ops[ctx_type]
    return ctx_func(value)


ctx_list_index = []
ctx_list_index.append(add_ctx_op(list_index, 0))

ctx_list_rank = []
ctx_list_rank.append(add_ctx_op(list_rank, -1))

ctx_list_value = []
ctx_list_value.append(add_ctx_op(list_value, 3))

ctx_map_index = []
ctx_map_index.append(add_ctx_op(map_index, 0))

ctx_map_key = []
ctx_map_key.append(add_ctx_op(map_key, "sb"))

ctx_map_rank = []
ctx_map_rank.append(add_ctx_op(map_rank, -1))

ctx_map_value = []
ctx_map_value.append(add_ctx_op(map_value, 3))


class TestQuery(TestBaseClass):
    # TODO: This fixture should be split up to speed up this test class
    def setup_class(cls):
        client = TestBaseClass.get_new_connection()

        try:
            client.index_integer_create("test", "demo", "test_age", "age_index")
        except e.IndexFoundError:
            pass

        try:
            client.index_string_create("test", "demo", "addr", "addr_index")
        except e.IndexFoundError:
            pass

        try:
            client.index_integer_create("test", "demo", "age1", "age_index1")
        except e.IndexFoundError:
            pass

        try:
            client.index_list_create("test", "demo", "numeric_list", aerospike.INDEX_NUMERIC, "numeric_list_index")
        except e.IndexFoundError:
            pass

        try:
            client.index_list_create("test", "demo", "string_list", aerospike.INDEX_STRING, "string_list_index")
        except e.IndexFoundError:
            pass

        try:
            client.index_map_keys_create("test", "demo", "numeric_map", aerospike.INDEX_NUMERIC, "numeric_map_index")
        except e.IndexFoundError:
            pass

        try:
            client.index_map_keys_create("test", "demo", "string_map", aerospike.INDEX_STRING, "string_map_index")
        except e.IndexFoundError:
            pass

        try:
            client.index_map_values_create(
                "test", "demo", "numeric_map", aerospike.INDEX_NUMERIC, "numeric_map_values_index"
            )
        except e.IndexFoundError:
            pass

        try:
            client.index_map_values_create(
                "test", "demo", "string_map", aerospike.INDEX_STRING, "string_map_values_index"
            )
        except e.IndexFoundError:
            pass

        try:
            client.index_integer_create("test", None, "test_age_none", "age_index_none")
        except e.IndexFoundError:
            pass

        try:
            client.index_integer_create("test", "demo", bytearray("sal\0kj", "utf-8"), "sal_index")
        except e.IndexFoundError:
            pass

        if (TestBaseClass.major_ver, TestBaseClass.minor_ver) >= (7, 0):
            # These indexes are only used for server 7.0+ tests
            try:
                client.index_list_create("test", "demo", "blob_list", aerospike.INDEX_BLOB, "blob_list_index")
            except e.IndexFoundError:
                pass

            try:
                client.index_map_keys_create("test", "demo", "blob_map", aerospike.INDEX_BLOB, "blob_map_keys_index")
            except e.IndexFoundError:
                pass

            try:
                client.index_map_values_create(
                    "test", "demo", "blob_map", aerospike.INDEX_BLOB, "blob_map_values_index"
                )
            except e.IndexFoundError:
                pass

        try:
            client.index_cdt_create(
                "test",
                "demo",
                "numeric_list",
                aerospike.INDEX_TYPE_DEFAULT,
                aerospike.INDEX_NUMERIC,
                "numeric_list_cdt_index",
                {"ctx": ctx_list_index},
            )
        except e.IndexFoundError:
            pass

        try:
            client.index_cdt_create(
                "test",
                "demo",
                "numeric_map",
                aerospike.INDEX_TYPE_DEFAULT,
                aerospike.INDEX_NUMERIC,
                "numeric_map_cdt_index",
                {"ctx": ctx_map_index},
            )
        except e.IndexFoundError:
            pass

        client.close()

    def teardown_class(cls):
        client = TestBaseClass.get_new_connection()

        policy = {}
        try:
            client.index_remove("test", "age_index", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "age_index1", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "addr_index", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "numeric_list_index", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "string_list_index", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "numeric_map_index", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "string_map_index", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "numeric_map_values_index", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "string_map_values_index", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "age_index_none", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "sal_index")
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "numeric_list_cdt_index", policy)
        except e.IndexNotFound:
            pass

        try:
            client.index_remove("test", "numeric_map_cdt_index", policy)
        except e.IndexNotFound:
            pass

        blob_index_names_to_remove = [
            "blob_list_index",
            "blob_map_keys_index",
            "blob_map_values_index"
        ]
        for name in blob_index_names_to_remove:
            try:
                client.index_remove("test", name, policy)
            except e.IndexNotFound:
                pass

        client.close()

    @pytest.fixture(autouse=True)
    def setup_method(self, request, as_connection):
        """
        Setup method.
        """
        for i in range(5):
            key = ("test", "demo", i)
            # 5x5 box, then 10x10 box, ... until 25x25 box
            box_coordinates = [[0, 0], [5 * (i + 1), 0], [5 * (i + 1), 5 * (i + 1)], [0, 5 * (i + 1)], [0, 0]]
            rec = {
                "geo_point": aerospike.GeoJSON({"type": "Point", "coordinates": [i, i]}),
                "geo_polygon": aerospike.GeoJSON({"type": "Polygon", "coordinates": [box_coordinates]}),
                "name": "name%s" % (str(i)),
                "addr": "name%s" % (str(i)),
                "numeric_list": [i, i + 1, i + 2],
                "string_list": ["str" + str(i), "str" + str(i + 1), "str" + str(i + 2)],
                "numeric_map": {"a": i, "b": i + 1, "c": i + 2},
                "string_map": {"a": "a" + str(i), "b": "b" + str(i + 1), "c": "c" + str(i + 2)},
                "blob_list": [i.to_bytes(length=1, byteorder='big')],
                "blob_map": {
                    i.to_bytes(length=1, byteorder='big'): i.to_bytes(length=1, byteorder='big')
                },
                "test_age_none": 1,
                "test_age": i,
                "no": i,
                "blob": i.to_bytes(length=1, byteorder='big')
            }
            as_connection.put(key, rec)
        for i in range(5, 10):
            key = ("test", "demo", i)
            rec = {"name": "name%s" % (str(i)), "addr": "name%s" % (str(i)), "test_age": i, "no": i}
            as_connection.put(key, rec)

        key = ("test", "demo", 122)
        llist = [{"op": aerospike.OPERATOR_WRITE, "bin": bytearray("sal\0kj", "utf-8"), "val": 80000}]
        as_connection.operate(key, llist)

        key = ("test", None, 145)
        rec = {"test_age_none": 1}
        as_connection.put(key, rec)

        def teardown():
            """
            Teardown method.
            """
            for i in range(10):
                key = ("test", "demo", i)
                as_connection.remove(key)

            key = ("test", "demo", 122)
            as_connection.remove(key)
            key = ("test", None, 145)
            as_connection.remove(key)

        request.addfinalizer(teardown)

    def test_query_with_correct_parameters_hi(self):
        """
        Invoke query() with correct arguments
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))

        records = []

        def callback(input_tuple):
            try:
                key, _, _ = input_tuple
                records.append(key)
            except Exception as ex:
                print(ex)

        query.foreach(callback)
        assert records
        assert len(records) == 1

    def test_query_with_incorrect_bin_name(self):
        """
        Invoke query() with predicate comparing using non-extant bin name
        """
        query = self.as_connection.query("test", "demo")
        query.select("name1", "age1")
        query.where(p.equals("age1", 1))
        records = []

        def callback(input_tuple):
            _, _, record = input_tuple
            records.append(record)

        query.foreach(callback)
        assert len(records) == 0

    def test_query_without_callback_parameter(self):
        """
        Invoke query.foreach() with without callback
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))

        def callback(input_tuple):
            _, metadata, _ = input_tuple
            assert metadata["gen"] is None

        with pytest.raises(TypeError) as typeError:
            query.foreach()

        assert "argument 'callback'" in str(typeError.value)

    def test_query_with_nonindexed_bin(self):
        """
        Invoke query() with non-indexed bin
        """

        query = self.as_connection.query("test", "demo")
        query.select("name", "no")
        query.where(p.equals("no", 1))
        time.sleep(1)

        def callback(input_tuple):
            pass

        with pytest.raises(e.IndexNotFound):
            query.foreach(callback)

    def test_query_with_where_incorrect(self):
        """
        Invoke query() with predicate which
        matches 0 records
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 165))
        records = []

        def callback(input_tuple):
            _, _, record = input_tuple
            records.append(record)

        query.foreach(callback)
        assert len(records) == 0

    def test_query_with_where_none_value(self):
        """
        Invoke query() with equality
        predicate comparing to None
        This is more of a predicate test

        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        try:
            query.where(p.equals("test_age", None))

        except e.ParamError as exception:
            assert exception.code == -2
            assert exception.msg == "predicate is invalid."

    def test_query_where_called_multiple_times(self):
        query = self.as_connection.query("test", "demo")
        query.where(p.equals("test_age", 165))
        with pytest.raises(e.ClientError):
            query.where(p.equals("test_age", 150))

    def test_query_with_policy(self):
        """
        Invoke query() with policy
        """
        policy = {"total_timeout": 180000, "replica": aerospike.POLICY_REPLICA_MASTER}
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))
        records = []

        def callback(input_tuple):
            _, _, record = input_tuple
            records.append(record)

        query.foreach(callback, policy)
        assert len(records) == 1

    @pytest.mark.xfail(reason="Query policies do not support rps")
    def test_query_with_policy_records_per_second(self):
        """
        Invoke query() with the records_per_second policy set
        """
        policy = {"records_per_second": 10}
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))
        records = []

        def callback(input_tuple):
            _, _, record = input_tuple
            records.append(record)

        query.foreach(callback, policy)
        assert len(records) == 1

    @pytest.mark.xfail(reason="Only pass with server version >= 3.15")
    def test_query_with_no_bins_option(self):
        """
        Invoke query() with policy
        """
        policy = {}
        options = {"nobins": True}
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))

        def callback(input_tuple):
            _, _, bins = input_tuple
            print(bins)
            assert len(bins) == 0

        query.foreach(callback, policy, options=options)

    def test_query_with_invalid_options_argument_type(self):
        """
        Invoke query() with incorrect options type, should be a dictionary
        """
        policy = {}
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))

        def callback(input_tuple):
            _, metadata, _ = input_tuple
            assert metadata["gen"] is None

        with pytest.raises(e.ParamError):
            query.foreach(callback, policy, "")

    def test_query_with_invalid_nobins_value(self):
        """
        Invoke query() with options['nobins'] type
        """
        policy = {}
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))

        def callback(input_tuple):
            pass

        with pytest.raises(e.ParamError):
            query.foreach(callback, policy, {"nobins": "False"})

    def test_query_with_put_in_callback(self):
        """
        Invoke query() with client.put in foreach callback
        """
        policy = {}
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))
        records = []

        def callback(input_tuple):
            _, _, record = input_tuple
            records.append(record)
            key = ("test", "demo", "put_in_callback")
            rec = {
                "name": "name%s" % (str(8)),
                "test_age": 8,
            }
            self.as_connection.put(key, rec)

        query.foreach(callback, policy)

        key = ("test", "demo", "put_in_callback")
        _, _, bins = self.as_connection.get(key)
        assert bins == {"test_age": 8, "name": "name8"}

        self.as_connection.remove(key)

    def test_query_with_correct_parameters_between(self):
        """
        Invoke query() with correct arguments and using predicate between
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.between("test_age", 1, 4))

        records = []

        def callback(input_tuple):
            _, _, record = input_tuple
            records.append(record)

        query.foreach(callback)
        assert len(records) == 4

    def test_query_with_callback_returning_false(self):
        """
        Invoke query() with callback function returns false
        This will stop iteration
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.between("test_age", 1, 5))
        lock = Lock()
        records = []

        def callback(input_tuple):
            key, _, _ = input_tuple
            with lock:
                if len(records) == 2:
                    return False
                records.append(key)

        query.foreach(callback)
        assert len(records) == 2

    def test_query_with_results_method(self):
        """
        Invoke query() with correct arguments
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))

        records = query.results()
        assert len(records) == 1

    def test_query_with_results_method_and_expressions(self):
        """
        Invoke query() with correct arguments
        """
        expr = exp.Eq(exp.IntBin("test_age"), 1)

        policy = {"expressions": expr.compile()}

        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")

        records = query.results(policy)
        assert len(records) == 1

    def test_query_with_results_method_and_invalid_exp(self):
        """
        Invoke query() with correct arguments
        """
        expr = exp.Eq(exp.IntBin("test_age"), "bad_arg")

        policy = {"expressions": expr.compile()}

        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")

        with pytest.raises(e.InvalidRequest):
            query.results(policy)

    def test_query_with_results_nobins_options(self):
        """
        Invoke query() with correct arguments
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))

        records = query.results(options={"nobins": True})
        assert len(records) == 1

    def test_query_with_results_invalid_options_type(self):
        """
        Invoke query() with correct arguments
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))
        with pytest.raises(e.ParamError):
            query.results(options=False)

    def test_query_with_results_invalid_nobins_options(self):
        """
        Invoke query() with correct arguments
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))
        with pytest.raises(e.ParamError):
            query.results(options={"nobins": "false"})

    def test_query_with_unicode_binnames_in_select_and_where(self):
        """
        Invoke query() with unicode bin names in select
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age", "addr")
        query.where(p.equals("test_age", 7))

        records = query.results()
        assert len(records) == 1
        assert records[0][2] == {"test_age": 7, "name": "name7", "addr": "name7"}

    @pytest.mark.parametrize(
        "predicate, expected_min_length, expected_max_length",
        (
            (p.contains("numeric_list", aerospike.INDEX_TYPE_LIST, 1), 2, 2),
            (p.contains("string_list", aerospike.INDEX_TYPE_LIST, "str3"), 3, 3),
            (p.contains("string_map", aerospike.INDEX_TYPE_MAPKEYS, "a"), 5, 5),
            (p.contains("string_map", aerospike.INDEX_TYPE_MAPVALUES, "a1"), 1, 1),
            (p.contains("numeric_map", aerospike.INDEX_TYPE_MAPVALUES, 1), 2, 2),
            (p.range("numeric_map", aerospike.INDEX_TYPE_MAPVALUES, 1, 3), 4, 8),
            (p.range("numeric_list", aerospike.INDEX_TYPE_LIST, 1, 3), 4, 8),
        ),
        ids=(
            "numeric_list contains",
            "string list contains",
            "string mapkeys contains",
            "string mapvalue contains",
            "numeric mapvalue contains",
            "numeric mapvalue range",
            "numeric list range",
        ),
    )
    def test_query_with_correct_parameters_predicates(self, predicate, expected_min_length, expected_max_length):
        """
        Invoke query() with correct arguments and using predicate contains
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(predicate)

        records = []

        def callback(input_tuple):
            _, _, record = input_tuple
            records.append(record)

        query.foreach(callback)

        # print("TestBaseClass.major_ver:", TestBaseClass.major_ver, "TestBaseClass.minor_ver:", TestBaseClass.minor_ver)  # noqa: E501
        assert len(records) >= expected_min_length or len(records) <= expected_max_length

    def test_query_with_correct_parameters_exp(self):
        """
        Invoke query() with correct arguments and using expressions
        """

        from .test_base_class import TestBaseClass

        if TestBaseClass.major_ver >= 6 or (TestBaseClass.major_ver >= 5 and TestBaseClass.minor_ver >= 7):
            # print("TestBaseClass.major_ver:", TestBaseClass.major_ver, "TestBaseClass.minor_ver:", TestBaseClass.minor_ver)  # noqa: E501
            pytest.skip("It deprecated and it only applies to < 5.7 earlier and enterprise edition")

        expr = exp.Eq(exp.IntBin("test_age"), 4)

        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        # query.where(predicate)

        records = []

        def callback(input_tuple):
            _, _, record = input_tuple
            records.append(record)

        query.foreach(callback, {"expressions": expr.compile()})
        assert len(records) == 1
        assert records[0]["test_age"] == 4

    def test_query_with_multiple_foreach_on_same_query_object(self):
        """
        Invoke query() with multple foreach() call on same query object
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))

        records = []

        def callback(input_tuple):
            key, _, _ = input_tuple
            records.append(key)

        query.foreach(callback)
        assert len(records) == 1

        query.foreach(callback)
        assert len(records) == 2

    #  All of these removed tests feature undocumented fragile behavior removed in 3.0.0
    def test_removed_query_with_correct_parameters_rangecontains_notuple(self):
        """ """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        with pytest.raises(TypeError):
            query.where("numeric_map", "range", aerospike.INDEX_TYPE_MAPVALUES, aerospike.INDEX_NUMERIC, 1, 3)

    def test_removed_query_with_correct_parameters_containsstring_mapvalues_notuple(self):
        """
        Invoke query() with mapvalue string index
        and using predicate contains
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        with pytest.raises(TypeError):
            query.where("string_map", "contains", aerospike.INDEX_TYPE_MAPVALUES, aerospike.INDEX_STRING, "a1")

    def test_test_removed_query_containsstring_notuple(self):
        """
        Invoke query() with correct a
        and a string list index arguments and using predicate contains
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        with pytest.raises(TypeError):
            query.where("string_list", "contains", aerospike.INDEX_TYPE_LIST, aerospike.INDEX_STRING, "str3")

    def test_removed_query_with_correct_parameters_between_notuple(self):
        """
        Invoke query() with correct arguments and using predicate between
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        with pytest.raises(TypeError):
            query.where("test_age", "between", 1, 4)

    def test_between_predicate_between_one_arg(self):
        """
        Invoke query and using predicate between with invalid predicate
        arguments
        """
        query = self.as_connection.query("test", "demo")
        with pytest.raises(TypeError):
            query.where("test_age", "between", 1)

    def test_between_predicate_between_no_args(self):
        """
        Invoke query and using predicate between with invalid predicate
        arguments
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        with pytest.raises(e.ParamError):
            query.where("test_age", "between")

    # TODO: incorrect test
    def test_query_with_policy_notuple(self):
        """
        Invoke query() with policy
        """
        policy = {"total_timeout": 1000}  # noqa: F841
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        with pytest.raises(TypeError):
            query.where("test_age", "equals", 1)

    def test_query_with_multiple_results_call_on_same_query_object(self):
        """
        Invoke query() with multple results() call on same query object
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age", "addr")
        query.where(p.equals("test_age", 7))

        records = query.results()
        assert len(records) == 1
        assert records[0][2] == {"test_age": 7, "name": "name7", "addr": "name7"}

        records = []
        records = query.results()
        assert len(records) == 1
        assert records[0][2] == {"test_age": 7, "name": "name7", "addr": "name7"}

    def test_query_with_policy_on_none_set_index(self):
        """
        Invoke query() with set specified as None
        """
        policy = {}
        query = self.as_connection.query("test", None)
        query.select("name", "test_age_none")
        query.where(p.equals("test_age_none", 1))
        records = []

        def callback(input_tuple):
            _, _, record = input_tuple
            records.append(record)

        # This should only select records in the ns with no set
        # that match the where clause, if it matched all in the ns,
        # this would return 6
        query.foreach(callback, policy)
        if TestBaseClass.major_ver < 6 or (TestBaseClass.major_ver == 6 and TestBaseClass.minor_ver == 0):
            assert len(records) == 1
        else:
            assert len(records) == 6

    def test_query_with_only_ns(self):
        """
        Invoke query() with only ns specified
        """
        policy = {}
        query = self.as_connection.query("test")
        query.select("name", "test_age_none")
        query.where(p.equals("test_age_none", 1))
        records = []

        def callback(input_tuple):
            _, _, record = input_tuple
            records.append(record)

        # This should only select records in the ns with no set
        # that match the where clause, if it matched all in the ns,
        # this would return 6
        query.foreach(callback, policy)
        if TestBaseClass.major_ver < 6 or (TestBaseClass.major_ver == 6 and TestBaseClass.minor_ver == 0):
            assert len(records) == 1
        else:
            assert len(records) == 6

    def test_query_with_select_bytearray(self):
        """
        Invoke query() with correct arguments
        and select and equality predicate comparing a byte array
        named bin
        """
        query = self.as_connection.query("test", "demo")
        query.select(bytearray("sal\0kj", "utf-8"))
        query.where(p.equals(bytearray("sal\0kj", "utf-8"), 80000))

        records = []

        def callback(input_tuple):
            key, _, _ = input_tuple
            records.append(key)

        query.foreach(callback)
        assert records
        assert len(records) == 1

    def test_query_with_no_parameters(self):
        """
        Invoke query() without any mandatory parameters.
        """
        with pytest.raises(e.ParamError) as err_info:
            self.as_connection.query()

        err_code = err_info.value.code
        assert err_code == AerospikeStatus.AEROSPIKE_ERR_PARAM

    def test_query_with_incorrect_policy_value(self):
        """
        Invoke query() with invalid policy passed to foreach
        """
        policy = {"total_timeout": ""}
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))

        def callback(input_tuple):
            _, metadata, _ = input_tuple
            print(metadata)
            assert metadata["gen"] is None

        with pytest.raises(e.ParamError) as err_info:
            query.foreach(callback, policy)

        err_code = err_info.value.code
        assert err_code == AerospikeStatus.AEROSPIKE_ERR_PARAM

    def test_query_with_select_bin_integer(self):
        """
        Invoke query() with select bin is of integer type.
        Try to select a bin with an integer name instead of a string
        """
        query = self.as_connection.query("test", "demo")

        with pytest.raises(e.ParamError) as err_info:
            query.select(22, "test_age")

        err_code = err_info.value.code
        assert err_code == AerospikeStatus.AEROSPIKE_ERR_PARAM

    def test_query_select_multiple_times(self):
        query: aerospike.Query = self.as_connection.query("test", "demo")
        query.select("a")
        with pytest.raises(e.ClientError):
            query.select("a")

    def test_query_with_argument_to_where_is_empty_string(self):

        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")

        with pytest.raises(e.ParamError) as err_info:
            query.where("")

        err_code = err_info.value.code
        assert err_code == AerospikeStatus.AEROSPIKE_ERR_PARAM

    def test_query_with_callback_contains_error(self):
        """
        Invoke query() with callback function contains an error
        """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where(p.equals("test_age", 1))

        def callback(input_tuple):
            raise Exception("error")

        with pytest.raises(e.ClientError) as err_info:
            query.foreach(callback)

        err_code = err_info.value.code
        assert err_code == AerospikeStatus.AEROSPIKE_ERR_CLIENT

    def test_query_with_incorrect_ns_set(self):
        """
        Invoke query() with incorrect ns and set
        foreach should raise an error since the ns
        is invalid
        """
        with pytest.raises(e.NamespaceNotFound) as err_info:
            query = self.as_connection.query("fake_namespace", "demo1")

            def callback(input_tuple):
                """
                no-op callback
                """

            query.foreach(callback)

        err_code = err_info.value.code
        assert err_code == AerospikeStatus.AEROSPIKE_ERR_NAMESPACE_NOT_FOUND

    def test_query_with_ns_not_string(self):
        """
        Invoke query() with incorrect ns and set. It
        should raises a param error.
        """
        with pytest.raises(e.ParamError) as err_info:
            self.as_connection.query(1, "demo")

        err_code = err_info.value.code
        assert err_code == AerospikeStatus.AEROSPIKE_ERR_PARAM

    def test_query_with_set_int(self):
        """
        Invoke query() with set argument not a string
        """
        with pytest.raises(e.ParamError) as err_info:
            self.as_connection.query("test", 1)

        err_code = err_info.value.code
        assert err_code == AerospikeStatus.AEROSPIKE_ERR_PARAM

    def test_query_with_correct_parameters_without_connection(self):
        """
        Invoke query() with correct arguments without connection
        """
        config = TestBaseClass.get_connection_config()
        client1 = aerospike.client(config)
        client1.close()

        with pytest.raises(e.ClusterError) as err_info:
            query = client1.query("test", "demo")
            query.select("name", "test_age")
            query.where(p.equals("test_age", 1))

            def callback(input_tuple):
                pass

            query.foreach(callback)

        err_code = err_info.value.code
        assert err_code == AerospikeStatus.AEROSPIKE_CLUSTER_ERROR

    @pytest.mark.skip(reason="segfault")
    def test_query_predicate_range_wrong_no_args(self):
        """ """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where("numeric_map", "range", aerospike.INDEX_TYPE_MAPVALUES, aerospike.INDEX_NUMERIC)

    @pytest.mark.skip(reason="segfault")
    def test_query_predicate_range_wrong_one_end_args(self):
        """ """
        query = self.as_connection.query("test", "demo")
        query.select("name", "test_age")
        query.where("numeric_map", "range", aerospike.INDEX_TYPE_MAPVALUES, aerospike.INDEX_NUMERIC, 1)

    def test_query_with_list_cdt_ctx(self):
        """
        Invoke query() with cdt_ctx and correct arguments
        """
        from .test_base_class import TestBaseClass

        if TestBaseClass.major_ver < 6 or (TestBaseClass.major_ver == 6 and TestBaseClass.minor_ver == 0):
            pytest.skip("It only applies to >= 6.1 enterprise edition")

        # ctx = []
        # ctx.append(add_ctx_op(list_index, 0))
        # query = self.as_connection.query('test', 'demo')
        # query.select('numeric_list')
        # query.where(p.range('numeric_list', aerospike.INDEX_TYPE_DEFAULT, 2,4), {'ctx':ctx})

        query = self.as_connection.query("test", "demo")
        query.select("numeric_list")
        query.where(p.range("numeric_list", aerospike.INDEX_TYPE_DEFAULT, 2, 4), {"ctx": ctx_list_index})

        records = []

        def callback(input_tuple):
            try:
                records.append(input_tuple)
            except Exception as ex:
                print(ex)

        query.foreach(callback)

        assert records
        assert len(records) == 3

    def test_query_with_list_cdt_ctx_and_invalid_bin(self):
        """
        Make sure that ctx is being cleaned up properly
        """
        query = self.as_connection.query("test", "demo")
        # Invalid bin
        with pytest.raises(e.ParamError):
            query.where(p.range(5, aerospike.INDEX_TYPE_DEFAULT, 2, 4), {"ctx": ctx_list_index})

    def test_query_with_invalid_predicate_type(self):
        query = self.as_connection.query("test", "demo")
        with pytest.raises(e.ParamError):
            # 2 is outside valid values for predicate types
            query.where((2, aerospike.INDEX_BLOB, "bin", bytearray(b'123')))

    def test_query_with_map_cdt_ctx(self):
        """
        Invoke query() with cdt_ctx and correct arguments
        """
        from .test_base_class import TestBaseClass

        if TestBaseClass.major_ver < 6 or (TestBaseClass.major_ver == 6 and TestBaseClass.minor_ver == 0):
            pytest.skip("It only applies to >= 6.1 enterprise edition")

        query = self.as_connection.query("test", "demo")
        query.select("numeric_map")

        query.where(p.range("numeric_map", aerospike.INDEX_TYPE_DEFAULT, 2, 4), {"ctx": ctx_map_index})

        records = []

        def callback(input_tuple):
            try:
                records.append(input_tuple)
            except Exception as ex:
                print(ex)

        query.foreach(callback)

        assert records
        assert len(records) == 3

    def test_query_with_base64_cdt_ctx(self):
        bs_b4_cdt = self.as_connection.get_cdtctx_base64(ctx_list_index)
        assert bs_b4_cdt == "khAA"

    def test_query_blob_bin_equal(self):
        if (TestBaseClass.major_ver, TestBaseClass.minor_ver) < (7, 0):
            pytest.skip("Blob indexes are only supported in server 7.0+")

        self.as_connection.index_blob_create("test", "demo", "blob", "blob_index")

        query = self.as_connection.query("test", "demo")
        blob_val = int.to_bytes(4, length=1, byteorder='big')
        query.where(p.equals("blob", blob_val))

        records = []

        def callback(input_tuple):
            try:
                records.append(input_tuple)
            except Exception as ex:
                print(ex)

        query.foreach(callback)

        assert records
        assert len(records) == 1

        self.as_connection.index_remove("test", "blob_index")

    @pytest.mark.parametrize(
        "indexed_bin_name, index_type",
        [
            ("blob_list", aerospike.INDEX_TYPE_LIST),
            ("blob_map", aerospike.INDEX_TYPE_MAPKEYS),
            ("blob_map", aerospike.INDEX_TYPE_MAPVALUES),
        ]
    )
    def test_query_blob_bin_contains(self, indexed_bin_name, index_type):
        if (TestBaseClass.major_ver, TestBaseClass.minor_ver) < (7, 0):
            pytest.skip("Blob indexes are only supported in server 7.0+")

        # Index creation is in the setup function
        query = self.as_connection.query("test", "demo")
        # In the previous test, we passed in a bytes for the predicate's value parameter
        # Now pass in a bytearray (which is also valid)
        blob_val = int.to_bytes(1, length=1, byteorder='big')
        blob_val = bytearray(blob_val)
        query.where(p.contains(indexed_bin_name, index_type, blob_val))

        records = []

        def callback(input_tuple):
            try:
                records.append(input_tuple)
            except Exception as ex:
                print(ex)

        query.foreach(callback)

        assert len(records) == 1

    def test_query_with_none_bin_and_where(self):
        query = self.as_connection.query("test", "demo")
        with pytest.raises(e.ParamError):
            query.where(p.equals(None, 1))

    @pytest.mark.parametrize(
        "duration",
        [
            aerospike.QUERY_DURATION_LONG,
            aerospike.QUERY_DURATION_SHORT,
            aerospike.QUERY_DURATION_LONG_RELAX_AP
        ]
    )
    def test_query_expected_duration(self, duration: int):
        if duration == aerospike.QUERY_DURATION_LONG_RELAX_AP and TestBaseClass.strong_consistency_enabled:
            pytest.skip("Using aerospike.QUERY_DURATION_LONG_RELAX_AP will fail if server is in SC mode")
        query: aerospike.Query = self.as_connection.query("test", "demo")
        policy = {
            "expected_duration": duration
        }
        query.results(policy=policy)

    def test_query_invalid_expected_duration(self):
        query: aerospike.Query = self.as_connection.query("test", "demo")
        policy = {
            "expected_duration": "t"
        }
        with pytest.raises(e.ParamError) as excinfo:
            query.results(policy=policy)
        assert excinfo.value.msg == "expected_duration is invalid"

    def test_query_with_invalid_expr(self):
        query: aerospike.Query = self.as_connection.query("test", "demo")
        with pytest.raises(e.ParamError):
            query.where_with_expr(4, p.equals("test_age", 165))

    INT_BIN_EXPR = Add(IntBin("test_age"), IntBin("no"))
    GEO_POLYGON_BIN_EXPR = GeoBin("geo_polygon")
    GEO_POINT_BIN_EXPR = GeoBin("geo_point")

    INDEX_EXPR_NAME = "index_expr"

    # Should contain each records' geo_point bin geographically
    # Note: we have to convert meters so that it covers the latitude and longitude (they are in different units)
    CIRCLE_RADIUS_METERS = 629000
    GEOJSON_CIRCLE = aerospike.GeoJSON({"type": "AeroCircle", "coordinates": [[0, 0], CIRCLE_RADIUS_METERS]})

    LIST_EXPR = ListBin("numeric_list")
    BLOB_EXPR = BlobBin("blob")

    @pytest.fixture
    def index_expr_cleanup(self):
        yield
        self.as_connection.index_remove("test", self.INDEX_EXPR_NAME)

    @pytest.mark.parametrize(
        "expr, index_type, index_datatype, predicate, expected_rec_count",
        [
            # Test every predicate to make sure it accepts a bin name of None

            # Only the first record
            (INT_BIN_EXPR, aerospike.INDEX_TYPE_DEFAULT, aerospike.INDEX_NUMERIC, p.equals(None, 2), 1),
            # The first two records
            (INT_BIN_EXPR, aerospike.INDEX_TYPE_DEFAULT, aerospike.INDEX_NUMERIC, p.between(None, 0, 2), 2),
            (
                GEO_POLYGON_BIN_EXPR,
                aerospike.INDEX_TYPE_DEFAULT,
                aerospike.INDEX_GEO2DSPHERE,
                p.geo_contains_geojson_point(
                    None,
                    # Only the 25x25 box should contain this point
                    aerospike.GeoJSON({"type": "Point", "coordinates": [23, 23]}).dumps()
                ),
                1
            ),
            # Same test as above, but with a different predicate
            (
                GEO_POLYGON_BIN_EXPR,
                aerospike.INDEX_TYPE_DEFAULT,
                aerospike.INDEX_GEO2DSPHERE, p.geo_contains_point(None, 23, 23), 1),
            (
                GEO_POINT_BIN_EXPR,
                aerospike.INDEX_TYPE_DEFAULT,
                aerospike.INDEX_GEO2DSPHERE,
                p.geo_within_geojson_region(None, GEOJSON_CIRCLE.dumps()),
                # The circle should cover all 5 points
                5
            ),
            # Same test as above but with a different pred
            (
                GEO_POINT_BIN_EXPR,
                aerospike.INDEX_TYPE_DEFAULT,
                aerospike.INDEX_GEO2DSPHERE,
                p.geo_within_radius(None, 0, 0, CIRCLE_RADIUS_METERS),
                5
            ),
            (
                LIST_EXPR,
                aerospike.INDEX_TYPE_LIST,
                aerospike.INDEX_NUMERIC,
                p.contains(None, aerospike.INDEX_TYPE_LIST, 0),
                # Only the first record
                1
            ),
            (
                LIST_EXPR,
                aerospike.INDEX_TYPE_LIST,
                aerospike.INDEX_NUMERIC,
                p.range(None, aerospike.INDEX_TYPE_LIST, 0, 1),
                # Only the first two records
                2
            ),
            # Test blobs with where_with_*() calls for code coverage
            (
                BLOB_EXPR,
                aerospike.INDEX_TYPE_DEFAULT,
                aerospike.INDEX_BLOB,
                p.equals(None, (0).to_bytes(length=1, byteorder='big')),
                # Only the first record
                1
            )
        ]
    )
    @pytest.mark.parametrize("use_index_name", [False, True])
    def test_query_with_expr_or_index_name(
        self,
        expr,
        index_type,
        index_datatype,
        predicate,
        expected_rec_count,
        index_expr_cleanup,
        use_index_name
    ):
        if (TestBaseClass.major_ver, TestBaseClass.minor_ver) < (8, 1):
            pytest.skip("Querying with expressions isn't supported yet")

        expr = expr.compile()
        self.as_connection.index_expr_create(ns="test", set="demo", index_type=index_type,
                                             index_datatype=index_datatype,
                                             expressions=expr, name=self.INDEX_EXPR_NAME, policy=None)

        # Verify where_*() methods return a Query object as well
        query: aerospike.Query = self.as_connection.query("test", "demo")
        if use_index_name:
            query = query.where_with_index_name(self.INDEX_EXPR_NAME, predicate)
        else:
            query = query.where_with_expr(expr, predicate)

        recs = query.results()
        assert len(recs) == expected_rec_count

        # We should also be able to query using the base64 encoded string for an expression
        query2: aerospike.Query = self.as_connection.query("test", "demo")
        if use_index_name is False:
            expr_base64_encoded = self.as_connection.get_expression_base64(expr)
            query2 = query2.where_with_expr(expr_base64_encoded, predicate)
            recs = query2.results()
            assert len(recs) == expected_rec_count
