import json
import os
import re
import shutil
import signal
from subprocess import Popen
import time
from typing import Tuple, Callable, Dict, List, Optional
import requests
import socket
from contextlib import closing
from pathlib import Path
import pytest
from .assertions import assert_http_ok


WAIT_TIME_SEC = 30
RETRY_INTERVAL_SEC = 0.2
PROJECT_ROOT = Path(__file__).parent.parent.parent

# Tracks processes that need to be killed at the end of the test
processes: List['PeerProcess'] = []
busy_ports = {}


class PeerProcess:
    def __init__(self, proc: Popen, http_port, grpc_port, p2p_port):
        self.proc = proc
        self.http_port = http_port
        self.grpc_port = grpc_port
        self.p2p_port = p2p_port
        self.pid = proc.pid

    def kill(self):
        self.proc.kill()
        self.proc.wait()

        # remove allocated ports from the dictionary
        # so they can be used afterwards
        del busy_ports[self.http_port]
        del busy_ports[self.grpc_port]
        del busy_ports[self.p2p_port]

    def interrupt(self):
        self.proc.send_signal(signal.SIGINT)
        self.proc.wait()

        del busy_ports[self.http_port]
        del busy_ports[self.grpc_port]
        del busy_ports[self.p2p_port]


def _occupy_port(port):
    if port in busy_ports:
        raise Exception(f'Port "{port}" was already allocated!')
    busy_ports[port] = True
    return port


def kill_all_processes():
    print()
    while len(processes) > 0:
        p = processes.pop(0)
        if is_coverage_mode():
            print(f"Interrupting {p.pid}")
            p.interrupt()
        else:
            print(f"Killing {p.pid}")
            p.kill()


@pytest.fixture(autouse=True)
def every_test():
    yield
    kill_all_processes()


def get_port() -> int:
    while True:
        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
            # get random port assigned by the OS
            s.bind(('', 0))
            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            allocated_port = s.getsockname()[1]
            if allocated_port in busy_ports:
                continue
            return allocated_port

def is_coverage_mode() -> bool:
    return os.getenv("COVERAGE") == "1"


def get_env(p2p_port: int, grpc_port: int, http_port: int) -> Dict[str, str]:
    env = os.environ.copy()
    env["QDRANT__CLUSTER__ENABLED"] = "true"
    env["QDRANT__CLUSTER__P2P__PORT"] = str(p2p_port)
    env["QDRANT__SERVICE__HTTP_PORT"] = str(http_port)
    env["QDRANT__SERVICE__GRPC_PORT"] = str(grpc_port)
    env["QDRANT__LOG_LEVEL"] = "TRACE,raft::raft=info,actix_http=info,tonic=info,want=info,mio=info"
    env["QDRANT__SERVICE__HARDWARE_REPORTING"] = "true"

    if is_coverage_mode():
        env["LLVM_PROFILE_FILE"] = get_llvm_profile_file()

    return env


def get_uri(port: int) -> str:
    return f"http://127.0.0.1:{port}"


def assert_project_root():
    """Deprecated: No longer needed as paths are resolved relative to __file__."""
    pass


def get_qdrant_exec() -> str:
    if is_coverage_mode():
        qdrant_exec = PROJECT_ROOT / "target" / "llvm-cov-target" / "debug" / "qdrant"
    else:
        qdrant_exec = PROJECT_ROOT / "target" / "debug" / "qdrant"
    return str(qdrant_exec)

def get_llvm_profile_file() -> str:
    # %m: keep merging results from each test into the same file
    # If you have multiple tests running in parallel, you can use -%p OR -%{thread_count}m to have different files
    # Not using -%p since each test will generate a new file
    llvm_profile_file = PROJECT_ROOT / "target" / "llvm-cov-target" / "qdrant-consensus-tests-%m.profraw"
    return str(llvm_profile_file)


def get_pytest_current_test_name() -> str:
    # https://docs.pytest.org/en/latest/example/simple.html#pytest-current-test-environment-variable
    return os.environ.get('PYTEST_CURRENT_TEST').split(':')[-1].split(' ')[0]


