"""High-level RPC client — async and sync APIs.

The async ``Client`` uses `anyio <https://anyio.readthedocs.io/>`_ for
structured concurrency: task groups for batched calls, memory object
streams for streaming with backpressure, and backend-agnostic execution
(asyncio, Trio, uvloop).
"""

from __future__ import annotations

from typing import Any, AsyncIterator, Iterator, Sequence

import anyio

from knafeh._native import RpcClient as _NativeClient
from knafeh._native import TlsConfig


class Client:
    """An async RPC client that connects to a Knafeh server.

    Uses anyio structured concurrency for safe concurrent operations.

    Example::

        async with Client("localhost:4433", tls=TlsConfig.client_insecure()) as client:
            response = await client.call("greeter/say_hello", b'{"name": "World"}')
            print(response)

            # Batch concurrent calls with structured cancellation
            responses = await client.call_many("echo/echo", [b'{"i":0}', b'{"i":1}'])
    """

    def __init__(
        self,
        endpoint: str,
        *,
        tls: TlsConfig | None = None,
        pool_size: int = 4,
    ) -> None:
        if tls is None:
            tls = TlsConfig.client_insecure()
        self._native = _NativeClient(endpoint, tls, pool_size)

    async def __aenter__(self) -> Client:
        self._native.connect()
        return self

    async def __aexit__(self, *exc: Any) -> None:
        self._native.close()

    def connect(self) -> None:
        """Explicitly connect to the server (synchronous one-time setup)."""
        self._native.connect()

    async def call(self, method: str, body: bytes) -> bytes:
        """Make a unary RPC call.

        Args:
            method: The fully-qualified method path (e.g., ``"greeter/say_hello"``).
            body: The request payload (codec-encoded).

        Returns:
            The response payload bytes.
        """
        return await self._native.call(method, body)

    async def call_many(
        self,
        method: str,
        bodies: Sequence[bytes],
        *,
        max_concurrency: int | None = None,
    ) -> list[bytes]:
        """Execute multiple calls concurrently with structured cancellation.

        Uses an anyio ``TaskGroup`` so that if any call fails, all
        in-flight calls are cancelled and the exception propagates
        immediately. A ``CapacityLimiter`` bounds the number of
        simultaneous in-flight RPCs to avoid overwhelming the
        connection pool.

        Args:
            method: The RPC method path (same for all calls).
            bodies: A sequence of request payloads.
            max_concurrency: Maximum simultaneous in-flight calls.
                Defaults to 50 (well within the QUIC stream limit).

        Returns:
            A list of response payloads, in the same order as *bodies*.
        """
        if max_concurrency is None:
            max_concurrency = 50
        if max_concurrency < 1:
            raise ValueError("max_concurrency must be at least 1")

        results: list[bytes | None] = [None] * len(bodies)
        iterator = iter(enumerate(bodies))
        lock = anyio.Lock()

        async def worker() -> None:
            while True:
                async with lock:
                    try:
                        idx, body = next(iterator)
                    except StopIteration:
                        return
                results[idx] = await self._native.call(method, body)

        async with anyio.create_task_group() as tg:
            for _ in range(min(max_concurrency, len(bodies))):
                tg.start_soon(worker)

        return results  # type: ignore[return-value]

    async def server_stream(
        self,
        method: str,
        body: bytes,
        *,
        max_buffer_size: int = 32,
    ) -> AsyncIterator[bytes]:
        """Initiate a server-streaming RPC call.

        Yields response chunks as they arrive. Uses an anyio memory
        object stream internally for backpressure between the Rust
        frame reader and the Python consumer.

        Args:
            method: The fully-qualified method path.
            body: The request payload.
            max_buffer_size: Backpressure buffer depth.

        Yields:
            Response payload chunks as bytes.
        """
        stream_receiver = await self._native.server_stream(method, body)
        send_stream, receive_stream = anyio.create_memory_object_stream[bytes](
            max_buffer_size=max_buffer_size,
        )

        async def _pump() -> None:
            """Read chunks from Rust and pump into the anyio stream."""
            async with send_stream:
                while True:
                    chunk = await stream_receiver.next_chunk()
                    if chunk is None:
                        break
                    await send_stream.send(bytes(chunk))

        async with anyio.create_task_group() as tg:
            tg.start_soon(_pump)
            async with receive_stream:
                async for chunk in receive_stream:
                    yield chunk
            # When the consumer exits (even via break/exception), the
            # task group cancels _pump, which drops stream_receiver.
            tg.cancel_scope.cancel()

    def close(self) -> None:
        """Close the client and release all connections."""
        self._native.close()


class SyncClient:
    """A synchronous RPC client for scripts, notebooks, and non-async contexts.

    Example::

        with SyncClient("localhost:4433", tls=TlsConfig.client_insecure()) as client:
            response = client.call("greeter/say_hello", b'{"name": "World"}')
            print(response)
    """

    def __init__(
        self,
        endpoint: str,
        *,
        tls: TlsConfig | None = None,
        pool_size: int = 4,
    ) -> None:
        if tls is None:
            tls = TlsConfig.client_insecure()
        self._native = _NativeClient(endpoint, tls, pool_size)
        self._native.connect()

    def __enter__(self) -> SyncClient:
        return self

    def __exit__(self, *exc: Any) -> None:
        self._native.close()

    def call(self, method: str, body: bytes) -> bytes:
        """Make a unary RPC call (blocking)."""
        return bytes(self._native.call_sync(method, body))

    def server_stream(self, method: str, body: bytes) -> Iterator[bytes]:
        """Initiate a server-streaming RPC call (blocking).

        Yields response chunks as they arrive.
        """
        receiver = self._native.server_stream_sync(method, body)
        try:
            for chunk in receiver:
                yield bytes(chunk)
        finally:
            receiver.close()

    def close(self) -> None:
        """Close the client and release all connections."""
        self._native.close()
