import pathlib

import requests

from .assertions import assert_http_ok
from .fixtures import create_collection, upsert_points
from .utils import *

N_PEERS = 3
COLLECTION_NAME = "test_collection"
# Keep in sync with `api::HTTP_HEADER_ROUTING_TOKEN`.
ROUTING_HEADER = "X-Qdrant-Route-Affinity"
POINT_ID = 1
VECTOR = [0.1, 0.2, 0.3, 0.4]


def upsert_marker(peer_uri, marker):
    """Upsert the single test point with the given `marker` payload on all replicas."""
    r = upsert_points(
        peer_uri,
        [{"id": POINT_ID, "vector": VECTOR, "payload": {"marker": marker}}],
        wait="true",
        ordering="strong",
    )
    assert_http_ok(r)


def retrieve_marker(peer_uri, headers={}):
    """Retrieve the test point and return its `marker`, routing with the given headers."""
    r = requests.post(
        f"{peer_uri}/collections/{COLLECTION_NAME}/points",
        json={"ids": [POINT_ID], "with_payload": True},
        headers=headers,
        timeout=10,
    )
    assert_http_ok(r)
    points = r.json()["result"]
    return points[0]["payload"]["marker"] if points else None


def retrieve_marker_on_all_peers(peer_api_uris, headers={}):
    """Retrieve the marker through *every* peer and assert they all agree.

    Routing is deterministic: a token is ordered by `hash(token, peer_id)`, which is
    identical on every node, so every peer must resolve the token to the same replica
    no matter which peer received the request. Returns that single agreed-upon marker.
    """
    by_peer = {uri: retrieve_marker(uri, headers=headers) for uri in peer_api_uris}
    distinct = set(by_peer.values())
    assert len(distinct) == 1, f"peers disagreed on the routed replica: {by_peer}"
    return distinct.pop()


def create_snapshot(peer_uri):
    r = requests.post(
        f"{peer_uri}/collections/{COLLECTION_NAME}/snapshots?wait=true",
        timeout=30,
    )
    assert_http_ok(r)
    return r.json()["result"]["name"]


def recover_snapshot_no_sync(peer_uri, snapshot_url):
    # `no_sync` keeps the recovered data authoritative on this peer *without* removing
    # or re-syncing the other replica, leaving the two replicas in a divergent (but
    # both Active) state — exactly the situation the routing token addresses.
    r = requests.put(
        f"{peer_uri}/collections/{COLLECTION_NAME}/snapshots/recover",
        json={"location": snapshot_url, "priority": "no_sync"},
        timeout=60,
    )
    assert_http_ok(r)


def split_by_local_replica(peer_api_uris):
    """Split peers into those that hold a local replica of shard 0 and those that don't."""
    hosts, non_hosts = [], []
    for uri in peer_api_uris:
        r = requests.get(f"{uri}/collections/{COLLECTION_NAME}/cluster", timeout=10)
        assert_http_ok(r)
        local_shards = r.json()["result"]["local_shards"]
        if any(shard["shard_id"] == 0 for shard in local_shards):
            hosts.append(uri)
        else:
            non_hosts.append(uri)
    return hosts, non_hosts


def replicas_diverged(peer_a, peer_b):
    """True once peer_a serves the recovered "A" and peer_b still serves "B".

    Read locally on each replica host (no token, so the local replica is preferred),
    confirming the two replicas hold the divergent state the test relies on.
    """
    return retrieve_marker(peer_a) == "A" and retrieve_marker(peer_b) == "B"