def init_pytest_log_folder() -> str:
    test_name = get_pytest_current_test_name()
    log_folder = f"consensus_test_logs/{test_name}"
    if not os.path.exists(log_folder):
        os.makedirs(log_folder)
    return log_folder


# Starts a peer and returns its api_uri
def start_peer(peer_dir: Path, log_file: str, bootstrap_uri: str, port=None, extra_env=None, reinit=False, uris_in_env=False) -> str:
    if extra_env is None:
        extra_env = {}
    p2p_port = get_port() if port is None else port + 0
    _occupy_port(p2p_port)
    grpc_port = get_port() if port is None else port + 1
    _occupy_port(grpc_port)
    http_port = get_port() if port is None else port + 2
    _occupy_port(http_port)

    test_log_folder = init_pytest_log_folder()
    log_file = open(f"{test_log_folder}/{log_file}", "w")
    this_peer_consensus_uri = get_uri(p2p_port)
    print(f"Starting follower peer with bootstrap uri {bootstrap_uri},"
          f" http: http://localhost:{http_port}/cluster, p2p: {p2p_port}")

    args = [get_qdrant_exec()]
    env = {
        **get_env(p2p_port, grpc_port, http_port),
        **extra_env
    }

    if uris_in_env:
        env["QDRANT_BOOTSTRAP"] = bootstrap_uri
        env["QDRANT_URI"] = this_peer_consensus_uri
    else:
        args.extend(["--bootstrap", bootstrap_uri, "--uri", this_peer_consensus_uri])

    if reinit:
        args.append("--reinit")

    # Wrap with systemd-run to throttle CPU to investigate issues
    # wrapped_cmd = ["systemd-run", "--user", "--scope", "-p", "CPUQuota=20%", "--"] + args
    # proc = Popen(wrapped_cmd, env=env, cwd=peer_dir, stdout=log_file)
    proc = Popen(args, env=env, cwd=peer_dir, stdout=log_file)
    processes.append(PeerProcess(proc, http_port, grpc_port, p2p_port))
    return get_uri(http_port)


# Starts a peer and returns its api_uri and p2p_uri
def start_first_peer(peer_dir: Path, log_file: str, port=None, extra_env=None, reinit=False, uris_in_env=False) -> Tuple[str, str]:
    if extra_env is None:
        extra_env = {}

    p2p_port = get_port() if port is None else port + 0
    _occupy_port(p2p_port)
    grpc_port = get_port() if port is None else port + 1
    _occupy_port(grpc_port)
    http_port = get_port() if port is None else port + 2
    _occupy_port(http_port)

    test_log_folder = init_pytest_log_folder()
    log_file = open(f"{test_log_folder}/{log_file}", "w")
    bootstrap_uri = get_uri(p2p_port)
    print(f"\nStarting first peer with uri {bootstrap_uri},"
          f" http: http://localhost:{http_port}/cluster, p2p: {p2p_port}")

    args = [get_qdrant_exec()]
    env = {
        **get_env(p2p_port, grpc_port, http_port),
        **extra_env
    }

    if uris_in_env:
        env["QDRANT_URI"] = bootstrap_uri
    else:
        args.extend(["--uri", bootstrap_uri])

    if reinit:
        args.append("--reinit")

    # Wrap with systemd-run to throttle CPU to investigate issues
    # wrapped_cmd = ["systemd-run", "--user", "--scope", "-p", "CPUQuota=20%", "--"] + args
    # proc = Popen(wrapped_cmd, env=env, cwd=peer_dir, stdout=log_file)
    proc = Popen(args, env=env, cwd=peer_dir, stdout=log_file)
    processes.append(PeerProcess(proc, http_port, grpc_port, p2p_port))
    return get_uri(http_port), bootstrap_uri


