import pytest
from sqlalchemy import event, select
from sqlalchemy.ext.asyncio import AsyncSession

from sqlalchemy_relay_pagination import paginate_many
from sqlalchemy_relay_pagination.exceptions import InvalidArgumentsError
from sqlalchemy_relay_pagination.types import RequestedFields
from tests.conftest import Post

# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------


def _make_query_counter(async_engine):
    counts = {"n": 0}

    def before_exec(conn, cursor, statement, parameters, context, executemany):
        stmt = statement.strip().upper()
        if stmt.startswith(("SELECT", "INSERT", "UPDATE", "DELETE")):
            counts["n"] += 1

    event.listen(async_engine.sync_engine, "before_cursor_execute", before_exec)
    return counts, lambda: event.remove(async_engine.sync_engine, "before_cursor_execute", before_exec)


# ---------------------------------------------------------------------------
# Forward pagination
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_basic_forward_first_page(async_session: AsyncSession, seeded_posts):
    result = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1, 2, 3],
        first=2,
    )
    assert set(result.keys()) == {1, 2, 3}
    for key in [1, 2, 3]:
        conn = result[key]
        assert len(conn.edges) == 2
        assert conn.pageInfo.hasNextPage is True


@pytest.mark.asyncio
async def test_forward_fits_on_one_page(async_session: AsyncSession, seeded_posts):
    result = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1, 2],
        first=5,
    )
    for key in [1, 2]:
        conn = result[key]
        assert len(conn.edges) == 3
        assert conn.pageInfo.hasNextPage is False


# ---------------------------------------------------------------------------
# Backward pagination
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_last_backward(async_session: AsyncSession, seeded_posts):
    result = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1, 2],
        last=2,
    )
    for key in [1, 2]:
        conn = result[key]
        assert len(conn.edges) == 2
        assert conn.pageInfo.hasPreviousPage is True
        # edges should be in ascending id order (reversed back by _assemble_connection)
        ids = [e.node.id for e in conn.edges]
        assert ids == sorted(ids)


@pytest.mark.asyncio
async def test_last_fits_on_one_page(async_session: AsyncSession, seeded_posts):
    result = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1, 2],
        last=5,
    )
    for key in [1, 2]:
        conn = result[key]
        assert len(conn.edges) == 3
        assert conn.pageInfo.hasPreviousPage is False


# ---------------------------------------------------------------------------
# total_count
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_total_count(async_session: AsyncSession, seeded_posts):
    fields = RequestedFields(
        total_count=True,
        edges=True,
        has_next_page=False,
        has_previous_page=False,
        start_cursor=False,
        end_cursor=False,
    )
    result = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1, 2, 3],
        first=5,
        requested_fields=fields,
    )
    for key in [1, 2, 3]:
        assert result[key].pageInfo.totalCount == 3


@pytest.mark.asyncio
async def test_total_count_single_query(async_session: AsyncSession, async_engine, seeded_posts):
    counts, remove = _make_query_counter(async_engine)
    try:
        fields = RequestedFields(
            total_count=True,
            edges=True,
            has_next_page=False,
            has_previous_page=False,
            start_cursor=False,
            end_cursor=False,
        )
        await paginate_many(
            async_session,
            select(Post),
            model=Post,
            partition_by=Post.user_id,
            order_by=Post.id,
            keys=[1, 2, 3],
            first=5,
            requested_fields=fields,
        )
    finally:
        remove()
    assert counts["n"] == 1


# ---------------------------------------------------------------------------
# Empty / missing keys
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_empty_key(async_session: AsyncSession, seeded_posts):
    # user_id=999 has no posts
    fields = RequestedFields(
        total_count=True,
        edges=True,
        has_next_page=False,
        has_previous_page=False,
        start_cursor=False,
        end_cursor=False,
    )
    result = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1, 999],
        first=5,
        requested_fields=fields,
    )
    assert result[999].edges == []
    assert result[999].pageInfo.totalCount == 0


@pytest.mark.asyncio
async def test_empty_keys_list(async_session: AsyncSession, async_engine, seeded_posts):
    counts, remove = _make_query_counter(async_engine)
    try:
        result = await paginate_many(
            async_session,
            select(Post),
            model=Post,
            partition_by=Post.user_id,
            order_by=Post.id,
            keys=[],
            first=5,
        )
    finally:
        remove()
    assert result == {}
    assert counts["n"] == 0


# ---------------------------------------------------------------------------
# Key ordering
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_preserves_key_order(async_session: AsyncSession, seeded_posts):
    result = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[3, 1, 2],
        first=5,
    )
    assert list(result.keys()) == [3, 1, 2]