def test_routing_token_sticky_reads(tmp_path: pathlib.Path):
    """
    A routing token pins a read to a deterministic replica, and the same token lands
    on the same replica regardless of which peer received the request.

    We build two replicas of one shard with *divergent* data: snapshot state "A",
    move both replicas to state "B", then recover the "A" snapshot onto a single
    replica with `no_sync` (so the two replicas stay Active but divergent). We then
    validate, issuing reads through *every* peer (not just the coordinator):

      1. without a routing token, a peer holding no local replica routes reads
         randomly, so repeated reads vary between the replicas (the baseline);
      2. with a routing token, every peer resolves it to the same replica, so reads
         are stable and identical across all peers;
      3. different routing tokens can land on different replicas, yet each individual
         token resolves to the same replica on every peer.
    """
    assert_project_root()

    peer_api_uris, _peer_dirs, _bootstrap_uri = start_cluster(tmp_path, N_PEERS)

    # One shard with two replicas (placed on two of the three peers).
    # - `write_consistency_factor=2` so every write reaches *both* replicas before we diverge them.
    # - `read_fan_out_factor=0` so each read hits exactly one replica (no fan-out),
    #   which is what makes the routing decision observable.
    create_collection(
        peer_api_uris[0],
        shard_number=1,
        replication_factor=2,
        write_consistency_factor=2,
        sparse_vectors=False,
        read_fan_out_factor=0,
    )
    wait_collection_exists_and_active_on_all_peers(
        collection_name=COLLECTION_NAME, peer_api_uris=peer_api_uris
    )

    # Both replicas start in state "A".
    upsert_marker(peer_api_uris[0], "A")

    hosts, non_hosts = split_by_local_replica(peer_api_uris)
    assert len(hosts) == 2, f"expected 2 replica hosts, got {hosts}"
    assert len(non_hosts) == 1, f"expected 1 peer without a local replica, got {non_hosts}"
    peer_a, peer_b = hosts
    # The coordinator holds no local replica, so reads issued here cannot be answered
    # locally and must be routed to one of the two remote replicas.
    coordinator = non_hosts[0]

    # Snapshot state "A" on peer_a, then move *both* replicas to state "B".
    snapshot_name = create_snapshot(peer_a)
    upsert_marker(peer_api_uris[0], "B")

    # Recover the "A" snapshot onto peer_a only, without resync:
    #   peer_a -> "A", peer_b -> "B"  (both Active, divergent)
    snapshot_url = f"{peer_a}/collections/{COLLECTION_NAME}/snapshots/{snapshot_name}"
    recover_snapshot_no_sync(peer_a, snapshot_url)
    # Reads are routed through every peer below, so every peer must see both replicas
    # as Active before we start (otherwise a not-yet-readable replica would be skipped
    # and break the deterministic routing we are asserting).
    for uri in peer_api_uris:
        wait_for_all_replicas_active(uri, COLLECTION_NAME)

    # Precondition: the replicas are genuinely divergent (local reads differ). Poll
    # rather than asserting immediately — `wait_for_all_replicas_active` only confirms
    # the replica *metadata* is Active, but there is a brief window afterwards where
    # peer_a still serves its pre-recovery "B" before the recovered "A" becomes
    # readable on the read path.
    wait_for(replicas_diverged, peer_a, peer_b)

    n_reads = 40

    # 1. No routing token: the coordinator routes each read to a random replica, so
    #    both are observed. (The non-determinism is only visible on the coordinator;
    #    the two replica-holding peers always answer such reads from their local copy.)
    without_token = {retrieve_marker(coordinator) for _ in range(n_reads)}
    assert without_token == {"A", "B"}, (
        f"expected varying results without a routing token, observed {without_token}"
    )

    # 2. With a routing token: the read is pinned to one replica, and *every* peer
    #    resolves the token to the same replica (`retrieve_marker_on_all_peers` asserts
    #    the peers agree). Repeated rounds must therefore yield a single stable marker.
    sticky_headers = {ROUTING_HEADER: "user-sticky"}
    with_token = {
        retrieve_marker_on_all_peers(peer_api_uris, headers=sticky_headers)
        for _ in range(n_reads)
    }
    assert len(with_token) == 1, (
        f"expected a stable result across all peers with a routing token, observed {with_token}"
    )

    # 3. Different routing tokens can land on different replicas, yet each individual
    #    token resolves to the same replica on every peer (asserted per token inside
    #    `retrieve_marker_on_all_peers`). Each token maps deterministically to a replica
    #    via `hash(token, peer_id)`; since peer ids are random per cluster, use enough
    #    tokens to make an all-same-replica split vanishingly unlikely (~2 * 0.5**32).
    per_token = {
        retrieve_marker_on_all_peers(peer_api_uris, headers={ROUTING_HEADER: f"user-{i}"})
        for i in range(32)
    }
    assert per_token == {"A", "B"}, (
        f"expected different routing tokens to reach different replicas, observed {per_token}"
    )