def start_cluster(tmp_path, num_peers, port_seed=None, extra_env=None, headers={}, uris_in_env=False, log_file_prefix=""):
    assert_project_root()
    peer_dirs = make_peer_folders(tmp_path, num_peers)

    # Gathers REST API uris
    peer_api_uris = []

    # Start bootstrap
    (bootstrap_api_uri, bootstrap_uri) = start_first_peer(peer_dirs[0], f"{log_file_prefix}peer_0_0.log", port=port_seed,
                                                          extra_env=extra_env, uris_in_env=uris_in_env)
    peer_api_uris.append(bootstrap_api_uri)

    # Wait for leader
    leader = wait_peer_added(bootstrap_api_uri, headers=headers)

    port = None
    # Start other peers
    for i in range(1, len(peer_dirs)):
        if port_seed is not None:
            port = port_seed + i * 100
        peer_api_uris.append(start_peer(peer_dirs[i], f"{log_file_prefix}peer_0_{i}.log", bootstrap_uri, port=port, extra_env=extra_env, uris_in_env=uris_in_env))

    # Wait for cluster
    wait_for_uniform_cluster_status(peer_api_uris, leader, headers=headers)

    return peer_api_uris, peer_dirs, bootstrap_uri


def make_peer_folder(base_path: Path, peer_number: int) -> Path:
    peer_dir = base_path / f"peer{peer_number}"
    peer_dir.mkdir()
    shutil.copytree(PROJECT_ROOT / "config", peer_dir / "config")
    return peer_dir


def make_peer_folders(base_path: Path, n_peers: int) -> List[Path]:
    peer_dirs = []
    for i in range(n_peers):
        peer_dir = make_peer_folder(base_path, i)
        peer_dirs.append(peer_dir)
    return peer_dirs


def get_cluster_info(peer_api_uri: str, headers={}) -> dict:
    r = requests.get(f"{peer_api_uri}/cluster", headers=headers)
    assert_http_ok(r)
    res = r.json()["result"]
    return res


def print_clusters_info(peer_api_uris: [str], headers={}):
    for uri in peer_api_uris:
        try:
            # do not crash if the peer is not online
            print(json.dumps(get_cluster_info(uri, headers=headers), indent=4))
        except requests.exceptions.ConnectionError:
            print(f"Can't retrieve cluster info for offline peer {uri}")


def fetch_highest_peer_id(peer_api_uris: [str]) -> str:
    max_peer_id = 0
    max_peer_url = None
    for uri in peer_api_uris:
        try:
            # do not crash if the peer is not online
            peer_id = get_cluster_info(uri)['peer_id']
            if peer_id > max_peer_id:
                max_peer_id = peer_id
                max_peer_url = uri
        except requests.exceptions.ConnectionError:
            print(f"Can't retrieve cluster info for offline peer {uri}")
    return max_peer_url


def get_collection_cluster_info(peer_api_uri: str, collection_name: str, headers={}) -> dict:
    r = requests.get(f"{peer_api_uri}/collections/{collection_name}/cluster", headers=headers)
    assert_http_ok(r)
    res = r.json()["result"]
    return res


def get_shard_transfer_count(peer_api_uri: str, collection_name: str) -> int:
    info = get_collection_cluster_info(peer_api_uri, collection_name)
    return len(info["shard_transfers"])


def get_collection_info(peer_api_uri: str, collection_name: str) -> dict:
    r = requests.get(f"{peer_api_uri}/collections/{collection_name}")
    assert_http_ok(r)
    res = r.json()["result"]
    return res


def get_collection_point_count(peer_api_uri: str, collection_name: str, exact: bool = False,
                               shard_key: Optional[str] = None, filter: Optional[dict] = None) -> int:
    r = requests.post(f"{peer_api_uri}/collections/{collection_name}/points/count", json={"exact": exact, "shard_key": shard_key, "filter": filter})
    assert_http_ok(r)
    res = r.json()["result"]["count"]
    return res


def print_collection_cluster_info(peer_api_uri: str, collection_name: str, headers={}):
    print(json.dumps(get_collection_cluster_info(peer_api_uri, collection_name, headers=headers), indent=4))


def get_leader(peer_api_uri: str, headers={}) -> str:
    r = requests.get(f"{peer_api_uri}/cluster", headers=headers)
    assert_http_ok(r)
    return r.json()["result"]["raft_info"]["leader"]


