from includes import *
try:
    from collections.abc import Iterable
except ImportError:
    from collections import Iterable
import time
from packaging import version
from functools import wraps
import signal
import platform
import itertools
import threading
from redis.client import NEVER_DECODE
from redis import exceptions as redis_exceptions
import RLTest
from typing import Any, Callable, List, Dict
from RLTest import Env, env_spec
from RLTest.env import Query
import numpy as np
from scipy import spatial
from pprint import pprint as pp
from deepdiff import DeepDiff
from unittest.mock import ANY, _ANY
from unittest import SkipTest
import inspect
import math
import tempfile
import faker

TEST_RDBS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'test_rdbs')
REDISEARCH_CACHE_DIR = os.path.join(tempfile.gettempdir(), 'redisearch-rdbs')
VECSIM_DATA_TYPES = ['FLOAT32', 'FLOAT64', 'FLOAT16', 'BFLOAT16']
VECSIM_ALGOS = ['FLAT', 'HNSW', 'SVS-VAMANA']

class TimeLimit(object):
    """
    A context manager that fires a TimeExpired exception if it does not
    return within the specified amount of time.
    """

    def __init__(self, timeout, message='operation timeout exceeded'):
        self.timeout = timeout
        self.message = message

    def __enter__(self):
        self.time_start = time.time()
        signal.signal(signal.SIGALRM, self.handler)
        signal.setitimer(signal.ITIMER_REAL, self.timeout, 0)

    def __exit__(self, exc_type, exc_value, traceback):
        signal.setitimer(signal.ITIMER_REAL, 0)
        signal.signal(signal.SIGALRM, signal.SIG_DFL)

    def handler(self, signum, frame):
        raise Exception(f'Timeout: {self.message} + after {time.time() - self.time_start}s')

def wait_for_condition(check_fn, message, timeout=120):
    """
    Wait for a condition with timeout and status reporting.

    Parameters:
        - env: Test environment
        - check_fn: Function that takes returns (status: bool, state: dict)
                   where state is a dict of the current state information
        - message: Message prefix for timeout exception
    """
    iter = 0
    timeout_msg = {}

    try:
        with TimeLimit(timeout):
            while True:
                done, state = check_fn()
                if done:
                    break
                time.sleep(0.01)
                iter += 1
                timeout_msg['iter'] = iter
                timeout_msg['state'] = state
    except Exception as e:
        log = f"{message}: {timeout_msg}"
        raise Exception(f'Error: {e}, log: {log}')

class DialectEnv(Env):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dialect = None

    def set_dialect(self, dialect):
        self.dialect = dialect
        result = run_command_on_all_shards(self, config_cmd(), 'SET', 'DEFAULT_DIALECT', dialect)
        expected_result = ['OK'] * self.shardsCount
        self.assertEqual(result, expected_result, message=f"Failed to set dialect to {dialect} on all shards")

    def get_dialect(self):
        return self.dialect

    def assertEqual(self, first, second, depth=0, message=None):
        if self.dialect is not None:
            if message is None:
                message = f'Dialect {self.dialect}'
            else:
                message = f'Dialect {self.dialect}, {message}'
        super().assertEqual(first, second, depth=depth+1, message=message)

def getConnectionByEnv(env):
    conn = None
    if env.env == 'oss-cluster':
        conn = env.envRunner.getClusterConnection()
    else:
        conn = env.getConnection()
    return conn

def waitForIndex(env, idx = 'idx'):
    waitForRdbSaveToFinish(env)
    while True:
        res = env.cmd('ft.info', idx)
        try:
            if res[res.index('indexing') + 1] == 0:
                break
        except:
            # RESP3
            if res['indexing'] == 0:
                break
        time.sleep(0.1)

def waitForNoCleanup(env, idx, max_wait=30):
    ''' Wait for the index to finish cleanup

    Parameters:
        max_wait - max duration in seconds to wait
    '''
    waitForRdbSaveToFinish(env)
    retry_wait = 0.1
    max_wait = max(max_wait, retry_wait)
    while max_wait >= 0:
        res = env.cmd('ft.info', idx)
        if int(res[res.index('cleaning') + 1]) == 0:
            break
        time.sleep(retry_wait)
        max_wait -= retry_wait

def py2sorted(x):
    it = iter(x)
    groups = [[next(it)]]
    for item in it:
        for group in groups:
            try:
                item < group[0]  # exception if not comparable
                group.append(item)
                break
            except TypeError:
                continue
        else:  # did not break, make new group
            groups.append([item])
    # print(groups)  # for debugging
    return list(itertools.chain.from_iterable(sorted(group) for group in groups))

def toSortedFlatList(res):
    if isinstance(res, str):
        return [res]
    if isinstance(res, Iterable):
        finalList = []
        for e in res:
            finalList += toSortedFlatList(e)

        return py2sorted(finalList)
    return [res]

def countFlatElements(arr):
    """Count elements without sorting (lighter than toSortedFlatList)"""
    if isinstance(arr, str):
        return 1
    if isinstance(arr, Iterable):
        count = 0
        for e in arr:
            count += countFlatElements(e)
        return count
    return 1

def assertInfoField(env, idx, field, expected, delta=None):
    d = index_info(env, idx)
    msg = f"field name: {field}"
    if delta is None:
        env.assertEqual(d[field], expected, message = msg)
    else:
        env.assertAlmostEqual(float(d[field]), float(expected), delta=delta, message = msg)

def sortedResults(res):
    n = res[0]
    res = res[1:]

    y = []
    data = []
    for x in res:
        y.append(x)
        if len(y) == 2:
            data.append(y)
            y = []

    data = py2sorted(data)
    res = [n] + [item for sublist in data for item in sublist]
    return res

def slice_at(v, val):
    try:
        i = v.index(val)
        return v[i+1:]
    except:
        return []

def numver_to_version(numver):
    v = numver
    v = "%d.%d.%d" % (int(v/10000), int(v/100)%100, v%100)
    return version.parse(v)

def arch_int_bits():
  arch = platform.machine()
  if arch == 'x86_64':
    return 128
  elif arch == 'aarch64':
    return 128
  elif arch == 'arm64':
    return 128
  else:
    return 64

module_ver = None
def module_version_at_least(env, ver):
    global module_ver
    if module_ver is None:
        v = env.cmd('MODULE LIST')[0][3]
        module_ver = numver_to_version(v)
    if not isinstance(ver, version.Version):
        ver = version.parse(ver)
    return module_ver >= ver

def module_version_less_than(env, ver):
    return not module_version_at_least(env, ver)

server_ver = None
def server_version_at_least(env: Env, ver):
    global server_ver
    if server_ver is None:
        v = env.cmd('INFO')['redis_version']
        server_ver = version.parse(v)
    if not isinstance(ver, version.Version):
        ver = version.parse(ver)
    return server_ver >= ver

def server_version_less_than(env: Env, ver):
    return not server_version_at_least(env, ver)

def server_version_is_at_least(ver):
    global server_ver
    if server_ver is None:
        import subprocess
        # Expecting something like "Redis server v=7.2.3 sha=******** malloc=jemalloc-5.3.0 bits=64 build=***************"
        v = subprocess.run([Defaults.binary, '--version'], stdout=subprocess.PIPE).stdout.decode().split()[2].split('=')[1]
        server_ver = version.parse(v)
    if not isinstance(ver, version.Version):
        ver = version.parse(ver)
    return server_ver >= ver

def server_version_is_less_than(ver):
    return not server_version_is_at_least(ver)

def index_info(env, idx='idx'):
    res = env.cmd('FT.INFO', idx)
    return to_dict(res)


def dump_numeric_index_tree(env, idx, numeric_field):
    tree_dump = env.cmd(debug_cmd(), 'DUMP_NUMIDXTREE', idx, numeric_field)
    return to_dict(tree_dump)


def dump_numeric_index_tree_root(env, idx, numeric_field):
    tree_root_stats = dump_numeric_index_tree(env, idx, numeric_field)['root']
    root_dump = {tree_root_stats[i]: tree_root_stats[i + 1]
                 for i in range(0, len(tree_root_stats), 2)}
    return root_dump

