"""
WebTransport stream utility functions.
"""

from ..exceptions import StreamError
from ..utils import get_logger
from .stream import WebTransportReceiveStream, WebTransportSendStream, WebTransportStream

__all__ = [
    "copy_stream_data",
    "echo_stream",
]

logger = get_logger("stream.utils")


async def copy_stream_data(
    *,
    source: "WebTransportReceiveStream",
    destination: "WebTransportSendStream",
    chunk_size: int = 8192,
) -> int:
    """Copy all data from a source stream to a destination stream."""
    total_bytes = 0
    try:
        async for chunk in source.read_iter(chunk_size=chunk_size):
            await destination.write(chunk)
            total_bytes += len(chunk)
        await destination.close()
    except StreamError as e:
        logger.error(f"Error copying stream data: {e}")
        await destination.abort(code=1)
        raise
    return total_bytes


async def echo_stream(*, stream: "WebTransportStream") -> None:
    """Echo all data received on a bidirectional stream back to the sender."""
    try:
        async for chunk in stream.read_iter():
            await stream.write(chunk)
        await stream.close()
    except StreamError as e:
        logger.error(f"Error in echo stream: {e}")
        await stream.abort(code=1)
