# The MIT License (MIT)
# Copyright (c) Microsoft Corporation. All rights reserved.
"""Async multi-partition feed_range tests."""

import asyncio
import unittest
import uuid
from typing import Iterable, List, Optional, Tuple

import pytest
import pytest_asyncio

import test_config
from azure.cosmos import _base, documents
from azure.cosmos import http_constants
from azure.cosmos.aio import CosmosClient
from azure.cosmos._routing.feed_range_continuation import _decode_token
from azure.cosmos.partition_key import PartitionKey

CONFIG = test_config.TestConfig()
DATABASE_ID = CONFIG.TEST_DATABASE_ID

REPRO_CONTAINER_ID = "FeedRangeMultiPartitionAsync-" + str(uuid.uuid4())
REPRO_PARTITION_KEY = "pk"
REPRO_THROUGHPUT = CONFIG.THROUGHPUT_FOR_5_PARTITIONS
REPRO_DOC_COUNT = 200
PAGE_SIZE = 5
MIN_DOCS_PER_PARTITION = 15


def _client() -> CosmosClient:
    """Return a data-plane client. Uses AAD when COSMOS_TEST_DATA_AUTH_MODE=aad,
    key auth otherwise, so the AAD lane actually exercises AAD."""
    return test_config.TestConfig.create_data_client_async()


def _get_container(client: CosmosClient):
    return client.get_database_client(DATABASE_ID).get_container_client(REPRO_CONTAINER_ID)


async def _sorted_partition_ranges(container) -> List[Tuple[str, str]]:
    pairs: List[Tuple[str, str]] = []
    async for fr in container.read_feed_ranges():
        r = fr["Range"]
        pairs.append((r["min"], r["max"]))
    pairs.sort(key=lambda p: p[0])
    return pairs


async def _count_in_range(container, range_min: str, range_max: str) -> int:
    fr = test_config.create_feed_range_in_dict(
        test_config.create_range(range_min=range_min, range_max=range_max,
                                 is_min_inclusive=True, is_max_inclusive=False))
    items = [it async for it in container.query_items(
        query="SELECT VALUE COUNT(1) FROM c", feed_range=fr)]
    return items[0] if items else 0


def _crossing_feed_range(range_min: str, range_max: str):
    return test_config.create_feed_range_in_dict(
        test_config.create_range(range_min=range_min, range_max=range_max,
                                 is_min_inclusive=True, is_max_inclusive=False))


async def _ids_via_per_partition_scan(container, partition_ranges: Iterable[Tuple[str, str]]):
    ground_truth = set()
    for (mn, mx) in partition_ranges:
        fr = _crossing_feed_range(mn, mx)
        async for item in container.query_items(query="SELECT c.id FROM c", feed_range=fr):
            ground_truth.add(item["id"])
    return ground_truth


async def _values_via_per_partition_scan(container, partition_ranges: Iterable[Tuple[str, str]]):
    values = []
    for (mn, mx) in partition_ranges:
        fr = _crossing_feed_range(mn, mx)
        async for value in container.query_items(query='SELECT VALUE c["value"] FROM c', feed_range=fr):
            values.append(value)
    return values


async def _drain_pages(pager) -> Tuple[List[List[dict]], List[str]]:
    pages: List[List[dict]] = []
    all_ids: List[str] = []
    async for page in pager:
        items = [it async for it in page]
        pages.append(items)
        all_ids.extend(it["id"] for it in items)
    return pages, all_ids


@pytest_asyncio.fixture(scope="class", autouse=True)
async def setup_and_teardown_async():
    client = _client()
    db = client.get_database_client(DATABASE_ID)
    container = await db.create_container_if_not_exists(
        id=REPRO_CONTAINER_ID,
        partition_key=PartitionKey(path="/" + REPRO_PARTITION_KEY, kind="Hash"),
        offer_throughput=REPRO_THROUGHPUT)
    for i in range(REPRO_DOC_COUNT):
        await container.upsert_item({
            REPRO_PARTITION_KEY: f"pk-{i:04d}",
            "id": f"doc-{i:04d}",
            "value": i,
        })
    yield
    try:
        await db.delete_container(REPRO_CONTAINER_ID)
    except Exception:  # pylint: disable=broad-except
        pass
    await client.close()