def numeric_tree_summary(env, idx, numeric_field):
    tree_summary = env.cmd(debug_cmd(), 'NUMIDX_SUMMARY', idx, numeric_field)
    return to_dict(tree_summary)


def getWorkersThpoolStats(env):
    return to_dict(env.cmd(debug_cmd(), "WORKERS", "stats"))

def getWorkersThpoolNumThreads(env):
    return env.cmd(debug_cmd(), "WORKERS", "n_threads")

def set_workers(env, workers):
    """Set the worker thread count and verify that the change took effect."""
    verify_command_OK_on_all_shards(env, config_cmd(), 'SET', 'WORKERS', workers)
    env.assertEqual(getWorkersThpoolNumThreadsFromAllShards(env), [workers] * env.shardsCount)

def getWorkersThpoolStatsFromShard(shard_conn):
    return to_dict(shard_conn.execute_command(debug_cmd(), "WORKERS", "stats"))

def getCoordThpoolStats(env):
    return to_dict(env.cmd(debug_cmd(), "COORD_THREADS", "stats"))

def getWorkersThpoolStatsFromAllShards(env):
    return [getWorkersThpoolStatsFromShard(shard_conn) for shard_conn in env.getOSSMasterNodesConnectionList()]

def getWorkersThpoolNumThreadsFromAllShards(env):
    return [shard_conn.execute_command(debug_cmd(), "WORKERS", "n_threads") for shard_conn in env.getOSSMasterNodesConnectionList()]

def skipOnExistingEnv(env):
    if 'existing' in env.env:
        env.skip()

def SkipOnNonCluster(env):
    if not env.isCluster():
        env.skip()

def skipOnCrdtEnv(env):
    if len([a for a in env.cmd('module', 'list') if a[1] == 'crdt']) > 0:
        env.skip()

def skipOnDialect(env, dialect):
    server_dialect = int(env.expect(config_cmd(), 'GET', 'DEFAULT_DIALECT').res[0][1])
    if dialect == server_dialect:
        env.skip()

def waitForRdbSaveToFinish(env):
    if env.isCluster():
        conns = env.getOSSMasterNodesConnectionList()
    else:
        conns = [env.getConnection()]

    # Busy wait until all connection are done rdb bgsave
    check_bgsave = True
    while check_bgsave:
        check_bgsave = False
        for conn in conns:
            if conn.execute_command('info', 'Persistence')['rdb_bgsave_in_progress']:
                check_bgsave = True
                break


def countKeys(env, pattern='*'):
    if not env.isCluster():
        return len(env.keys(pattern))
    keys = 0
    for shard in range(0, env.shardsCount):
        conn = env.getConnection(shard)
        keys += len(conn.keys(pattern))
    return keys

def collectKeys(env, pattern='*'):
    if not env.isCluster():
        return sorted(env.keys(pattern))
    keys = []
    for shard in range(0, env.shardsCount):
        conn = env.getConnection(shard)
        keys.extend(conn.keys(pattern))
    return sorted(keys)


def debug_cmd():
    return '_FT.DEBUG'

def config_cmd():
    return '_FT.CONFIG'

def enable_unstable_features(env):
    run_command_on_all_shards(env, 'CONFIG', 'SET', 'search-enable-unstable-features', 'yes')

def run_command_on_all_shards(env, *args):
    return [con.execute_command(*args) for con in env.getOSSMasterNodesConnectionList()]

def verify_command_OK_on_all_shards(env, *args):
    res = run_command_on_all_shards(env, *args)
    env.assertEqual(res, ['OK'] * env.shardsCount)

def allShards_set_info_on_zero_indexes(env, enabled: bool):
    """
    Enable/disable INFO MODULES full output when there are zero indexes.

    In cluster mode, applies to all OSS shards. In standalone mode, applies to the single node.
    Asserts success (all replies are OK).
    """
    val = 'yes' if enabled else 'no'
    if env.isCluster():
        verify_command_OK_on_all_shards(env, 'CONFIG', 'SET', 'search-_info-on-zero-indexes', val)
        return
    res = env.cmd('CONFIG', 'SET', 'search-_info-on-zero-indexes', val)
    env.assertEqual(res, 'OK')
    return

def shard_set_info_on_zero_indexes(env, enabled: bool):
    """
    Enable/disable INFO MODULES full output when there are zero indexes on the current node.

    Uses `getConnectionByEnv(env)` so callers don't need to pass a shard id.
    """
    val = 'yes' if enabled else 'no'
    conn = getConnectionByEnv(env)
    res = conn.execute_command('CONFIG', 'SET', 'search-_info-on-zero-indexes', val)
    env.assertEqual(res, 'OK')
    return res

def get_vecsim_debug_dict(env, index_name, vector_field):
    return to_dict(env.cmd(debug_cmd(), "VECSIM_INFO", index_name, vector_field))

def forceInvokeGC(env, idx='idx', timeout=None):
    waitForRdbSaveToFinish(env)
    if timeout is not None:
        # Note: timeout==0 means infinite (no timeout)
        env.cmd(debug_cmd(), 'GC_FORCEINVOKE', idx, timeout)
    else:
        env.cmd(debug_cmd(), 'GC_FORCEINVOKE', idx)

def forceBGInvokeGC(env, idx='idx'):
    waitForRdbSaveToFinish(env)
    env.cmd(debug_cmd(), 'GC_FORCEBGINVOKE', idx)

def no_msan(f):
    @wraps(f)
    def wrapper(env, *args, **kwargs):
        if SANITIZER == 'memory':
            fname = f.__name__
            env.debugPrint(f"skipping {fname} due to memory sanitizer", force=True)
            env.skip()
            return
        return f(env, *args, **kwargs)
    return wrapper

def unstable(f):
    @wraps(f)
    def wrapper(env, *args, **kwargs):
        if UNSTABLE == True:
            fname = f.__name__
            env.debugPrint(f"skipping {fname} because it is unstable", force=True)
            env.skip()
            return
        return f(env, *args, **kwargs)
    return wrapper

# Wraps the decorator `skip` for calling from within a test function
def skipTest(**kwargs):
    skip(**kwargs)(lambda: None)()