def check_leader(peer_api_uri: str, expected_leader: str, headers={}) -> bool:
    try:
        r = requests.get(f"{peer_api_uri}/cluster", headers=headers)
        assert_http_ok(r)
        leader = r.json()["result"]["raft_info"]["leader"]
        correct_leader = leader == expected_leader
        if not correct_leader:
            print(f"Cluster leader invalid for peer {peer_api_uri} {leader}/{expected_leader}")
        return correct_leader
    except requests.exceptions.ConnectionError:
        # the api is not yet available - caller needs to retry
        print(f"Could not contact peer {peer_api_uri} to fetch cluster leader")
        return False


def leader_is_defined(peer_api_uri: str, headers={}) -> bool:
    try:
        r = requests.get(f"{peer_api_uri}/cluster", headers=headers)
        assert_http_ok(r)
        leader = r.json()["result"]["raft_info"]["leader"]
        return leader is not None
    except requests.exceptions.ConnectionError:
        # the api is not yet available - caller needs to retry
        print(f"Could not contact peer {peer_api_uri} to fetch leader info")
        return False


def check_cluster_size(peer_api_uri: str, expected_size: int, headers={}) -> bool:
    try:
        r = requests.get(f"{peer_api_uri}/cluster", headers=headers)
        assert_http_ok(r)
        peers = r.json()["result"]["peers"]
        correct_size = len(peers) == expected_size
        if not correct_size:
            print(f"Cluster size invalid for peer {peer_api_uri} {len(peers)}/{expected_size}")
        return correct_size
    except requests.exceptions.ConnectionError:
        # the api is not yet available - caller needs to retry
        print(f"Could not contact peer {peer_api_uri} to fetch cluster size")
        return False


def all_nodes_cluster_info_consistent(peer_api_uris: [str], expected_leader: str, headers={}) -> bool:
    expected_size = len(peer_api_uris)
    for uri in peer_api_uris:
        if check_leader(uri, expected_leader, headers=headers) and check_cluster_size(uri, expected_size, headers=headers):
            continue
        else:
            return False
    return True

def peers_have_version(peer_api_uris: [str]) -> bool:
    for peer_api_uri in peer_api_uris:
        try:
            # Check versions in local telemetry of each peer
            # Not using cluster level telemetry because it shows versions before
            # they are fully propagated through consensus
            r = requests.get(f"{peer_api_uri}/telemetry?details_level=3")
            assert_http_ok(r)
            cluster = r.json()["result"]["cluster"]
            peers = cluster["peers"]
            peer_metadata = cluster["peer_metadata"]

            for peer_id in peers.keys():
                # Peers without metadata are not listed
                if peer_id not in peer_metadata:
                    return False
                if peer_metadata[peer_id].get("version") is None:
                    return False
        except requests.exceptions.ConnectionError:
            print(f"Could not contact peer {peer_api_uri} to fetch versions")
            return False
    return True

def all_nodes_have_same_commit(peer_api_uris: [str]) -> bool:
    commits = []
    for uri in peer_api_uris:
        try:
            r = requests.get(f"{uri}/cluster")
            assert_http_ok(r)
            commits.append(r.json()["result"]["raft_info"]["commit"])
        except requests.exceptions.ConnectionError:
            print(f"Could not contact peer {uri} to fetch commit")
            return False
    return len(set(commits)) == 1


def all_nodes_respond(peer_api_uris: [str]) -> bool:
    for uri in peer_api_uris:
        try:
            r = requests.get(f"{uri}/collections")
            assert_http_ok(r)
        except requests.exceptions.ConnectionError:
            print(f"Could not contact peer {uri} to fetch collections")
            return False
    return True


def collection_exists_on_all_peers(collection_name: str, peer_api_uris: [str]) -> bool:
    for uri in peer_api_uris:
        r = requests.get(f"{uri}/collections")
        assert_http_ok(r)
        collections = r.json()["result"]["collections"]
        filtered_collections = [c for c in collections if c['name'] == collection_name]
        if len(filtered_collections) == 0:
            print(
                f"Collection '{collection_name}' does not exist on peer {uri} found {json.dumps(collections, indent=4)}")
            return False
        else:
            continue
    return True


