import pathlib

from .utils import *
from .assertions import assert_http_ok

N_PEERS = 4
COLLECTION_NAME = "test_collection"


def test_peer_snapshot_bootstrap(tmp_path: pathlib.Path):
    assert_project_root()

    peer_dirs = make_peer_folders(tmp_path, N_PEERS)

    # Gathers REST API uris
    peer_api_uris = []

    # Get stable port for first peer
    #
    # NOTE:
    #
    # This peer will be killed and restarted later in the test.
    #
    # If the port changes and the peer have to report URI change,
    # *4 out of 5 times it will fail to do so*, and the test will fail.
    first_peer_port = get_port()

    # Start bootstrap
    (bootstrap_api_uri, bootstrap_uri) = start_first_peer(
        peer_dirs[0], "peer_0_0.log", port=first_peer_port)

    peer_api_uris.append(bootstrap_api_uri)

    # Wait for leader
    leader = wait_peer_added(bootstrap_api_uri)

    # Start other peers
    for i in range(1, len(peer_dirs) - 1):
        peer_api_uris.append(start_peer(peer_dirs[i], f"peer_0_{i}.log", bootstrap_uri))

    # Wait for cluster
    wait_for_uniform_cluster_status(peer_api_uris, leader)

    # Check that there are no collections on all peers
    for uri in peer_api_uris:
        r = requests.get(f"{uri}/collections")
        assert_http_ok(r)
        assert len(r.json()["result"]["collections"]) == 0

    # Create collection on first peer
    r = requests.put(f"{peer_api_uris[0]}/collections/{COLLECTION_NAME}", json={
        "vectors": {
            "size": 4,
            "distance": "Dot"
        }
    })
    assert_http_ok(r)

    # Check that it exists on all peers
    wait_collection_exists_and_active_on_all_peers(
        collection_name=COLLECTION_NAME,
        peer_api_uris=peer_api_uris,
    )

    # Recover all follower peers
    for peer_api_uri in peer_api_uris[1:]:
        r = requests.post(f"{peer_api_uri}/cluster/recover")
        assert_http_ok(r)

    # Kill leader
    processes.pop(0).kill()

    # Wait for new leader
    #
    # NOTE: `leader_is_reelected(3)` is not an error! It returns a stateful closure.
    wait_for(leader_is_reelected(3), peer_api_uris[1], leader)

    # Restart killed peer
    (bootstrap_api_uri, bootstrap_uri) = start_first_peer(
        peer_dirs[0], "peer_0_0_restarted.log", port=first_peer_port)

    peer_api_uris[0] = bootstrap_api_uri

    # Wait for restarted peer
    leader = wait_peer_added(bootstrap_api_uri, expected_size=N_PEERS - 1)

    # Recover restarted peer
    r = requests.post(f"{bootstrap_api_uri}/cluster/recover")
    assert_http_ok(r)

    # Start new peer
    peer_api_uris.append(start_peer(
        peer_dirs[-1], f"peer_0_{len(peer_dirs) - 1}.log", bootstrap_uri))

    # Wait for new peer
    wait_peer_added(bootstrap_api_uri, expected_size=N_PEERS)
    wait_for_peer_online(peer_api_uris[-1])

    # Check that collection exists on new peer
    wait_collection_exists_and_active_on_all_peers(
        collection_name=COLLECTION_NAME,
        peer_api_uris=peer_api_uris,
    )


def leader_is_reelected(count):
    current_leader = None
    current_leader_count = 0

    def leader_is_reelected(uri, previous_leader):
        r = requests.get(f"{uri}/cluster")
        assert_http_ok(r)

        leader = r.json()["result"]["raft_info"]["leader"]

        if leader != previous_leader and not leader is None:
            nonlocal current_leader
            nonlocal current_leader_count

            if leader == current_leader:
                current_leader_count += 1
            else:
                current_leader = leader
                current_leader_count = 1

            if current_leader_count >= count:
                return True

        return False

    return leader_is_reelected