@pytest.mark.cosmosQuery
@pytest.mark.cosmosAADQuery
@pytest.mark.asyncio
@pytest.mark.usefixtures("setup_and_teardown_async")
class TestFeedRangeMultiPartitionAsync:
    """Async end-to-end tests for feed_range queries that overlap multiple
    physical partitions."""

    async def test_single_partition_feed_range_async(self):
        """Single-partition regression guard."""
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if not partitions:
                pytest.skip("Container has no physical partitions")

            chosen_pp = None
            for (mn, mx) in partitions:
                if await _count_in_range(container, mn, mx) >= MIN_DOCS_PER_PARTITION:
                    chosen_pp = (mn, mx)
                    break
            if chosen_pp is None:
                pytest.skip("No single partition populated with ≥ "
                            f"{MIN_DOCS_PER_PARTITION} docs")
            single = _crossing_feed_range(chosen_pp[0], chosen_pp[1])
            ground_truth = await _ids_via_per_partition_scan(container, [chosen_pp])

            pager = container.query_items(
                query="SELECT * FROM c",
                feed_range=single,
                max_item_count=PAGE_SIZE,
            ).by_page()
            pages, all_ids = await _drain_pages(pager)

            for idx, page in enumerate(pages):
                if idx < len(pages) - 1:
                    assert len(page) == PAGE_SIZE, (
                        f"page {idx} returned {len(page)} items, expected "
                        f"exactly {PAGE_SIZE}")
                else:
                    assert len(page) <= PAGE_SIZE

            unique = set(all_ids)
            assert len(all_ids) == len(unique), (
                f"single-partition path returned duplicates: "
                f"{len(all_ids) - len(unique)} duplicate id(s)")
            assert unique == ground_truth, (
                f"single-partition coverage mismatch: returned={len(unique)}, "
                f"ground_truth={len(ground_truth)}, "
                f"missing={len(ground_truth - unique)}, "
                f"unexpected={len(unique - ground_truth)}")

            assert pager.continuation_token in (None, "", b""), (
                f"expected empty continuation after draining all pages; got "
                f"{pager.continuation_token!r}")
        finally:
            await client.close()


    async def test_two_partition_feed_range_async(self):
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 2:
                pytest.skip("Need a container with ≥ 2 physical partitions")

            chosen = None
            for i in range(len(partitions) - 1):
                p0, p1 = partitions[i], partitions[i + 1]
                if (await _count_in_range(container, p0[0], p0[1]) >= MIN_DOCS_PER_PARTITION
                        and await _count_in_range(container, p1[0], p1[1]) >= MIN_DOCS_PER_PARTITION):
                    chosen = (p0, p1)
                    break
            if chosen is None:
                pytest.skip("No adjacent partition pair both populated with ≥ "
                            f"{MIN_DOCS_PER_PARTITION} docs")
            (p0_min, _), (_, p1_max) = chosen
            crossing = _crossing_feed_range(p0_min, p1_max)
            ground_truth = await _ids_via_per_partition_scan(
                container, [chosen[0], chosen[1]])

            pager = container.query_items(
                query="SELECT * FROM c",
                feed_range=crossing,
                max_item_count=PAGE_SIZE,
            ).by_page()
            pages, all_ids = await _drain_pages(pager)

            oversized = [(i, len(p)) for i, p in enumerate(pages) if len(p) > PAGE_SIZE]
            assert not oversized, (
                f"page-size limit violated (max_item_count={PAGE_SIZE}); "
                f"got pages with sizes {[len(p) for p in pages]}; "
                f"oversized pages: {oversized}.")

            unique = set(all_ids)
            assert len(all_ids) == len(unique), (
                f"duplicates across pages: {len(all_ids)} returned, "
                f"{len(unique)} distinct, "
                f"{len(all_ids) - len(unique)} duplicate(s).")

            assert unique == ground_truth, (
                f"item-set mismatch: returned={len(unique)}, "
                f"ground_truth={len(ground_truth)}, "
                f"missing={len(ground_truth - unique)}, "
                f"unexpected={len(unique - ground_truth)}.")
        finally:
            await client.close()

    async def test_two_partition_feed_range_count_aggregate_pagination_async(self):
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 2:
                pytest.skip("Need a container with ≥ 2 physical partitions")

            chosen = None
            for i in range(len(partitions) - 1):
                p0, p1 = partitions[i], partitions[i + 1]
                if (await _count_in_range(container, p0[0], p0[1]) >= MIN_DOCS_PER_PARTITION
                        and await _count_in_range(container, p1[0], p1[1]) >= MIN_DOCS_PER_PARTITION):
                    chosen = (p0, p1)
                    break
            if chosen is None:
                pytest.skip("No adjacent partition pair both populated with ≥ "
                            f"{MIN_DOCS_PER_PARTITION} docs")

            (p0_min, _), (_, p1_max) = chosen
            crossing = _crossing_feed_range(p0_min, p1_max)
            expected_count = len(await _ids_via_per_partition_scan(container, [chosen[0], chosen[1]]))

            pager = container.query_items(
                query="SELECT VALUE COUNT(1) FROM c",
                feed_range=crossing,
                max_item_count=1,
            ).by_page()

            pages: List[List[object]] = []
            merged_rows: List[object] = []
            async for page in pager:
                items = [it async for it in page]
                pages.append(items)
                merged_rows.extend(items)

            oversized = [(i, len(p)) for i, p in enumerate(pages) if len(p) > 1]
            assert not oversized, (
                "aggregate page-size limit violated (max_item_count=1); "
                f"page sizes={[len(p) for p in pages]}, oversized={oversized}.")
            assert len(merged_rows) == 1, (
                "aggregate merge leaked partial fragments or dropped final value; "
                f"expected one merged row, got {len(merged_rows)} rows: {merged_rows}")
            assert merged_rows[0] == expected_count, (
                "merged COUNT result mismatch for two-partition crossing feed_range; "
                f"returned={merged_rows[0]}, expected={expected_count}")
            assert pager.continuation_token in (None, "", b""), (
                f"expected empty continuation after draining aggregate query; got "
                f"{pager.continuation_token!r}")
        finally:
            await client.close()

    @pytest.mark.parametrize("merge_error_type", [TypeError, KeyError])
    async def test_two_partition_feed_range_merge_fallback_preserves_rows_async(
        self, monkeypatch, caplog, merge_error_type
    ):
        """Force merge failures and verify fallback extends docs without loss."""
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 2:
                pytest.skip("Need a container with >= 2 physical partitions")

            chosen = None
            for i in range(len(partitions) - 1):
                p0, p1 = partitions[i], partitions[i + 1]
                if (await _count_in_range(container, p0[0], p0[1]) >= MIN_DOCS_PER_PARTITION
                        and await _count_in_range(container, p1[0], p1[1]) >= MIN_DOCS_PER_PARTITION):
                    chosen = (p0, p1)
                    break
            if chosen is None:
                pytest.skip("No adjacent partition pair both populated with >= "
                            f"{MIN_DOCS_PER_PARTITION} docs")

            (p0_min, _), (_, p1_max) = chosen
            crossing = _crossing_feed_range(p0_min, p1_max)
            ground_truth = await _ids_via_per_partition_scan(container, [chosen[0], chosen[1]])

            merge_call_count = 0

            def _raising_merge(*_args, **_kwargs):
                nonlocal merge_call_count
                merge_call_count += 1
                raise merge_error_type("injected-merge-failure")

            monkeypatch.setattr(_base, "_merge_query_results", _raising_merge)

            with caplog.at_level("WARNING", logger="azure.cosmos.aio._cosmos_client_connection_async"):
                pager = container.query_items(
                    query="SELECT * FROM c",
                    feed_range=crossing,
                    max_item_count=PAGE_SIZE,
                ).by_page()
                pages, all_ids = await _drain_pages(pager)

            assert merge_call_count > 0
            assert any(
                "Falling back to non-aggregate merge after aggregate merge failure" in record.getMessage()
                for record in caplog.records
            ), "Expected warning log for merge fallback path"

            oversized = [(i, len(p)) for i, p in enumerate(pages) if len(p) > PAGE_SIZE]
            assert not oversized, (
                f"page-size limit violated (max_item_count={PAGE_SIZE}); "
                f"got pages with sizes {[len(p) for p in pages]}; "
                f"oversized pages (index, size): {oversized}.")

            unique = set(all_ids)
            assert len(all_ids) == len(unique), (
                f"fallback path produced duplicates: {len(all_ids)} items returned, "
                f"{len(unique)} distinct, "
                f"{len(all_ids) - len(unique)} duplicate(s).")
            assert unique == ground_truth, (
                f"fallback path dropped/added items: returned={len(unique)} ids, "
                f"ground_truth={len(ground_truth)} ids, "
                f"missing={len(ground_truth - unique)}, "
                f"unexpected={len(unique - ground_truth)}.")
        finally:
            await client.close()

    async def test_two_partition_feed_range_min_aggregate_pagination_async(self):
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 2:
                pytest.skip("Need a container with ≥ 2 physical partitions")

            chosen = None
            for i in range(len(partitions) - 1):
                p0, p1 = partitions[i], partitions[i + 1]
                if (await _count_in_range(container, p0[0], p0[1]) >= MIN_DOCS_PER_PARTITION
                        and await _count_in_range(container, p1[0], p1[1]) >= MIN_DOCS_PER_PARTITION):
                    chosen = (p0, p1)
                    break
            if chosen is None:
                pytest.skip("No adjacent partition pair both populated with ≥ "
                            f"{MIN_DOCS_PER_PARTITION} docs")

            (p0_min, _), (_, p1_max) = chosen
            crossing = _crossing_feed_range(p0_min, p1_max)
            expected_values = await _values_via_per_partition_scan(container, [chosen[0], chosen[1]])
            expected_min = min(expected_values)

            pager = container.query_items(
                query='SELECT VALUE MIN(c["value"]) FROM c',
                feed_range=crossing,
                max_item_count=1,
            ).by_page()

            pages: List[List[object]] = []
            merged_rows: List[object] = []
            async for page in pager:
                items = [it async for it in page]
                pages.append(items)
                merged_rows.extend(items)

            oversized = [(i, len(p)) for i, p in enumerate(pages) if len(p) > 1]
            assert not oversized, (
                "aggregate page-size limit violated (max_item_count=1); "
                f"page sizes={[len(p) for p in pages]}, oversized={oversized}.")
            assert len(merged_rows) == 1, (
                "aggregate merge leaked partial fragments or dropped final value; "
                f"expected one merged row, got {len(merged_rows)} rows: {merged_rows}")
            assert merged_rows[0] == expected_min, (
                "merged MIN result mismatch for two-partition crossing feed_range; "
                f"returned={merged_rows[0]}, expected={expected_min}")
            assert pager.continuation_token in (None, "", b""), (
                f"expected empty continuation after draining aggregate query; got "
                f"{pager.continuation_token!r}")
        finally:
            await client.close()

    async def test_two_partition_feed_range_max_aggregate_pagination_async(self):
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 2:
                pytest.skip("Need a container with ≥ 2 physical partitions")

            chosen = None
            for i in range(len(partitions) - 1):
                p0, p1 = partitions[i], partitions[i + 1]
                if (await _count_in_range(container, p0[0], p0[1]) >= MIN_DOCS_PER_PARTITION
                        and await _count_in_range(container, p1[0], p1[1]) >= MIN_DOCS_PER_PARTITION):
                    chosen = (p0, p1)
                    break
            if chosen is None:
                pytest.skip("No adjacent partition pair both populated with ≥ "
                            f"{MIN_DOCS_PER_PARTITION} docs")

            (p0_min, _), (_, p1_max) = chosen
            crossing = _crossing_feed_range(p0_min, p1_max)
            expected_values = await _values_via_per_partition_scan(container, [chosen[0], chosen[1]])
            expected_max = max(expected_values)

            pager = container.query_items(
                query='SELECT VALUE MAX(c["value"]) FROM c',
                feed_range=crossing,
                max_item_count=1,
            ).by_page()

            pages: List[List[object]] = []
            merged_rows: List[object] = []
            async for page in pager:
                items = [it async for it in page]
                pages.append(items)
                merged_rows.extend(items)

            oversized = [(i, len(p)) for i, p in enumerate(pages) if len(p) > 1]
            assert not oversized, (
                "aggregate page-size limit violated (max_item_count=1); "
                f"page sizes={[len(p) for p in pages]}, oversized={oversized}.")
            assert len(merged_rows) == 1, (
                "aggregate merge leaked partial fragments or dropped final value; "
                f"expected one merged row, got {len(merged_rows)} rows: {merged_rows}")
            assert merged_rows[0] == expected_max, (
                "merged MAX result mismatch for two-partition crossing feed_range; "
                f"returned={merged_rows[0]}, expected={expected_max}")
            assert pager.continuation_token in (None, "", b""), (
                f"expected empty continuation after draining aggregate query; got "
                f"{pager.continuation_token!r}")
        finally:
            await client.close()

    async def test_exception_during_post_preserves_resume_checkpoint_async(self):
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 2:
                pytest.skip("Need a container with ≥ 2 physical partitions")

            chosen = None
            for i in range(len(partitions) - 1):
                p0, p1 = partitions[i], partitions[i + 1]
                if (await _count_in_range(container, p0[0], p0[1]) >= MIN_DOCS_PER_PARTITION
                        and await _count_in_range(container, p1[0], p1[1]) >= MIN_DOCS_PER_PARTITION):
                    chosen = (p0, p1)
                    break
            if chosen is None:
                pytest.skip("No adjacent partition pair both populated with ≥ "
                            f"{MIN_DOCS_PER_PARTITION} docs")

            (p0_min, _), (_, p1_max) = chosen
            crossing = _crossing_feed_range(p0_min, p1_max)

            client_conn = container.client_connection
            original_post = client_conn._CosmosClientConnection__Post
            call_count = 0

            async def _failing_post(*args, **kwargs):
                nonlocal call_count
                call_count += 1
                if call_count == 2:
                    raise RuntimeError("injected-post-failure")
                return await original_post(*args, **kwargs)

            client_conn._CosmosClientConnection__Post = _failing_post
            try:
                pager = container.query_items(
                    query="SELECT VALUE COUNT(1) FROM c",
                    feed_range=crossing,
                    max_item_count=1,
                ).by_page()
                with pytest.raises(RuntimeError, match="injected-post-failure"):
                    page_iter = await pager.__anext__()
                    _ = [it async for it in page_iter]
            finally:
                client_conn._CosmosClientConnection__Post = original_post

            token = client_conn.last_response_headers.get(http_constants.HttpHeaders.Continuation)
            assert token, "Expected continuation checkpoint to be stamped on POST failure"
            decoded = _decode_token(token)
            assert decoded is not None
            assert decoded["c"][0]["min"] == chosen[1][0], (
                "Checkpoint should resume from the second sub-range after the first "
                "sub-range completed successfully before failure."
            )
        finally:
            await client.close()

    async def test_three_way_overlap_async(self):
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 3:
                pytest.skip("Need a container with ≥ 3 physical partitions")

            chosen: Optional[List[Tuple[str, str]]] = None
            for i in range(len(partitions) - 2):
                triple = partitions[i:i + 3]
                ok = True
                for mn, mx in triple:
                    if await _count_in_range(container, mn, mx) < MIN_DOCS_PER_PARTITION:
                        ok = False
                        break
                if ok:
                    chosen = triple
                    break
            if chosen is None:
                pytest.skip("No three adjacent partitions all populated with ≥ "
                            f"{MIN_DOCS_PER_PARTITION} docs")
            assert chosen is not None  # narrows for type checkers; pytest.skip raises
            crossing = _crossing_feed_range(chosen[0][0], chosen[2][1])
            ground_truth = await _ids_via_per_partition_scan(container, chosen)

            pager = container.query_items(
                query="SELECT * FROM c",
                feed_range=crossing,
                max_item_count=PAGE_SIZE,
            ).by_page()
            pages, all_ids = await _drain_pages(pager)

            oversized = [(i, len(p)) for i, p in enumerate(pages) if len(p) > PAGE_SIZE]
            assert not oversized, (
                f"page-size limit violated; sizes={[len(p) for p in pages]}; "
                f"oversized={oversized}.")

            unique = set(all_ids)
            assert len(all_ids) == len(unique), (
                f"duplicates across pages: {len(all_ids)} returned, "
                f"{len(unique)} distinct, "
                f"{len(all_ids) - len(unique)} duplicate(s).")

            assert unique == ground_truth, (
                f"item-set mismatch: returned={len(unique)}, "
                f"ground_truth={len(ground_truth)}, "
                f"missing={len(ground_truth - unique)}, "
                f"unexpected={len(unique - ground_truth)}.")
        finally:
            await client.close()

    @pytest.mark.cosmosSplit
    @pytest.mark.cosmosAADSplit
    async def test_post_split_resume_async(self):
        client = _client()
        try:
            container = _get_container(client)
            partitions_before = await _sorted_partition_ranges(container)
            if len(partitions_before) < 2:
                pytest.skip("Need a container with ≥ 2 physical partitions")

            chosen = None
            for i in range(len(partitions_before) - 1):
                p0, p1 = partitions_before[i], partitions_before[i + 1]
                if (await _count_in_range(container, p0[0], p0[1]) >= MIN_DOCS_PER_PARTITION
                        and await _count_in_range(container, p1[0], p1[1]) >= MIN_DOCS_PER_PARTITION):
                    chosen = (p0, p1)
                    break
            if chosen is None:
                pytest.skip("No adjacent partition pair both populated with ≥ "
                            f"{MIN_DOCS_PER_PARTITION} docs")
            (p0_min, _), (_, p1_max) = chosen
            crossing = _crossing_feed_range(p0_min, p1_max)

            # Step 1 — drain page 1 only and save the outbound continuation.
            pager_pre = container.query_items(
                query="SELECT * FROM c",
                feed_range=crossing,
                max_item_count=PAGE_SIZE,
            ).by_page()
            page_1_iter = await pager_pre.__anext__()
            page_1 = [it async for it in page_1_iter]
            ct_after_page_1 = pager_pre.continuation_token
            page_1_ids = [item["id"] for item in page_1]
            assert ct_after_page_1, (
                "expected a non-empty continuation token after page 1")

            # Step 2 — trigger a real split.
            target_throughput = max(REPRO_THROUGHPUT * 2, 60000)
            # Split trigger is a control-plane throughput operation; run it via
            # key-auth container even in AAD lanes (same pattern as other split tests).
            key_client_for_split = CosmosClient(CONFIG.host, CONFIG.masterKey)
            try:
                key_container_for_split = _get_container(key_client_for_split)
                await test_config.TestConfig.trigger_split_async(
                    key_container_for_split,
                    target_throughput,
                )
            except unittest.SkipTest:
                raise
            finally:
                await key_client_for_split.close()
            await asyncio.sleep(10)
            _ = [fr async for fr in container.read_feed_ranges(force_refresh=True)]

            # Step 3 — resume with the saved continuation, same feed_range.
            pager_post = container.query_items(
                query="SELECT * FROM c",
                feed_range=crossing,
                max_item_count=PAGE_SIZE,
            ).by_page(continuation_token=ct_after_page_1)
            post_pages, post_ids = await _drain_pages(pager_post)

            combined_ids = page_1_ids + post_ids
            unique = set(combined_ids)
            duplicate_count = len(combined_ids) - len(unique)
            # When the parent's backend continuation is dropped during
            # the post-split explode (children won't accept the
            # parent's bc), the children restart at offset 0 of their
            # slice. The lower child can therefore re-emit up to
            # PAGE_SIZE rows that page 1 already returned. The strict
            # no-dup invariant only holds when page 1 happened to fully
            # drain the parent slice; the bounded-replay invariant
            # always holds and is what we assert here. The strict
            # no-loss / no-out-of-range guarantee is still enforced by
            # the ``unique == ground_truth`` check below.
            assert duplicate_count <= PAGE_SIZE, (
                f"unexpected duplicate count across the split boundary: "
                f"{len(combined_ids)} ids returned across page 1 + "
                f"{len(post_pages)} post-split page(s), {len(unique)} distinct, "
                f"{duplicate_count} duplicate(s) (max allowed: {PAGE_SIZE}).")

            oversized = [(i, len(p)) for i, p in enumerate(post_pages) if len(p) > PAGE_SIZE]
            assert not oversized, (
                f"page-size limit violated post-split; sizes="
                f"{[len(p) for p in post_pages]}; oversized={oversized}.")

            partitions_after = await _sorted_partition_ranges(container)
            post_split_overlaps = [(mn, mx) for (mn, mx) in partitions_after
                                   if min(p1_max, mx) > max(p0_min, mn)]
            ground_truth = await _ids_via_per_partition_scan(container, post_split_overlaps)
            assert unique == ground_truth, (
                f"item-set mismatch after post-split resume: returned="
                f"{len(unique)}, ground_truth={len(ground_truth)}, "
                f"missing={len(ground_truth - unique)}, "
                f"unexpected={len(unique - ground_truth)}.")
        finally:
            await client.close()

    async def test_legacy_opaque_token_compat_async(self, caplog):
        """Use an opaque continuation token and verify restart behavior."""
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 2:
                pytest.skip("Need a container with ≥ 2 physical partitions")

            chosen = None
            for i in range(len(partitions) - 1):
                p0, p1 = partitions[i], partitions[i + 1]
                if (await _count_in_range(container, p0[0], p0[1]) >= MIN_DOCS_PER_PARTITION
                        and await _count_in_range(container, p1[0], p1[1]) >= MIN_DOCS_PER_PARTITION):
                    chosen = (p0, p1)
                    break
            if chosen is None:
                pytest.skip("No adjacent partition pair both populated with ≥ "
                            f"{MIN_DOCS_PER_PARTITION} docs")
            (p0_min, _), (_, p1_max) = chosen
            crossing = _crossing_feed_range(p0_min, p1_max)
            ground_truth = await _ids_via_per_partition_scan(
                container, [chosen[0], chosen[1]])

            # cspell:ignore AOXB BAAAAAAAAAA EAAAAFAAAA
            legacy_token = "+RID:~Yxs1AOXBSp4BAAAAAAAAAA==#RT:1#TRC:5#ISV:2#IEO:65567#FPC:AgEAAAAFAAAA"

            with caplog.at_level("WARNING", logger="azure.cosmos.aio._cosmos_client_connection_async"):
                pager = container.query_items(
                    query="SELECT * FROM c",
                    feed_range=crossing,
                    max_item_count=PAGE_SIZE,
                ).by_page(continuation_token=legacy_token)
                pages, all_ids = await _drain_pages(pager)

            assert any(
                "not in the supported structured format" in record.getMessage()
                for record in caplog.records
            ), "Expected warning log when a non-structured continuation token is supplied"
            assert all(
                legacy_token not in record.getMessage() for record in caplog.records
            ), "Warning log must not include raw continuation token text"

            oversized = [(i, len(p)) for i, p in enumerate(pages) if len(p) > PAGE_SIZE]
            assert not oversized, (
                f"page-size limit violated under legacy-token resume; "
                f"sizes={[len(p) for p in pages]}; oversized={oversized}")

            unique = set(all_ids)
            assert len(all_ids) == len(unique), (
                f"legacy-token restart produced duplicates: "
                f"{len(all_ids) - len(unique)} duplicate id(s)")
            assert unique == ground_truth, (
                f"legacy-token restart coverage mismatch: returned="
                f"{len(unique)}, ground_truth={len(ground_truth)}, "
                f"missing={len(ground_truth - unique)}, "
                f"unexpected={len(unique - ground_truth)}")

            assert pager.continuation_token in (None, "", b""), (
                f"expected empty continuation after draining all pages on "
                f"legacy-token restart; got {pager.continuation_token!r}")
        finally:
            await client.close()

    async def test_token_identity_mismatch_rejected_async(self):
        """Live identity-mismatch rejection test."""
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 2:
                pytest.skip("Need a container with ≥ 2 physical partitions")

            chosen = None
            for i in range(len(partitions) - 1):
                p0, p1 = partitions[i], partitions[i + 1]
                if (await _count_in_range(container, p0[0], p0[1]) >= MIN_DOCS_PER_PARTITION
                        and await _count_in_range(container, p1[0], p1[1]) >= MIN_DOCS_PER_PARTITION):
                    chosen = (p0, p1)
                    break
            if chosen is None:
                pytest.skip("No adjacent partition pair both populated with ≥ "
                            f"{MIN_DOCS_PER_PARTITION} docs")
            (p0_min, _), (_, p1_max) = chosen
            crossing = _crossing_feed_range(p0_min, p1_max)

            # Use bracket notation: ``value`` is a reserved word in Cosmos SQL.
            original_query = "SELECT * FROM c WHERE c[\"value\"] >= @v"
            original_params = [{"name": "@v", "value": 0}]
            pager = container.query_items(
                query={"query": original_query, "parameters": original_params},
                feed_range=crossing,
                max_item_count=PAGE_SIZE,
            ).by_page()
            page_1_iter = await pager.__anext__()
            _ = [it async for it in page_1_iter]
            token = pager.continuation_token
            assert token, ("expected a non-empty continuation after page 1; "
                           "test cannot exercise resume validation otherwise")

            async def _drain(p):
                async for page in p:
                    _ = [it async for it in page]

            # (a) Different query TEXT.
            with pytest.raises(ValueError) as excinfo_q:
                await _drain(container.query_items(
                    query={"query": "SELECT * FROM c WHERE c[\"value\"] >= @v AND c.id != ''",
                           "parameters": original_params},
                    feed_range=crossing,
                    max_item_count=PAGE_SIZE,
                ).by_page(continuation_token=token))
            assert "query" in str(excinfo_q.value).lower(), (
                f"ValueError on query-text mismatch must name the failing "
                f"field; got: {excinfo_q.value!r}")

            # (b) Different parameter VALUE.
            with pytest.raises(ValueError) as excinfo_p:
                await _drain(container.query_items(
                    query={"query": original_query,
                           "parameters": [{"name": "@v", "value": 999999}]},
                    feed_range=crossing,
                    max_item_count=PAGE_SIZE,
                ).by_page(continuation_token=token))
            assert "query" in str(excinfo_p.value).lower(), (
                f"ValueError on parameter-value mismatch must name the "
                f"query field; got: {excinfo_p.value!r}")

            # (c) Different feed_range.
            single_p0 = _crossing_feed_range(chosen[0][0], chosen[0][1])
            with pytest.raises(ValueError) as excinfo_f:
                await _drain(container.query_items(
                    query={"query": original_query, "parameters": original_params},
                    feed_range=single_p0,
                    max_item_count=PAGE_SIZE,
                ).by_page(continuation_token=token))
            assert "feed_range" in str(excinfo_f.value).lower(), (
                f"ValueError on feed_range mismatch must name the "
                f"feed_range field; got: {excinfo_f.value!r}")
        finally:
            await client.close()

    async def test_full_partition_key_query_pagination_resume_async(self):
        """Query with a full partition key on a hierarchical container, drain
        page 1, resume from the returned continuation token, and confirm the
        resumed pages match the remaining pages of a fresh iterator and
        cover the same documents as a baseline scan in the same order.
        """
        client = _client()
        try:
            db = client.get_database_client(DATABASE_ID)
            container_id = "FeedRangeMultiPartitionAsyncFullPK-" + str(uuid.uuid4())
            created_container = await db.create_container_if_not_exists(
                id=container_id,
                partition_key=PartitionKey(
                    path=['/state', '/city', '/zipcode'],
                    kind=documents.PartitionKind.MultiHash),
                offer_throughput=400,
            )
            try:
                full_key = ['CA', 'Oxnard', '93033']
                for i in range(25):
                    await created_container.upsert_item({
                        'id': f'full-pk-doc-async-{i:03d}',
                        'state': full_key[0],
                        'city': full_key[1],
                        'zipcode': full_key[2],
                        'value': i,
                    })
                for i in range(5):
                    await created_container.upsert_item({
                        'id': f'other-doc-async-{i:03d}',
                        'state': 'WA',
                        'city': 'Seattle',
                        'zipcode': f'98{i:03d}',
                        'value': i,
                    })

                query = 'SELECT c.id FROM c ORDER BY c.id'
                query_iterable = created_container.query_items(
                    query=query,
                    partition_key=full_key,
                    max_item_count=7,
                )

                pager = query_iterable.by_page()
                first_page_iter = await pager.__anext__()
                first_page = [it async for it in first_page_iter]
                assert first_page, (
                    "first page must contain at least one item to exercise resume")
                token = pager.continuation_token
                assert token, (
                    "expected a non-empty continuation token after page 1 to "
                    "exercise the resume path")

                expected_remaining_ids: List[str] = []
                async for page in pager:
                    expected_remaining_ids.extend(
                        [it['id'] async for it in page])

                resumed_remaining_ids: List[str] = []
                async for page in query_iterable.by_page(token):
                    resumed_remaining_ids.extend(
                        [it['id'] async for it in page])

                assert expected_remaining_ids == resumed_remaining_ids, (
                    "Pages returned after resuming from the continuation token "
                    "must match the pages returned by draining a fresh iterator")

                baseline_ids: List[str] = []
                async for item in created_container.query_items(
                        query=query, partition_key=full_key):
                    baseline_ids.append(item['id'])
                fetched_ids = [item['id'] for item in first_page] + resumed_remaining_ids
                assert baseline_ids == fetched_ids, (
                    "Page 1 plus the resumed pages must equal the baseline "
                    "documents for this partition key in the same order")
            finally:
                try:
                    await db.delete_container(created_container.id)
                except Exception:  # pylint: disable=broad-except
                    pass
        finally:
            await client.close()

    async def test_prefix_partition_key_query_pagination_resume_async(self):
        """Query with a partition key prefix on a hierarchical container,
        drain page 1, resume from the returned token, and confirm the
        resumed pages match a fresh iterator and cover the baseline
        documents in the same order.
        """
        client = _client()
        try:
            db = client.get_database_client(DATABASE_ID)
            container_id = "FeedRangeMultiPartitionAsyncPrefixPK-" + str(uuid.uuid4())
            created_container = await db.create_container_if_not_exists(
                id=container_id,
                partition_key=PartitionKey(
                    path=['/state', '/city', '/zipcode'],
                    kind=documents.PartitionKind.MultiHash),
                offer_throughput=400,
            )
            try:
                for i in range(30):
                    await created_container.upsert_item({
                        'id': f'ca-doc-async-{i:03d}',
                        'state': 'CA',
                        'city': f'city-{i % 5}',
                        'zipcode': f'zip-{i:03d}',
                        'value': i,
                    })
                for i in range(6):
                    await created_container.upsert_item({
                        'id': f'wa-doc-async-{i:03d}',
                        'state': 'WA',
                        'city': f'city-{i % 2}',
                        'zipcode': f'zip-{i:03d}',
                        'value': i,
                    })

                query = 'SELECT c.id FROM c ORDER BY c.id'
                query_iterable = created_container.query_items(
                    query=query,
                    partition_key=['CA'],
                    max_item_count=7,
                )

                pager = query_iterable.by_page()
                first_page_iter = await pager.__anext__()
                first_page = [it async for it in first_page_iter]
                assert first_page
                token = pager.continuation_token
                assert token, (
                    "Expected a non-empty continuation token after page 1 "
                    "to exercise the resume path")

                expected_remaining_ids: List[str] = []
                async for page in pager:
                    expected_remaining_ids.extend(
                        [it['id'] async for it in page])

                resumed_remaining_ids: List[str] = []
                async for page in query_iterable.by_page(token):
                    resumed_remaining_ids.extend(
                        [it['id'] async for it in page])

                assert expected_remaining_ids == resumed_remaining_ids, (
                    "Pages returned after resuming from the continuation token "
                    "must match the pages returned by draining a fresh iterator")

                baseline_ids: List[str] = []
                async for item in created_container.query_items(
                        query=query, partition_key=['CA']):
                    baseline_ids.append(item['id'])
                fetched_ids = [item['id'] for item in first_page] + resumed_remaining_ids
                assert baseline_ids == fetched_ids, (
                    "Page 1 plus the resumed pages must equal the baseline "
                    "documents for this partition key prefix in the same order")
            finally:
                try:
                    await db.delete_container(created_container.id)
                except Exception:  # pylint: disable=broad-except
                    pass
        finally:
            await client.close()

    async def test_explode_iteration_guard_raises_in_query_loop_async(self, monkeypatch):
        """If the routing lookup keeps returning multiple children for the same
        range, the async query loop must give up after a fixed number of
        retries and raise instead of looping forever.
        """
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 2:
                pytest.skip("Need a container with >= 2 physical partitions")

            p0, p1 = partitions[0], partitions[1]
            crossing = _crossing_feed_range(p0[0], p1[1])
            client_conn = container.client_connection

            # Always return two children whose ranges equal the requested head
            # so the loop never makes progress and the guard must trip.
            async def _always_multi_overlap(_rid, feed_ranges, _opts):
                head = feed_ranges[0]
                return [
                    {"id": "left", "minInclusive": head.min, "maxExclusive": head.max},
                    {"id": "right", "minInclusive": head.min, "maxExclusive": head.max},
                ]

            monkeypatch.setattr(
                client_conn._routing_map_provider,
                "get_overlapping_ranges",
                _always_multi_overlap,
            )
            monkeypatch.setattr(
                "azure.cosmos._routing.feed_range_continuation."
                "_MAX_MULTI_OVERLAP_EXPLODE_ITERATIONS",
                2,
            )

            with pytest.raises(RuntimeError) as excinfo:
                pager = container.query_items(
                    query="SELECT * FROM c",
                    feed_range=crossing,
                    max_item_count=PAGE_SIZE,
                ).by_page()
                async for page in pager:
                    _ = [it async for it in page]
            assert "split re-resolution" in str(excinfo.value), (
                "Expected the safety-guard error message; "
                f"got: {excinfo.value!r}")
        finally:
            await client.close()

    async def test_no_progress_guard_logs_warning_in_query_loop_async(
        self, monkeypatch, caplog
    ):
        """If the async query loop keeps receiving empty pages with the same
        continuation token, a warning must be logged so the situation is
        visible to operators.
        """
        client = _client()
        try:
            container = _get_container(client)
            partitions = await _sorted_partition_ranges(container)
            if len(partitions) < 2:
                pytest.skip("Need a container with >= 2 physical partitions")

            p0, p1 = partitions[0], partitions[1]
            crossing = _crossing_feed_range(p0[0], p1[1])
            client_conn = container.client_connection

            post_call_count = 0

            async def _stalled_post(*_args, **_kwargs):
                nonlocal post_call_count
                post_call_count += 1
                continuation = "stalled-token-async" if post_call_count <= 3 else None
                return (
                    {"Documents": []},
                    {http_constants.HttpHeaders.Continuation: continuation},
                )

            monkeypatch.setattr(
                client_conn,
                "_CosmosClientConnection__Post",
                _stalled_post,
            )
            monkeypatch.setattr(
                "azure.cosmos.aio._cosmos_client_connection_async."
                "_MAX_CONSECUTIVE_NO_PROGRESS_PAGES",
                2,
            )

            with caplog.at_level(
                "WARNING",
                logger="azure.cosmos.aio._cosmos_client_connection_async",
            ):
                pager = container.query_items(
                    query="SELECT * FROM c",
                    feed_range=crossing,
                    max_item_count=PAGE_SIZE,
                ).by_page()
                async for page in pager:
                    _ = [it async for it in page]

            assert post_call_count >= 3, (
                "Expected at least 3 stalled __Post calls to trigger the "
                f"no-progress guard; got {post_call_count}")
            assert any(
                "same continuation token" in record.getMessage()
                for record in caplog.records
            ), (
                "Expected warning log from the async no-progress guard "
                "mentioning the unchanged continuation token")
        finally:
            await client.close()


if __name__ == "__main__":
    unittest.main()