def check_collection_local_shards_count(peer_api_uri: str, collection_name: str,
                                        expected_local_shard_count: int) -> bool:
    return get_collection_local_shards_count(peer_api_uri, collection_name) == expected_local_shard_count


def get_collection_local_shards_count(peer_api_uri: str, collection_name: str) -> int:
    collection_cluster_info = get_collection_cluster_info(peer_api_uri, collection_name)
    return len(collection_cluster_info["local_shards"])


def check_collection_local_shards_point_count(peer_api_uri: str, collection_name: str,
                                            expected_count: int) -> int:
    collection_cluster_info = get_collection_cluster_info(peer_api_uri, collection_name)
    point_count = sum(map(lambda shard: shard["points_count"], collection_cluster_info["local_shards"]))

    is_correct = point_count == expected_count
    if not is_correct:
        print(f"Collection '{collection_name}' on peer {peer_api_uri} ({point_count} != {expected_count}): {json.dumps(collection_cluster_info, indent=4)}")

    return is_correct


def check_collection_shard_transfers_count(peer_api_uri: str, collection_name: str,
                                           expected_shard_transfers_count: int, headers={}) -> bool:
    collection_cluster_info = get_collection_cluster_info(peer_api_uri, collection_name, headers=headers)
    local_shard_count = len(collection_cluster_info["shard_transfers"])
    return local_shard_count == expected_shard_transfers_count


def check_collection_resharding_operations_count(peer_api_uri: str, collection_name: str,
                                           expected_resharding_operations_count: int, headers={}) -> bool:
    collection_cluster_info = get_collection_cluster_info(peer_api_uri, collection_name, headers=headers)

    # TODO(resharding): until resharding release, the resharding operations are not always exposed
    # Once we do release, we can remove the zero fallback here
    # See: <https://github.com/qdrant/qdrant/pull/4599>
    local_resharding_count = len(collection_cluster_info["resharding_operations"]) if "resharding_operations" in collection_cluster_info else 0
    return local_resharding_count == expected_resharding_operations_count


def check_collection_resharding_operation_stage(peer_api_uri: str, collection_name: str, expected_stage: str, headers={}) -> bool:
    collection_cluster_info = get_collection_cluster_info(peer_api_uri, collection_name, headers=headers)
    if "resharding_operations" not in collection_cluster_info:
        return False
    for resharding in collection_cluster_info["resharding_operations"]:
        if "comment" in resharding and resharding["comment"].startswith(expected_stage):
            return True
    return False


def check_collection_shard_transfer_method(peer_api_uri: str, collection_name: str,
                                           expected_method: str) -> bool:
    collection_cluster_info = get_collection_cluster_info(peer_api_uri, collection_name)

    # Shortcut if no transfers
    if len(collection_cluster_info["shard_transfers"]) == 0:
        print(f"check_collection_shard_transfer_method: no transfers for collection '{collection_name}'")
        return False

    # Check method on each transfer
    for transfer in collection_cluster_info["shard_transfers"]:
        if "method" not in transfer:
            print(f"check_collection_shard_transfer_method: unknown method not match expected '{expected_method}'")
            continue
        method = transfer["method"]
        if method == expected_method:
            return True
        else:
            print(f"check_collection_shard_transfer_method: method '{method}' does not match expected '{expected_method}'")

    return False


def check_collection_shard_transfer_progress(peer_api_uri: str, collection_name: str,
                                             expected_transfer_progress: int,
                                             expected_transfer_total: int) -> bool:
    collection_cluster_info = get_collection_cluster_info(peer_api_uri, collection_name)

    # Check progress on each transfer
    for transfer in collection_cluster_info["shard_transfers"]:
        if "comment" not in transfer:
            continue
        comment = transfer["comment"]

        # Compare progress or total
        m = re.search(r"Transferring records \((\d+)/(\d+)\)", comment)
        if m is None:
            continue
        current, total = m.groups()
        if current is not None and expected_transfer_progress is not None and int(
                current) >= expected_transfer_progress:
            return True
        if total is not None and expected_transfer_total is not None and int(total) >= expected_transfer_total:
            return True

    return False