def skip_until(date_str, reason=None):
    """
    Decorator to skip a test until a specific date.
    After the date passes, the test will run normally.

    This is useful for temporarily skipping flaky tests while ensuring
    they are not forgotten - the test will automatically start running
    again after the specified date.

    Args:
        date_str: A date string in ISO format "YYYY-MM-DD" (e.g., "2024-06-15")
        reason: Optional reason for skipping the test

    Usage:
        @skip_until("2024-06-15", reason="Flaky test, investigating MOD-1234")
        def testSomething(env):
            ...
    """
    from datetime import datetime

    def decorate(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            skip_date = datetime.strptime(date_str, "%Y-%m-%d").date()
            today = datetime.now().date()
            if today < skip_date:
                reason_msg = f" ({reason})" if reason else ""
                print(f"Skipping {f.__name__} until {date_str}{reason_msg}")
                raise SkipTest(f"Skipped until {date_str}{reason_msg}")
            # Date has passed, run the test
            return f(*args, **kwargs)
        return wrapper
    return decorate

# Wraps the decorator `skip_until` for calling from within a test function
def skipTestUntil(date_str, reason=None):
    """
    Skip the current test until a specific date.
    Call this from within a test function.

    Args:
        date_str: A date string in ISO format "YYYY-MM-DD" (e.g., "2024-06-15")
        reason: Optional reason for skipping the test

    Usage:
        def testSomething(env):
            if some_condition:
                skipTestUntil("2024-06-15", reason="Flaky under certain conditions")
            ...
    """
    skip_until(date_str, reason)(lambda: None)()

def _any_skip_condition_set(*, cluster, macos, asan, msan, redis_less_than,
                            redis_greater_equal, min_shards, arch, gc_no_fork,
                            no_json):
    """True if the caller provided at least one skip condition.

    With no conditions, @skip's legacy behaviour is to always skip — used as a
    "temporarily disable this test" marker.
    """
    return ((cluster is not None) or macos or asan or msan or redis_less_than
            or redis_greater_equal or min_shards or (arch is not None)
            or gc_no_fork or no_json)


def _skip_fires_statically(*, cluster, macos, asan, msan, min_shards, arch,
                            no_json):
    """Evaluate the subset of @skip predicates that don't need a live Redis.

    Excludes redis_less_than/redis_greater_equal/gc_no_fork — those genuinely
    need a running env, so they can't decide at module-load time.
    """
    if cluster is not None and cluster == CLUSTER:
        return True
    if macos and OS == 'macos':
        return True
    if arch == platform.machine():
        return True
    if asan and SANITIZER == 'address':
        return True
    if msan and SANITIZER == 'memory':
        return True
    if min_shards and Defaults.num_shards < min_shards:
        return True
    if no_json and not REJSON:
        return True
    return False


def _skip_fires_at_runtime(*, redis_less_than, redis_greater_equal, gc_no_fork):
    """Evaluate the @skip predicates that need a live Redis.

    Spawns a transient Env for gc_no_fork; the version helpers maintain their
    own connection.
    """
    if redis_less_than and server_version_is_less_than(redis_less_than):
        return True
    if redis_greater_equal and server_version_is_at_least(redis_greater_equal):
        return True
    if gc_no_fork and Env().cmd(config_cmd(), 'GET', 'GC_POLICY')[0][1] != 'fork':
        return True
    return False


def skip(cluster=None, macos=False, asan=False, msan=False, redis_less_than=None, redis_greater_equal=None, min_shards=None, arch=None, gc_no_fork=None, no_json=False):
    static_kwargs = dict(cluster=cluster, macos=macos, asan=asan, msan=msan,
                         min_shards=min_shards, arch=arch, no_json=no_json)
    runtime_kwargs = dict(redis_less_than=redis_less_than,
                          redis_greater_equal=redis_greater_equal,
                          gc_no_fork=gc_no_fork)
    any_condition = _any_skip_condition_set(
        **static_kwargs, **runtime_kwargs,
    )

    def decorate(target):
        if isinstance(target, type):
            # Class decoration. We must decide whether to skip BEFORE RLTest
            # provisions a class's @env_spec env — otherwise a cluster-only
            # test class would spin up a 3-shard env on standalone runs just
            # to be thrown away. Evaluate the static predicates now; if any
            # fires, strip the @env_spec marker (so no env is built) and
            # replace __init__ with one that raises SkipTest the moment
            # RLTest tries to instantiate the class.
            if redis_less_than or redis_greater_equal or gc_no_fork:
                raise TypeError(
                    "@skip on a class cannot use redis_less_than, "
                    "redis_greater_equal, or gc_no_fork because those need a "
                    "live Redis and would defeat the purpose of skipping "
                    "before env provisioning. Use skipTest() inside a method "
                    "instead, or apply @skip to individual functions."
                )
            # No conditions = legacy "always skip" marker.
            fires = (not any_condition) or _skip_fires_statically(**static_kwargs)
            if fires:
                if hasattr(target, '_rltest_env_spec'):
                    delattr(target, '_rltest_env_spec')
                def _skipping_init(self, *args, **kwargs):
                    raise SkipTest()
                target.__init__ = _skipping_init
            return target

        f = target
        def wrapper():
            if not any_condition:
                raise SkipTest()
            if _skip_fires_statically(**static_kwargs):
                raise SkipTest()
            if _skip_fires_at_runtime(**runtime_kwargs):
                raise SkipTest()
            if len(inspect.signature(f).parameters) > 0:
                # Honor a declared @env_spec when constructing the env so the
                # spec stays load-bearing even when @skip is stacked on top.
                spec = getattr(f, '_rltest_env_spec', {})
                env = Env(**spec)
                return f(env)
            else:
                return f()
        # Propagate identifying metadata + the env_spec marker. Deliberately
        # NOT functools.wraps: that would set wrapper.__wrapped__ = f, which
        # exposes f's signature to inspect.signature(follow_wrapped=True) and
        # could trick callers into passing an env arg to this zero-arg
        # wrapper. We only need __name__/__qualname__ for debugging and
        # _rltest_env_spec so RLTest's loader can read the declared spec off
        # the wrapper.
        wrapper.__name__ = f.__name__
        wrapper.__qualname__ = f.__qualname__
        wrapper.__doc__ = f.__doc__
        spec = getattr(f, '_rltest_env_spec', None)
        if spec is not None:
            wrapper._rltest_env_spec = spec
        return wrapper
    return decorate

def to_dict(res):
    if type(res) == dict:
        return res
    if len(res) % 2 != 0:
        raise ValueError(f"to_dict expects even-length array (key-value pairs), got {len(res)} elements")
    d = {res[i]: res[i + 1] for i in range(0, len(res), 2)}
    return d

def to_list(input_dict: dict):
    return [item for pair in input_dict.items() for item in pair]

def get_redis_memory_in_mb(env):
    return float(env.cmd('info', 'memory')['used_memory'])/0x100000

MAX_DIALECT = 0
def set_max_dialect(env):
    global MAX_DIALECT
    if MAX_DIALECT == 0:
        # Ensure INFO MODULES is not in minimal suppression mode when there are zero indexes.
        # This keeps dialect discovery simple and consistent across tests.
        # We only query INFO MODULES on the current connection, so it's enough to set this locally
        # (no need to broadcast to all shards).
        shard_set_info_on_zero_indexes(env, True)

        info = env.cmd('INFO', 'MODULES')
        prefix = 'search_dialect_'
        MAX_DIALECT = max([int(key.replace(prefix, '')) for key in info.keys() if prefix in key])
    return MAX_DIALECT

def get_redisearch_index_memory(env, index_key):
    return float(index_info(env, index_key)["inverted_sz_mb"])

def module_ver_filter(env, module_name, ver_filter):
    info = env.getConnection().info()
    for module in info['modules']:
        if module['name'] == module_name:
            ver = int(module['ver'])
            return ver_filter(ver)
    return False

def has_json_api_v2(env):
    return module_ver_filter(env, 'ReJSON', lambda ver: True if ver == 999999 or ver >= 20200 else False)

# A very simple implementation of a bfloat16 array type.
# wrap a numpy array (for basic operations) and override `tobytes` to convert to bfloat16
# This saves us the need to install a new package for bfloat16 support (e.g. tensorflow, torch, bfloat16 numpy extension)
# and deal with dependencies and compatibility issues.
class Bfloat16Array(np.ndarray):
    offset = 2 if sys.byteorder == 'little' else 0
    def __new__(cls, input_array):
        return np.asarray(input_array).view(cls)

    def tobytes(self):
        b32 = np.ndarray.tobytes(self.astype(np.float32))
        # Generate a byte string from every other pair of bytes in b32
        return b''.join(b32[i:i+2] for i in range(Bfloat16Array.offset, len(b32), 4))

# Helper function to create numpy array vector with a specific type
def create_np_array_typed(data, data_type='FLOAT32'):
    if data_type.upper() == 'BFLOAT16':
        return Bfloat16Array(data)
    return np.array(data, dtype=data_type.lower())

np.random.seed(42)
def create_random_np_array_typed(dim, data_type='FLOAT32', normalize=False):
    vector = create_np_array_typed(np.random.rand(dim), data_type)
    if normalize:
        vector /= np.linalg.norm(vector)
    return vector
def compare_lists_rec(var1, var2, delta):
    if type(var1) != type(var2):
        return False
    try:
        if type(var1) is not str and len(var1) != len(var2):
            return False
    except:
        pass

    if isinstance(var1, list):
        #print("compare_lists_rec: list {}".format(var1))
        for i in range(len(var1)):
            #print("compare_lists_rec: list: i = {}".format(i))
            res = compare_lists_rec(var1[i], var2[i], delta)
            #print("list: var1 = {}, var2 = {}, res = {}".format(var1[i], var2[i], res))
            if res is False:
                return False

    elif isinstance(var1, dict):
        for k in var1:
            res = compare_lists_rec(var1[k], var2[k], delta)
            if res is False:
                return False

    elif isinstance(var1, set):
        for v in var1:
            if v not in var2:
                return False

    elif isinstance(var1, tuple):
        for i in range(len(var1)):
            compare_lists_rec(var1[i], var2[i], delta)
            if res is False:
                return False

    elif isinstance(var1, float):
        diff = var1 - var2
        if diff < 0:
            diff = -diff
        #print("diff {} delta {}".format(diff, delta))
        return diff <= delta

    elif isinstance(var1, str): # float as string
        try:
            diff = float(var1) - float(var2)
            if diff < 0:
                diff = -diff
        except:
            return var1 == var2

        #print("var1 {} var2 {} diff {} delta {}".format(var1, var2, diff, delta))
        return diff <= delta

    else: # int() | bool() | None:
        return var1 == var2

    return True

def compare_lists(env, list1, list2, delta=0.01, _assert=True):
    res = compare_lists_rec(list1, list2, delta + 0.000001)
    if res:
        if _assert:
            env.assertTrue(True, message=f'{str(list1)} ~ {str(list2)}')
        return True
    else:
        if _assert:
            env.assertTrue(False, message=f'{str(list1)} ~ {str(list2)}')
        return False

class ConditionalExpected:
    def __init__(self, env, cond):
        self.env = env
        self.cond_val = cond(env)
        self.query = None

    def call(self, *query):
        self.query = query
        return self

    def expect_when(self, cond_val, func: Callable[[Query], Any]):
        if cond_val == self.cond_val:
            func(self.env.expect(*self.query))
        return self

def load_vectors_to_redis(env, n_vec, query_vec_index, vec_size, data_type='FLOAT32', ids_offset=0, seed=10):
    conn = getConnectionByEnv(env)
    np.random.seed(seed)
    p = conn.pipeline(transaction=False)
    query_vec = None
    for i in range(n_vec):
        vector = create_np_array_typed(np.random.rand(vec_size), data_type)
        if i == query_vec_index:
            query_vec = vector
        p.execute_command('HSET', ids_offset + i, 'vector', vector.tobytes())
    p.execute()
    return query_vec

def sortResultByKeyName(res, start_index=1):
  '''
    Sorts the result by NAMEs
    res = [<COUNT>, '<NAME_1>, '<VALUE_1>', '<NAME_2>, '<VALUE_2>', ...]

    If VALUEs are lists, they are sorted by name as well
  '''
  # Sort name and value pairs by name
  pairs = [(name,sortResultByKeyName(value, 0) if isinstance(value, list) else value) for name,value in zip(res[start_index::2], res[start_index+1::2])]
  pairs = [i for i in sorted(pairs, key=lambda x: x[0])]
  # Flatten the sorted pairs to a list
  pairs = [i for pair in pairs for i in pair]
  if start_index == 1:
    # Bring the COUNT back to the beginning
    res = [res[0], *pairs]
  else:
    res = [*pairs]
  return res

def dict_diff(res, exp, show=False, ignore_order=True, significant_digits=7,
              ignore_numeric_type_changes=True, exclude_paths=None,
              exclude_regex_paths=None):
    dd = DeepDiff(res, exp, exclude_types={_ANY}, ignore_order=ignore_order, significant_digits=significant_digits,
                  ignore_numeric_type_changes=ignore_numeric_type_changes, exclude_paths=exclude_paths,
                  exclude_regex_paths=exclude_regex_paths)
    if dd != {} and show:
        pp(dd)
    return dd

def number_to_ordinal(n: int) -> str:
    if 11 <= (n % 100) <= 13:
        suffix = 'th'
    else:
        suffix = ['th', 'st', 'nd', 'rd', 'th'][min(n % 10, 4)]
    return str(n) + suffix

def populate_db(env: Env, idx_name: str = 'idx', text: bool = False, numeric: bool = False, tag: bool = False, n_per_shard=10000):
    """
    Creates a simple index called `idx`, and populates the database with
    `n * n_shards` matching documents.
    The names of the fields will be 'text1', 'numeric1', 'tag1' corresponding to
    the field type.

    Parameters:
    -----------
        env (Env): Environment to populate.
        idx_name: The name of the index to create.
        text (bool): Whether to create a text field in the index.
        numeric (bool): Whether to create a numeric field in the index.
        tag (bool): Whether to create a tag field in the index.
        n_per_shard (int): Number of documents to create per shard.

    Returns:
    -----------
        None
    """
    conn = getConnectionByEnv(env)
    text_f = 'text1 TEXT' if text else ''
    numeric_f = 'numeric1 NUMERIC' if numeric else ''
    tag_f = 'tag1 TAG' if tag else ''

    index_creation = f'FT.CREATE {idx_name} SCHEMA'
    if text:
        index_creation += f' {text_f}'
    if numeric:
        index_creation += f' {numeric_f}'
    if tag:
        index_creation += f' {tag_f}'

    conn.execute_command(*index_creation.split(' '))

    num_docs = n_per_shard * env.shardsCount
    pipeline = conn.pipeline(transaction=False)
    for i in range(num_docs):
        population_command = f'HMSET doc:{i}'
        if text:
            population_command += f' text1 lala:{i}'
        if numeric:
            population_command += f' numeric1 {i}'
        if tag:
            population_command += f' tag1 MOVIE'

        pipeline.execute_command(*population_command.split(' '))
        if i % 1000 == 0:
            pipeline.execute()
            pipeline = conn.pipeline(transaction=False)
    pipeline.execute()

def get_TLS_args():
    root = os.environ.get('ROOT', None)
    if root is None:
        root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # go up 3 levels from common.py

    cert_file       = os.path.join(root, 'bin', 'tls', 'redis.crt')
    key_file        = os.path.join(root, 'bin', 'tls', 'redis.key')
    ca_cert_file    = os.path.join(root, 'bin', 'tls', 'ca.crt')
    passphrase_file = os.path.join(root, 'bin', 'tls', '.passphrase')

    with_pass = server_version_is_at_least('6.2')

    # If any of the files are missing, generate them
    import subprocess
    subprocess.run([os.path.join(root, 'sbin', 'gen-test-certs'), str(1 if with_pass else 0)]).check_returncode()

    def get_passphrase():
        with open(passphrase_file, 'r') as f:
            return f.read()

    passphrase = get_passphrase() if with_pass else None

    return cert_file, key_file, ca_cert_file, passphrase

# Dispatch a command to make sure that the module is loaded and initialized
# We need to dispatch a command that will activate the topology updater, by
# sending a command to the shards. Otherwise the cluster.refresh command will
# not be effective, due to lazy initialization of the topology updater.
# Thus we dispatch a command that does not have an index, as it is not stopped
# in the coordinator level.
def verify_shard_init(shard):
    # One of the following errors can be raised (timing), yet they
    # mean the same thing in this case - the command was dispatched
    # to the shards before the connections were ready. Continue to
    # try until success\timeout.
    uninitialized_errors = [
        'ERRCLUSTER Uninitialized cluster state, could not perform command',
        'Could not distribute command'
    ]
    # The following error means that the cluster is initialized, as it was
    # returned from the shards.
    initialized_error = 'Alias does not exist'

    with TimeLimit(5, 'Failed to verify shard initialization'):
        while True:
            try:
                shard.execute_command('FT.ALIASDEL', 'non-existing-alias')
                break
            except redis_exceptions.ResponseError as e:
                if any([err in str(e) for err in uninitialized_errors]):
                    continue
                elif initialized_error in str(e):
                    break
                # Unexpected error, raise it.
                raise

def cmd_assert(env, cmd, res, message=None):
    db_res = env.cmd(*cmd)
    env.assertEqual(db_res, res, message=message)

# fields should be in capital letters
def getInvertedIndexInitialSize(env, fields, depth=0):
    total_size = 0
    for field in fields:
        if field in ['GEO', 'NUMERIC']:
            inverted_index_size = 24
            inverted_index_meta_data = 8
            total_size += inverted_index_size + inverted_index_meta_data
            continue
        env.assertTrue(field in ['TEXT', 'TAG', 'GEOMETRY', 'VECTOR'], message=f"type {field} is not supported", depth=depth+1)

    return total_size

# fields should be in capital letters
def getInvertedIndexInitialSize_MB(env, fields, depth=0) -> float:
    return getInvertedIndexInitialSize(env, fields, depth=depth+1) / float(1024 * 1024)

def check_index_info(env, idx, exp_num_records, exp_inv_idx_size, msg="", depth=0):
    d = index_info(env, idx)
    env.assertEqual(float(d['num_records']), exp_num_records, message=msg + ", num_records", depth=depth+1)

    if(exp_inv_idx_size != None):
        env.assertEqual(float(d['inverted_sz_mb']), exp_inv_idx_size, message=msg + ", inverted_sz_mb", depth=depth+1)

# Iterates items in d1 and compare their keys[value] with d2
# asserts when a key is missing in d2
# For simplicity, all values are compared as floats
def compare_numeric_dicts(env, d1, d2, d1_name="d1", d2_name="d2", msg="", _assert=True, depth=0):
    for key, value in d1.items():
        try:
            res = float(d2[key]) == float(value)
            if _assert:
                env.assertTrue(res, message=msg + " value is different in key: " + key + " expected " + str(value) + " got " + str(d2[key]), depth=depth+1)
            else:
                if res == False:
                    return False
        except KeyError:
            if _assert:
                env.assertTrue(False, message=msg + f" key {key} exists in {d1_name} but doesn't exist in {d2_name}")
            else:
                raise KeyError
    return True

def compare_index_info_dict(env, idx, expected_info_dict, msg="", depth=0):
    d = index_info(env, idx)
    compare_numeric_dicts(env, expected_info_dict, d, "expected_info_dict", "index_info", msg, depth=depth+1)

# expected info for index that was initialized and *emptied*
def check_index_info_empty(env, idx, fields, msg="after delete all and gc", depth=0):
    expected_size = getInvertedIndexInitialSize_MB(env, fields, depth=depth+1)
    check_index_info(env, idx, exp_num_records=0, exp_inv_idx_size=expected_size, msg=msg, depth=depth+1)

def recursive_index(lst, target):
    for i, element in enumerate(lst):
        if isinstance(element, list):
            sublist_index = recursive_index(element, target)
            if sublist_index != -1:
                return [i] + sublist_index
        elif element == target:
            return [i]
    return -1

def recursive_contains(lst, target):
    return recursive_index(lst, target) != -1

def access_nested_list(lst, index):
    result = lst
    for entry in index:
        result = result[entry]
    return result

def getRDBFile(env, file_name, depth=0):
    # Materialise a bundled RDB fixture from tests/pytests/test_rdbs/<file_name>.zip
    # into REDISEARCH_CACHE_DIR/<file_name>. Extraction is idempotent: if the
    # target file already exists with non-zero size we skip re-extracting.
    src = os.path.join(TEST_RDBS_DIR, file_name + '.zip')
    dst = os.path.join(REDISEARCH_CACHE_DIR, file_name)
    if os.path.exists(dst) and os.path.getsize(dst) > 0:
        return True
    if not os.path.exists(src):
        env.assertTrue(
            False,
            message=f"bundled RDB fixture {src} is missing",
            depth=depth + 1,
        )
        return False
    import zipfile
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    try:
        with zipfile.ZipFile(src, 'r') as z:
            z.extract(os.path.basename(file_name), os.path.dirname(dst))
    except (zipfile.BadZipFile, KeyError, OSError) as e:
        env.assertTrue(
            False,
            message=f"failed to extract bundled RDB fixture {src}: {e}",
            depth=depth + 1,
        )
        return False
    return True

def getRDBFiles(env, rdbs=None, depth=0):
    if rdbs is None:
        return False
    for f in rdbs:
        if not getRDBFile(env, f, depth=depth + 1):
            return False
    return True

def index_errors(env, idx = 'idx'):
    return to_dict(index_info(env, idx)['Index Errors'])
def field_errors(env, idx = 'idx', fld_index = 0):
    return to_dict(to_dict(to_dict(index_info(env, idx)['field statistics'][fld_index]))['Index Errors'])

def VerifyTimeoutWarningResp3(env, res, message="", depth=0):
    env.assertTrue(res['warning'], message=message + " expected warning", depth=depth+1)
    if (res['warning']):
        env.assertContains("Timeout", res["warning"][0], message=message + " expected timeout warning", depth=depth+1)

def parseDebugQueryCommandArgs(query_cmd, debug_params):
    return [*query_cmd, *debug_params, 'DEBUG_PARAMS_COUNT', len(debug_params)]

def runDebugQueryCommand(env, query_cmd, debug_params):
    # Use the helper function to build the argument list
    args = parseDebugQueryCommandArgs(query_cmd, debug_params)
    return env.cmd(debug_cmd(), *args)


def runDebugQueryCommandTimeoutAfterN(env, query_cmd, timeout_res_count, internal_only=False):
    debug_params = ['TIMEOUT_AFTER_N', timeout_res_count]
    if internal_only:
        debug_params.append("INTERNAL_ONLY")
    return runDebugQueryCommand(env, query_cmd, debug_params)


def runDebugQueryCommandAndCrash(env, query_cmd, crash_in_rust=False):
    debug_params = ["CRASH_IN_RUST" if crash_in_rust else "CRASH"]
    return env.expect(
        debug_cmd(), *query_cmd, *debug_params, "DEBUG_PARAMS_COUNT", len(debug_params)
    ).error()


def runDebugQueryCommandPauseAfterRPAfterN(env, query_cmd, rp_type, pause_after_n):
    debug_params = ['PAUSE_AFTER_RP_N', rp_type, pause_after_n]
    return runDebugQueryCommand(env, query_cmd, debug_params)

def runDebugQueryCommandPauseBeforeRPAfterN(env, query_cmd, rp_type, pause_after_n, extra_args=None):
    debug_params = ['PAUSE_BEFORE_RP_N', rp_type, pause_after_n]
    if extra_args:
        debug_params.extend(extra_args)
    return runDebugQueryCommand(env, query_cmd, debug_params)

def getIsRPPaused(env):
    return env.cmd(debug_cmd(), 'QUERY_CONTROLLER', 'GET_IS_RP_PAUSED')

def setPauseRPResume(env):
    return env.cmd(debug_cmd(), 'QUERY_CONTROLLER', 'SET_PAUSE_RP_RESUME')

def allShards_getIsRPPaused(env):
    results = []
    for shardId in range(1, env.shardsCount + 1):
        result = env.getConnection(shardId).execute_command(debug_cmd(), 'QUERY_CONTROLLER', 'GET_IS_RP_PAUSED')
        results.append(result)
    return results

def allShards_setPauseRPResume(env, start_shard=1):
    results = []
    for shardId in range(start_shard, env.shardsCount + 1):
        result = env.getConnection(shardId).execute_command(debug_cmd(), 'QUERY_CONTROLLER', 'SET_PAUSE_RP_RESUME')
        results.append(result)
    return results

# Coordinator Reduce Pause helpers (only available when built with ENABLE_ASSERT)

# Named constants for the N parameter of setPauseBeforeReduce
NO_PAUSE = 0                    # Disable pause (no pause point set)
PAUSE_AFTER_LAST_RESULT = -1    # Pause after the last result is reduced
PAUSE_BEFORE_REDUCER_INIT = -2  # Pause after claiming reducing but before reducer context init

def setPauseBeforeReduce(env, N):
    """
    Set the coordinator to pause before reducing the Nth result.
    PAUSE_BEFORE_REDUCER_INIT (-2): pause after claiming reducing but before reducer context init
    PAUSE_AFTER_LAST_RESULT (-1): pause after the last result is reduced
    NO_PAUSE (0): no pause
    N>0: pause before the Nth result (1-based index)
    """
    env.expect(debug_cmd(), 'QUERY_CONTROLLER', 'SET_PAUSE_BEFORE_REDUCE', N).ok()

def getIsCoordReducePaused(env):
    """Check if the coordinator is currently paused during reduce."""
    return env.cmd(debug_cmd(), 'QUERY_CONTROLLER', 'GET_IS_COORD_REDUCE_PAUSED')

def setCoordReduceResume(env):
    """Resume the coordinator from a reduce pause."""
    env.expect(debug_cmd(), 'QUERY_CONTROLLER', 'SET_COORD_REDUCE_RESUME').ok()

def getCoordReduceCount(env):
    """Get the current count of results reduced so far."""
    return env.cmd(debug_cmd(), 'QUERY_CONTROLLER', 'GET_COORD_REDUCE_COUNT')

def resetCoordReduceDebug(env):
    """Reset the coordinator reduce debug context (set N=0 and resume).

    Note: setCoordReduceResume will error if the coordinator is not currently paused,
    which is expected in cleanup scenarios where the coordinator already resumed.
    """
    setPauseBeforeReduce(env, NO_PAUSE)
    try:
        # Use env.cmd here since we need to catch the exception
        env.cmd(debug_cmd(), 'QUERY_CONTROLLER', 'SET_COORD_REDUCE_RESUME')
    except Exception:
        pass  # Ignore error if coordinator is not paused

# AggregateResults loop pause helpers (only available when built with ENABLE_ASSERT).
# These drive the AggregateResultsDebugCtx in src/debug_commands.{h,c} which the
# AggregateResults loop in aggregate_exec_common.c consults after each extracted
# result, busy-spinning until the test calls setAggregateResultsResume.
# `target` may be an Env (uses .expect().ok() / .cmd()) or a raw connection
# (uses .execute_command()), so per-shard tests can drive the same hooks on
# a specific shard's process.
def _qc_set(target, *args):
    cmd = (debug_cmd(), 'QUERY_CONTROLLER') + args
    if hasattr(target, 'expect'):
        target.expect(*cmd).ok()
    else:
        target.execute_command(*cmd)

def _qc_get(target, *args):
    cmd = (debug_cmd(), 'QUERY_CONTROLLER') + args
    if hasattr(target, 'cmd'):
        return target.cmd(*cmd)
    return target.execute_command(*cmd)

def setPauseAfterAggregateResult(target, N):
    """Pause the AggregateResults loop after the Nth result is extracted.

    N == 0 disables the pause; N > 0 pauses after the Nth result (1-based).
    Resets the internal results counter so successive tests start from zero.
    """
    _qc_set(target, 'SET_PAUSE_AFTER_AGGREGATE_RESULT', N)

def getIsAggregateResultsPaused(target):
    """Check if the AggregateResults loop is currently paused."""
    return _qc_get(target, 'GET_IS_AGGREGATE_RESULTS_PAUSED')

def setAggregateResultsResume(target):
    """Resume the AggregateResults loop from a pause."""
    _qc_set(target, 'SET_AGGREGATE_RESULTS_RESUME')

def getAggregateResultsCount(target):
    """Get the number of results extracted so far by the AggregateResults loop."""
    return _qc_get(target, 'GET_AGGREGATE_RESULTS_COUNT')

def resetAggregateResultsDebug(target):
    """Reset the AggregateResults debug context (clear pause point and resume).

    Mirrors resetCoordReduceDebug: tolerates the "not paused" state so cleanup
    is safe to call regardless of the loop's current state. The pause loop in
    debugCheckAndPauseAfterAggregateResult self-releases when AREQ_TimedOut is
    observed, so by the time tests reach cleanup the loop may already be
    unpaused -- in that case SET_AGGREGATE_RESULTS_RESUME would error, which
    must not be surfaced as a test failure.
    """
    setPauseAfterAggregateResult(target, 0)
    if getIsAggregateResultsPaused(target) == 1:
        _qc_set(target, 'SET_AGGREGATE_RESULTS_RESUME')

# Store Results Pause helpers (only available when built with ENABLE_ASSERT)
def setPauseBeforeStoreResults(env, enabled, internal):
    """Enable/disable pausing before AREQ_StoreResults/HREQ_StoreResults.

    internal: True restricts the pause to internal (coordinator-dispatched)
    requests; False restricts it to non-internal (user-facing) requests.
    """
    scope = 'INTERNAL_ONLY' if internal else 'NON_INTERNAL_ONLY'
    env.expect(debug_cmd(), 'QUERY_CONTROLLER', 'SET_PAUSE_BEFORE_STORE_RESULTS',
               'true' if enabled else 'false', scope).ok()

def setPauseAfterStoreResults(env, enabled, internal):
    """Enable/disable pausing after AREQ_StoreResults/HREQ_StoreResults.

    internal: True restricts the pause to internal (coordinator-dispatched)
    requests; False restricts it to non-internal (user-facing) requests.
    """
    scope = 'INTERNAL_ONLY' if internal else 'NON_INTERNAL_ONLY'
    env.expect(debug_cmd(), 'QUERY_CONTROLLER', 'SET_PAUSE_AFTER_STORE_RESULTS',
               'true' if enabled else 'false', scope).ok()

def getIsStoreResultsPaused(env):
    """Check if the query is currently paused during store results."""
    return env.cmd(debug_cmd(), 'QUERY_CONTROLLER', 'GET_IS_STORE_RESULTS_PAUSED')

def resetStoreResultsDebug(env):
    """Reset the store results debug context (disable pauses and resume).

    Disabling does not depend on the scope (the C-level enable flags gate the
    pause), so the raw FT.DEBUG commands are sent here without a scope token.
    """
    env.expect(debug_cmd(), 'QUERY_CONTROLLER', 'SET_PAUSE_BEFORE_STORE_RESULTS', 'false').ok()
    env.expect(debug_cmd(), 'QUERY_CONTROLLER', 'SET_PAUSE_AFTER_STORE_RESULTS', 'false').ok()
    try:
        env.cmd(debug_cmd(), 'QUERY_CONTROLLER', 'SET_STORE_RESULTS_RESUME')
    except Exception:
        pass  # Ignore error if not paused

# Hybrid Store Cursors Pause helpers (only available when built with ENABLE_ASSERT)
# These are separate from Store Results and only affect cursor storage in HybridRequest_StartCursors
def setPauseBeforeHybridStoreCursors(env, enabled):
    """Enable/disable pausing before hybrid cursor storage (HybridRequest_StartCursors)."""
    env.expect(debug_cmd(), 'QUERY_CONTROLLER', 'SET_PAUSE_BEFORE_HYBRID_STORE_CURSORS', 'true' if enabled else 'false').ok()

def setPauseAfterHybridStoreCursors(env, enabled):
    """Enable/disable pausing after hybrid cursor storage (HybridRequest_StartCursors)."""
    env.expect(debug_cmd(), 'QUERY_CONTROLLER', 'SET_PAUSE_AFTER_HYBRID_STORE_CURSORS', 'true' if enabled else 'false').ok()

def getIsHybridStoreCursorsPaused(env):
    """Check if the hybrid is currently paused during cursor storage."""
    return env.cmd(debug_cmd(), 'QUERY_CONTROLLER', 'GET_IS_HYBRID_STORE_CURSORS_PAUSED')

def resetHybridStoreCursorsDebug(env):
    """Reset the hybrid store cursors debug context (disable pauses and resume)."""
    setPauseBeforeHybridStoreCursors(env, False)
    setPauseAfterHybridStoreCursors(env, False)
    try:
        env.cmd(debug_cmd(), 'QUERY_CONTROLLER', 'SET_HYBRID_STORE_CURSORS_RESUME')
    except Exception:
        pass  # Ignore error if not paused

def isEnableAssertEnabled(env):
    """
    Check if ENABLE_ASSERT is enabled in the build.
    Returns True if ENABLE_ASSERT commands are available, False otherwise.
    """
    try:
        env.cmd(debug_cmd(), 'QUERY_CONTROLLER', 'GET_IS_COORD_REDUCE_PAUSED')
        return True
    except Exception:
        return False

def skipIfNoEnableAssert(env):
    """
    Skip the current test if ENABLE_ASSERT is not enabled in the build.
    Call this at the beginning of tests that require ENABLE_ASSERT functionality.
    """
    if not isEnableAssertEnabled(env):
        env.debugPrint("Skipping test: ENABLE_ASSERT is not enabled", force=True)
        env.skip()

def require_enable_assert(f):
    """
    Decorator to skip tests if ENABLE_ASSERT is not enabled in the build.
    Usage: @require_enable_assert
    """
    @wraps(f)
    def wrapper(env, *args, **kwargs):
        if not isEnableAssertEnabled(env):
            env.debugPrint(f"Skipping {f.__name__}: ENABLE_ASSERT is not enabled", force=True)
            env.skip()
            return
        return f(env, *args, **kwargs)
    return wrapper

class vecsimMockTimeoutContext:
    """Context manager for enabling/disabling VECSIM mock timeout on all shards"""
    def __init__(self, env):
        self.env = env

    def __enter__(self):
        run_command_on_all_shards(self.env, debug_cmd(), 'VECSIM_MOCK_TIMEOUT', 'enable')
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        run_command_on_all_shards(self.env, debug_cmd(), 'VECSIM_MOCK_TIMEOUT', 'disable')

def shardsConnections(env):
  for s in range(1, env.shardsCount + 1):
      yield env.getConnection(shardId=s)

def _normalize_cluster_shards(reply):
    """Coerce a CLUSTER SHARDS reply (RESP2 nested arrays or RESP3 maps) into a
    list of {'slots': [s1,e1,...], 'nodes': [{'id','role',...}]} dicts."""
    shards = []
    for sh in reply:
        d = sh if isinstance(sh, dict) else {sh[i]: sh[i + 1] for i in range(0, len(sh), 2)}
        nodes = [n if isinstance(n, dict) else {n[i]: n[i + 1] for i in range(0, len(n), 2)}
                 for n in d['nodes']]
        shards.append({'slots': list(d['slots']), 'nodes': nodes})
    return shards

def distinct_shard_tags(conn):
    """Yield hash tags that each land on a different shard, verified live via
    CLUSTER SHARDS + CLUSTER KEYSLOT. One tag per shard, deterministic order."""
    shards = _normalize_cluster_shards(conn.execute_command('CLUSTER', 'SHARDS'))
    def owner(slot):
        for sh in shards:
            r = sh['slots']
            if any(a <= slot <= b for a, b in zip(r[::2], r[1::2])):
                return next(n['id'] for n in sh['nodes'] if n['role'] in ('master', 'primary'))
        raise AssertionError(f'no shard owns slot {slot}')
    seen = set()
    for n in itertools.count():
        tag = str(n)
        sid = owner(int(conn.execute_command('CLUSTER', 'KEYSLOT', tag)))
        if sid not in seen:
            seen.add(sid)
            yield tag
            if len(seen) == len(shards):
                return

def waitForIndexFinishScan(env, idx = 'idx'):
    # Wait for the index to finish scan
    # Check if equals 1 for RESP3 support
    with TimeLimit(60, 'Timeout while waiting for index to finish scan'):
        while index_info(env, idx)['percent_indexed'] not in (1, '1'):
            time.sleep(0.1)

def bgScanCommand():
    return debug_cmd() + ' BG_SCAN_CONTROLLER'

def getDebugScannerStatus(env, idx = 'idx'):
    return env.cmd(bgScanCommand(), 'GET_DEBUG_SCANNER_STATUS', idx)

def checkDebugScannerStatusError(env, idx = 'idx', expected_error = ''):
    env.expect(bgScanCommand(), 'GET_DEBUG_SCANNER_STATUS', idx).error() \
        .contains(expected_error)

def checkDebugScannerUpdateError(env, idx = 'idx', expected_error = ''):
    env.expect(bgScanCommand(), 'DEBUG_SCANNER_UPDATE_CONFIG', idx).error() \
        .contains(expected_error)

def set_tight_maxmemory_for_oom(env, used_memory_ratio = 1.001):
    # Get current memory consumption value
    memory_usage = env.cmd('INFO', 'MEMORY')['used_memory']
    # Set memory limit based on the current memory usage, according to the used ratio
    required_limit = memory_usage / used_memory_ratio

    env.expect('config', 'set', 'maxmemory', int(required_limit)).ok()

def set_unlimited_maxmemory_for_oom(env):
    env.expect('config', 'set', 'maxmemory', 0).ok()


def waitForIndexStatus(env, status, idx='idx'):
    with TimeLimit(60, 'Timeout while waiting for index status'):
        while getDebugScannerStatus(env, idx) != status:
            time.sleep(0.1)

def waitForIndexPauseScan(env,idx = 'idx'):
    waitForIndexStatus(env,'PAUSED', idx)

def shard_getDebugScannerStatus(env, shardId, idx = 'idx'):
    return env.getConnection(shardId).execute_command(bgScanCommand(), 'GET_DEBUG_SCANNER_STATUS', idx)

def shard_waitForIndexStatus(env, shardId, status, idx='idx'):
    with TimeLimit(60, 'Timeout while waiting for index status'):
        while shard_getDebugScannerStatus(env, shardId, idx) != status:
            time.sleep(0.1)

def shard_waitForIndexPauseScan(env, shardId, idx = 'idx'):
    shard_waitForIndexStatus(env, shardId, 'PAUSED', idx)

def allShards_waitForIndexPauseScan(env, idx = 'idx'):
    for shardId in range(1, env.shardsCount + 1):
        shard_waitForIndexPauseScan(env, shardId, idx)

def allShards_waitForIndexStatus(env, status, idx='idx'):
    for shardId in range(1, env.shardsCount + 1):
        shard_waitForIndexStatus(env, shardId, status, idx)

def shard_waitForIndexFinishScan(env, shardId, idx = 'idx'):
    # Wait for the index to finish scan
    # Check if equals 1 for RESP3 support
    with TimeLimit(60, 'Timeout while waiting for index to finish scan'):
        while index_info(env, idx)['percent_indexed'] not in (1, '1'):
            time.sleep(0.1)

def allShards_waitForIndexFinishScan(env, idx = 'idx'):
    for shardId in range(1, env.shardsCount + 1):
        shard_waitForIndexFinishScan(env, shardId, idx)

def shard_set_tight_maxmemory_for_oom(env, shardId, memory_limit_per = 1.0):
    # Get current memory consumption value
    memory_usage = env.getConnection(shardId).execute_command('INFO', 'MEMORY')['used_memory']
    # Set memory limit to less then memory limit
    required_memory = memory_usage * (1/memory_limit_per)
    # Round up and add 1
    new_memory = math.ceil(required_memory) + 1
    res = env.getConnection(shardId).execute_command('config', 'set', 'maxmemory', new_memory)
    env.assertEqual(res, 'OK')

def allShards_set_tight_maxmemory_for_oom(env, memory_limit_per = 1.0):
    for shardId in range(1, env.shardsCount + 1):
        shard_set_tight_maxmemory_for_oom(env, shardId, memory_limit_per)

def shard_set_unlimited_maxmemory_for_oom(env, shardId):
    res = env.getConnection(shardId).execute_command('config', 'set', 'maxmemory', 0)
    env.assertEqual(res, 'OK')

def allShards_set_unlimited_maxmemory_for_oom(env):
    for shardId in range(1, env.shardsCount + 1):
        shard_set_unlimited_maxmemory_for_oom(env, shardId)

def assertEqual_dicts_on_intersection(env, d1, d2, message=None, depth=0):
    for k in d1:
        if k in d2:
            env.assertEqual(d1[k], d2[k], message=message, depth=depth+1)

def get_results_from_hybrid_response(response) -> Dict[str, Dict[str, any]]:
    """Extract all fields from hybrid response results

    Args:
        response: Hybrid search response containing results

    Returns:
        Dict mapping key -> dict of all fields from the results list
        Example: {'doc:1': {'__score': '0.5', 'vector_distance': '0.3'}}
    """
    # Handle RESP3 format (dict)
    if isinstance(response, dict):
        results = {}
        for result in response.get('results', []):
            if '__key' in result:
                key = result['__key']
                results[key] = result
        total_results = response.get('total_results', 0)
        return results, total_results

    res_results_index = recursive_index(response, 'results')
    res_count_index = recursive_index(response, 'total_results')
    res_results_index[-1] += 1
    res_count_index[-1] += 1

    results = {}
    for result in access_nested_list(response, res_results_index):
        # Each result has structure: ['attributes', [flat_key_value_list]]
        result = dict(zip(result[::2], result[1::2]))
        if '__key' in result:
            key = result['__key']
            results[key] = result
    total_results = access_nested_list(response, res_count_index)
    return results, total_results

def populate_db_with_faker_text(env, num_docs, doc_len=5, seed=12345, offset=0):
    """Populate database with faker-generated text documents

    Args:
        env: Test environment
        num_docs: Number of documents to create
        doc_len: Number of words per document (equivalent to dim parameter)
        seed: Random seed for reproducibility
        offset: Starting offset for document IDs (equivalent to ids_offset)
    """
    conn = getConnectionByEnv(env)
    fake = faker.Faker()
    fake.seed_instance(seed)

    # Use pipeline for better performance
    pipeline = conn.pipeline(transaction=False)
    for i in range(num_docs):
        # Generate sentences with specified number of words
        text = fake.sentence(nb_words=doc_len, variable_nb_words=False).rstrip('.')
        pipeline.execute_command('HSET', f'{offset + i}', 'description', text)

        # Execute pipeline every 1000 docs to avoid memory issues
        if i % 1000 == 0:
            pipeline.execute()
            pipeline = conn.pipeline(transaction=False)

    # Execute remaining docs
    pipeline.execute()


def call_and_store(fn, args, out_list):
    """
    Helper function for threading: calls a function and stores its return value in a list.

    Args:
        fn: Function to call
        args: Tuple of arguments to pass to the function
        out_list: List to append the function's return value to
    """
    out_list.append(fn(*args))

def launch_cmds_in_bg_with_exception_check(env, command, num_triggers, exception_timeout=1):
    """
    Launch the same Redis command multiple times in background threads with exception monitoring.

    Args:
        env: Redis test environment for executing commands.
        command: A list containing the Redis command to execute (e.g., ['FT.SEARCH', 'idx', 'query']).
        num_triggers: Number of background threads to spawn, each executing the same command.
        exception_timeout: Seconds to wait for exception detection (default: 1).

    Returns:
        tuple[list[Thread] | None, list[Exception]]: (threads, exceptions). `threads` is the
        list of started thread objects, or None if any thread raised within `exception_timeout`.
        `exceptions` is the live list that background threads append to on failure; callers may
        re-inspect it after a later wait to surface errors that arrive past the fast-fail window.
    """
    threads = []
    exceptions = []
    exception_event = threading.Event()

    def run_cmd():
        try:
            env.cmd(*command)
        except Exception as e:
            exceptions.append(e)
            exception_event.set()

    for i in range(num_triggers):
        t = threading.Thread(target=run_cmd)
        threads.append(t)
        t.start()

    # Check for exceptions before proceeding
    if exception_event.wait(timeout=exception_timeout):
        error_msg = f"Background command {command} failed with {len(exceptions)} error(s): {exceptions}"
        env.assertTrue(False, message=error_msg)
        return None, exceptions

    return threads, exceptions

def generate_slots(slots = range(2**14)) -> bytes:
    """Generate slot ranges in binary format matching RedisModuleSlotRangeArray serialization.

    Args:
        slots: Iterable of slot numbers (default: 0-16383)

    Returns:
        bytes: Binary format with:
            - First 4 bytes: int32 number of ranges (little-endian)
            - Following bytes: pairs of uint16 (start, end) for each range (little-endian)
    """
    slots = set(slots)
    ranges_list = []

    for slot in range(2**14):
        if slot in slots:
            if ranges_list and slot == ranges_list[-1][1] + 1:
                ranges_list[-1][1] = slot
            else:
                ranges_list.append([slot, slot])

    # Convert list to numpy array of uint16 pairs
    ranges_array = np.array(ranges_list, dtype=np.uint16)

    # Create the output: 4 bytes for count (int32) + flattened uint16 pairs
    num_ranges = np.int32(len(ranges_list))

    # Use sys.byteorder to handle endianness properly, but force little-endian
    count_bytes = num_ranges.tobytes() if sys.byteorder == 'little' else num_ranges.byteswap().tobytes()
    ranges_bytes = ranges_array.tobytes() if sys.byteorder == 'little' else ranges_array.byteswap().tobytes()

    return count_bytes + ranges_bytes

def change_oom_policy(env, policy):
    env.expect(config_cmd(), 'SET', 'ON_OOM', policy).ok()

def shard_change_oom_policy(env, shardId, policy):
    res = env.getConnection(shardId).execute_command(config_cmd(), 'SET', 'ON_OOM', policy)
    env.assertEqual(res, 'OK')

def allShards_change_oom_policy(env, policy):
    for shardId in range(1, env.shardsCount + 1):
        shard_change_oom_policy(env, shardId, policy)

def allShards_change_maxmemory_low(env):
    for shardId in range(1, env.shardsCount + 1):
        res = env.getConnection(shardId).execute_command('config', 'set', 'maxmemory', 1)
        env.assertEqual(res, 'OK')

def shard_change_timeout_policy(env, shardId, policy):
    res = env.getConnection(shardId).execute_command(config_cmd(), 'SET', 'ON_TIMEOUT', policy)
    env.assertEqual(res, 'OK')

def allShards_change_timeout_policy(env, policy):
    for shardId in range(1, env.shardsCount + 1):
        shard_change_timeout_policy(env, shardId, policy)

def get_shards_profile(env, res):
  """Extract shard profiles from FT.PROFILE AGGREGATE response."""
  if env.protocol == 3:
    return res['Profile']['Shards']
  else:
    return [to_dict(p) for p in res[-1][1]]


# --- Coordinator / cluster timeout test helpers (shared across timeout test files) ---
ON_TIMEOUT_CONFIG = 'search-on-timeout'


def pid_cmd(conn):
    """Get the process ID of a Redis connection."""
    return conn.execute_command('info', 'server')['process_id']


def non_coord_shard_conns(env):
    """Return shard connections whose process id differs from the coordinator's."""
    coord_pid = pid_cmd(env.con)
    conns = []
    for shardId in range(1, env.shardsCount + 1):
        conn = env.getConnection(shardId)
        if pid_cmd(conn) != coord_pid:
            conns.append(conn)
    return conns


def split_shards_pick_one_paused(env):
    """Pick one non-coordinator shard to designate as paused and split the rest.

    Returns ``(all_shard_conns, paused_conn, paused_pid, responsive_conns)``.
    Asserts that at least one non-coordinator shard exists.
    """
    all_shard_conns = [env.getConnection(i) for i in range(1, env.shardsCount + 1)]
    non_coord_conns = non_coord_shard_conns(env)
    env.assertGreater(len(non_coord_conns), 0,
                      message="Test requires at least one non-coordinator shard")
    paused_conn = non_coord_conns[0]
    paused_pid = pid_cmd(paused_conn)
    responsive_conns = [c for c in all_shard_conns if pid_cmd(c) != paused_pid]
    return all_shard_conns, paused_conn, paused_pid, responsive_conns


def assert_timeout_warning(env, res, message=''):
    warnings = res.get('warning', res.get('warnings', []))
    env.assertTrue(warnings, message=message + " expected timeout warning")
    env.assertContains('Timeout', warnings[0], message=message + " expected timeout warning")