# ---------------------------------------------------------------------------
# Descending order
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_desc_order_by(async_session: AsyncSession, seeded_posts):
    # user 1 has posts 1, 2, 3 — desc + first=2 should yield [3, 2]
    result = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id.desc(),
        keys=[1],
        first=2,
    )
    ids = [e.node.id for e in result[1].edges]
    assert ids == [3, 2]
    assert result[1].pageInfo.hasNextPage is True


# ---------------------------------------------------------------------------
# No-query short-circuit
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_no_requested_fields_no_query(async_session: AsyncSession, async_engine, seeded_posts):
    counts, remove = _make_query_counter(async_engine)
    try:
        fields = RequestedFields(
            edges=False,
            has_next_page=False,
            has_previous_page=False,
            start_cursor=False,
            end_cursor=False,
            total_count=False,
        )
        result = await paginate_many(
            async_session,
            select(Post),
            model=Post,
            partition_by=Post.user_id,
            order_by=Post.id,
            keys=[1, 2],
            first=5,
            requested_fields=fields,
        )
    finally:
        remove()
    assert counts["n"] == 0
    assert result[1].edges == []
    assert result[2].edges == []


# ---------------------------------------------------------------------------
# Cursor encoding
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_cursor_values(async_session: AsyncSession, seeded_posts):
    result = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1],
        first=3,
    )
    conn = result[1]
    # user 1 has post IDs 1, 2, 3
    assert len(conn.edges) == 3
    for edge in conn.edges:
        assert edge.cursor == str(edge.node.id)


# ---------------------------------------------------------------------------
# after / before fallback (N individual queries, still correct)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_after_cursor_second_page(async_session: AsyncSession, seeded_posts):
    # First page: first=2 → posts 1, 2 for user 1; posts 4, 5 for user 2
    page1 = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1, 2],
        first=2,
    )
    after_dict = {key: page1[key].pageInfo.endCursor for key in [1, 2]}

    # Second page using after cursors: only 1 post remains per user
    page2 = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1, 2],
        first=2,
        after=after_dict,
    )
    # user 1: post 3; user 2: post 6
    assert [e.node.id for e in page2[1].edges] == [3]
    assert [e.node.id for e in page2[2].edges] == [6]
    assert page2[1].pageInfo.hasNextPage is False
    assert page2[2].pageInfo.hasNextPage is False


@pytest.mark.asyncio
async def test_after_cursor_uses_n_queries(async_session: AsyncSession, async_engine, seeded_posts):
    keys = [1, 2, 3]
    page1 = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=keys,
        first=1,
    )
    after_dict = {key: page1[key].pageInfo.endCursor for key in keys}

    counts, remove = _make_query_counter(async_engine)
    try:
        await paginate_many(
            async_session,
            select(Post),
            model=Post,
            partition_by=Post.user_id,
            order_by=Post.id,
            keys=keys,
            first=2,
            after=after_dict,
        )
    finally:
        remove()
    # One query per key (N queries, no batching)
    assert counts["n"] == len(keys)


@pytest.mark.asyncio
async def test_before_cursor_last_page(async_session: AsyncSession, seeded_posts):
    # Get cursor at post 3 (end of user 1's list)
    page1 = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1],
        first=3,
    )
    before_cursor = page1[1].pageInfo.endCursor  # cursor at post 3

    result = await paginate_many(
        async_session,
        select(Post),
        model=Post,
        partition_by=Post.user_id,
        order_by=Post.id,
        keys=[1],
        last=2,
        before={1: before_cursor},
    )
    ids = [e.node.id for e in result[1].edges]
    # last=2 before post 3 → posts 1, 2
    assert ids == [1, 2]


# ---------------------------------------------------------------------------
# Query validation
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_rejects_query_with_order_by(async_session: AsyncSession):
    bad_query = select(Post).order_by(Post.id)
    with pytest.raises(InvalidArgumentsError, match="ORDER BY"):
        await paginate_many(
            async_session,
            bad_query,
            model=Post,
            partition_by=Post.user_id,
            order_by=Post.id,
            keys=[1],
            first=5,
        )


@pytest.mark.asyncio
async def test_rejects_query_with_limit(async_session: AsyncSession):
    bad_query = select(Post).limit(10)
    with pytest.raises(InvalidArgumentsError, match="LIMIT"):
        await paginate_many(
            async_session,
            bad_query,
            model=Post,
            partition_by=Post.user_id,
            order_by=Post.id,
            keys=[1],
            first=5,
        )


@pytest.mark.asyncio
async def test_rejects_query_with_offset(async_session: AsyncSession):
    bad_query = select(Post).offset(5)
    with pytest.raises(InvalidArgumentsError, match="OFFSET"):
        await paginate_many(
            async_session,
            bad_query,
            model=Post,
            partition_by=Post.user_id,
            order_by=Post.id,
            keys=[1],
            first=5,
        )