def check_all_replicas_active(peer_api_uri: str, collection_name: str, headers={}) -> bool:
    try:
        collection_cluster_info = get_collection_cluster_info(peer_api_uri, collection_name, headers=headers)
        for shard in collection_cluster_info["local_shards"]:
            if shard['state'] != 'Active':
                return False
        for shard in collection_cluster_info["remote_shards"]:
            if shard['state'] != 'Active':
                return False
    except requests.exceptions.ConnectionError:
        return False
    return True


def check_some_replicas_not_active(peer_api_uri: str, collection_name: str) -> bool:
    return not check_all_replicas_active(peer_api_uri, collection_name)


def check_collection_cluster(peer_url, collection_name):
    res = requests.get(f"{peer_url}/collections/{collection_name}/cluster", timeout=10)
    assert_http_ok(res)
    return res.json()["result"]['local_shards'][0]


def check_strict_mode_enabled(peer_api_uri: str, collection_name: str) -> bool:
    collection_info = get_collection_info(peer_api_uri, collection_name)
    if "strict_mode_config" not in collection_info["config"]:
        return False
    strict_mode_enabled = collection_info["config"]["strict_mode_config"]["enabled"]
    return strict_mode_enabled == True

def check_strict_mode_disabled(peer_api_uri: str, collection_name: str) -> bool:
    collection_info = get_collection_info(peer_api_uri, collection_name)
    if "strict_mode_config" not in collection_info["config"]:
        return True
    strict_mode_enabled = collection_info["config"]["strict_mode_config"]["enabled"]
    return strict_mode_enabled == False

def wait_peer_added(peer_api_uri: str, expected_size: int = 1, headers={}) -> str:
    wait_for(check_cluster_size, peer_api_uri, expected_size, headers=headers)
    wait_for(leader_is_defined, peer_api_uri, headers=headers)
    return get_leader(peer_api_uri, headers=headers)

def wait_for_collection(peer_api_uri: str, collection_name: str):
    def is_collection_listed() -> bool:
        try:
            res = requests.get(f"{peer_api_uri}/collections")
            if not res.ok:
                return False
            collections = set(collection['name'] for collection in res.json()["result"]['collections'])
            return collection_name in collections
        except requests.exceptions.ConnectionError:
            return False

    wait_for(is_collection_listed)

def wait_collection_green(peer_api_uri: str, collection_name: str):
    try:
        wait_for(check_collection_green, peer_api_uri, collection_name)
    except Exception as e:
        print_clusters_info([peer_api_uri])
        raise e


def wait_for_some_replicas_not_active(peer_api_uri: str, collection_name: str):
    try:
        wait_for(check_some_replicas_not_active, peer_api_uri, collection_name)
    except Exception as e:
        print_clusters_info([peer_api_uri])
        raise e


def wait_for_all_replicas_active(peer_api_uri: str, collection_name: str, headers={}):
    try:
        wait_for(check_all_replicas_active, peer_api_uri, collection_name, headers=headers)
    except Exception as e:
        print_collection_cluster_info(peer_api_uri, collection_name, headers=headers)
        raise e


def wait_for_uniform_cluster_status(peer_api_uris: [str], expected_leader: str, headers={}):
    try:
        wait_for(all_nodes_cluster_info_consistent, peer_api_uris, expected_leader, headers=headers)
    except Exception as e:
        print_clusters_info(peer_api_uris)
        raise e

def wait_for_all_peers_versions(peer_api_uris: [str]):
    try:
        wait_for(peers_have_version, peer_api_uris)
    except Exception as e:
        print_clusters_info(peer_api_uris)
        raise e

def wait_for_same_commit(peer_api_uris: [str]):
    try:
        wait_for(all_nodes_have_same_commit, peer_api_uris)
    except Exception as e:
        print_clusters_info(peer_api_uris)
        raise e


def wait_all_peers_up(peer_api_uris: [str]):
    try:
        wait_for(all_nodes_respond, peer_api_uris)
    except Exception as e:
        print_clusters_info(peer_api_uris)
        raise e


def wait_for_uniform_collection_existence(collection_name: str, peer_api_uris: [str]):
    try:
        wait_for(collection_exists_on_all_peers, collection_name, peer_api_uris)
    except Exception as e:
        print_clusters_info(peer_api_uris)
        raise e


def wait_for_collection_shard_transfers_count(peer_api_uri: str, collection_name: str,
                                              expected_shard_transfer_count: int, headers={}):
    try:
        wait_for(check_collection_shard_transfers_count, peer_api_uri, collection_name, expected_shard_transfer_count, headers=headers)
    except Exception as e:
        print_collection_cluster_info(peer_api_uri, collection_name, headers=headers)
        raise e


def wait_for_collection_shard_transfer_method(peer_api_uri: str, collection_name: str,
                                              expected_method: str):
    try:
        wait_for(check_collection_shard_transfer_method, peer_api_uri, collection_name, expected_method)
    except Exception as e:
        print_collection_cluster_info(peer_api_uri, collection_name)
        raise e


def wait_for_collection_shard_transfer_progress(peer_api_uri: str, collection_name: str,
                                                expected_transfer_progress: int = None,
                                                expected_transfer_total: int = None):
    try:
        wait_for(check_collection_shard_transfer_progress, peer_api_uri, collection_name, expected_transfer_progress,
                 expected_transfer_total, wait_for_interval=0.1)
    except Exception as e:
        print_collection_cluster_info(peer_api_uri, collection_name)
        raise e


def wait_for_collection_resharding_operations_count(peer_api_uri: str,
                                                    collection_name: str,
                                                    expected_resharding_operations_count:
                                                    int, headers={}, **kwargs):
    try:
        wait_for(check_collection_resharding_operations_count, peer_api_uri, collection_name, expected_resharding_operations_count, headers=headers, **kwargs)
    except Exception as e:
        print_collection_cluster_info(peer_api_uri, collection_name, headers=headers)
        raise e


def wait_for_collection_resharding_operation_stage(peer_api_uri: str, collection_name: str, expected_stage: str, headers={}):
    try:
        wait_for(check_collection_resharding_operation_stage, peer_api_uri, collection_name, expected_stage, headers=headers)
    except Exception as e:
        print_collection_cluster_info(peer_api_uri, collection_name, headers=headers)
        raise e


def wait_for_collection_local_shards_count(peer_api_uri: str, collection_name: str, expected_local_shard_count: int):
    try:
        wait_for(check_collection_local_shards_count, peer_api_uri, collection_name, expected_local_shard_count)
    except Exception as e:
        print_collection_cluster_info(peer_api_uri, collection_name)
        raise e


def wait_for_strict_mode_enabled(peer_api_uri: str, collection_name: str):
    try:
        wait_for(check_strict_mode_enabled, peer_api_uri, collection_name)
    except Exception as e:
        print_collection_cluster_info(peer_api_uri, collection_name)
        raise e

def wait_for_strict_mode_disabled(peer_api_uri: str, collection_name: str):
    try:
        wait_for(check_strict_mode_disabled, peer_api_uri, collection_name)
    except Exception as e:
        print_collection_cluster_info(peer_api_uri, collection_name)
        raise e


def wait_for(condition: Callable[..., bool], *args, wait_for_timeout=WAIT_TIME_SEC, wait_for_interval=RETRY_INTERVAL_SEC, **kwargs):
    start = time.time()
    while not condition(*args, **kwargs):
        elapsed = time.time() - start
        if elapsed > wait_for_timeout:
            raise Exception(
                f"Timeout waiting for condition {condition.__name__} to be satisfied in {wait_for_timeout} seconds")
        else:
            time.sleep(wait_for_interval)

def peer_is_online(peer_api_uri: str, path: str = "/readyz") -> bool:
    try:
        r = requests.get(f"{peer_api_uri}{path}")
        return r.status_code == 200
    except:
        return False


def wait_for_peer_online(peer_api_uri: str, path="/readyz"):
    try:
        wait_for(peer_is_online, peer_api_uri, path=path)
    except Exception as e:
        print_clusters_info([peer_api_uri])
        raise e


def check_collection_green(peer_api_uri: str, collection_name: str, expected_status: str = "green") -> bool:
    collection_cluster_info = get_collection_info(peer_api_uri, collection_name)
    return collection_cluster_info['status'] == expected_status


def check_collection_points_count(peer_api_uri: str, collection_name: str, expected_size: int) -> bool:
    collection_cluster_info = get_collection_info(peer_api_uri, collection_name)
    return collection_cluster_info['points_count'] == expected_size


def wait_collection_points_count(peer_api_uri: str, collection_name: str, expected_size: int):
    try:
        wait_for(check_collection_points_count, peer_api_uri, collection_name, expected_size)
    except Exception as e:
        print_collection_cluster_info(peer_api_uri, collection_name)
        raise e


def wait_collection_on_all_peers(collection_name: str, peer_api_uris: [str], max_wait=30, headers={}):
    # Check that it exists on all peers
    while True:
        exists = True
        for url in peer_api_uris:
            r = requests.get(f"{url}/collections", headers=headers)
            assert_http_ok(r)
            collections = r.json()["result"]["collections"]
            exists &= any(collection["name"] == collection_name for collection in collections)
        if exists:
            break
        else:
            # Wait until collection is created on all peers
            # Consensus guarantees that collection will appear on majority of peers, but not on all of them
            # So we need to wait a bit extra time
            time.sleep(1)
            max_wait -= 1
        if max_wait <= 0:
            raise Exception("Collection was not created on all peers in time")


def wait_collection_exists_and_active_on_all_peers(collection_name: str, peer_api_uris: [str], max_wait=30, headers={}):
    wait_collection_on_all_peers(collection_name, peer_api_uris, max_wait, headers=headers)
    for peer_uri in peer_api_uris:
        # Collection is active on all peers
        wait_for_all_replicas_active(collection_name=collection_name, peer_api_uri=peer_uri, headers=headers)


def create_shard_key(
    shard_key,
    peer_url,
    collection="test_collection",
    shard_number=None,
    replication_factor=None,
    placement=None,
    timeout=10,
    headers={},
):
    r_create = requests.put(
        f"{peer_url}/collections/{collection}/shards?timeout={timeout}",
        json={
            "shard_key": shard_key,
            "shards_number": shard_number,
            "replication_factor": replication_factor,
            "placement": placement,
        },
        headers=headers,
    )
    assert_http_ok(r_create)


def move_shard(source_uri, collection_name, shard_id, source_peer_id, target_peer_id):
    r = requests.post(
        f"{source_uri}/collections/{collection_name}/cluster", json={
            "move_shard": {
                "shard_id": shard_id,
                "from_peer_id": source_peer_id,
                "to_peer_id": target_peer_id
            }
        })
    assert_http_ok(r)

def replicate_shard(source_uri, collection_name, shard_id, source_peer_id, target_peer_id, method=None):
    payload = {
        "shard_id": shard_id,
        "from_peer_id": source_peer_id,
        "to_peer_id": target_peer_id
    }
    if method:
        payload["method"] = method
    r = requests.post(
        f"{source_uri}/collections/{collection_name}/cluster", json={
            "replicate_shard": payload
        })
    assert_http_ok(r)


def check_data_consistency(data):

    assert(len(data) > 1)

    for i in range(len(data) - 1):
        j = i + 1

        data_i = data[i]
        data_j = data[j]

        if data_i != data_j:
            ids_i = set(x.get("id") for x in data_i)
            ids_j = set(x.get("id") for x in data_j)

            diff = ids_i - ids_j

            if len(diff) < 100:
                print(f"Diff between {i} and {j}: {diff}")
            else:
                sample = list(diff)[:32]
                print(f"Diff len between {i} and {j}: {len(diff)}, sample: {sample}")

            assert False, "Data on all nodes should be consistent"


def check_feature_enabled(peer_api_uri, feature) -> bool:
    r = requests.get(f"{peer_api_uri}/telemetry", params={"details_level": 10})
    assert_http_ok(r)
    features = r.json()['result']['app']['features']
    result = features[feature]
    return result


def skip_if_no_feature(peer_api_uri, feature):
    feature_is_enabled = check_feature_enabled(peer_api_uri, feature)
    if not feature_is_enabled:
        pytest.skip(f"Skipping because the feature {feature} is disabled at runtime.")
