"""Recording transports for Spotify's httpx2 and Anthropic's httpx."""

from collections.abc import AsyncIterator, Callable
from typing import cast

import httpx
import httpx2

from recording.model import RecordedInteraction
from recording.redaction import is_sensitive_url


class HttpxRecordingTransport(httpx.AsyncBaseTransport):
    """Wrap an httpx transport and retain completed response chunk sequences."""

    def __init__(self, transport: httpx.AsyncBaseTransport) -> None:
        """Bind the real transport and start an empty in-memory cassette."""
        self.transport = transport
        self.interactions: list[RecordedInteraction] = []

    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
        """Forward one request and wrap its response stream for capture."""
        request.headers["Accept-Encoding"] = "identity"
        request_body = await request.aread()
        response = await self.transport.handle_async_request(request)
        if is_sensitive_url(str(request.url)):
            return response
        stream = _HttpxRecordingStream(
            cast(httpx.AsyncByteStream, response.stream),
            lambda chunks: self._finish(request, request_body, response.status_code, chunks),
        )
        return httpx.Response(
            response.status_code,
            headers=response.headers,
            stream=stream,
            extensions=response.extensions,
            request=request,
        )

    async def aclose(self) -> None:
        """Close the wrapped transport."""
        await self.transport.aclose()

    def _finish(
        self,
        request: httpx.Request,
        request_body: bytes,
        status: int,
        chunks: tuple[bytes, ...],
    ) -> None:
        self.interactions.append(
            RecordedInteraction(request.method, str(request.url), status, request_body, chunks)
        )


class Httpx2RecordingTransport(httpx2.AsyncBaseTransport):
    """Wrap an httpx2 transport and retain completed response chunk sequences."""

    def __init__(self, transport: httpx2.AsyncBaseTransport) -> None:
        """Bind the real transport and start an empty in-memory cassette."""
        self.transport = transport
        self.interactions: list[RecordedInteraction] = []

    async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
        """Forward one request and wrap its response stream for capture."""
        request.headers["Accept-Encoding"] = "identity"
        request_body = await request.aread()
        response = await self.transport.handle_async_request(request)
        if is_sensitive_url(str(request.url)):
            return response
        stream = _Httpx2RecordingStream(
            cast(httpx2.AsyncByteStream, response.stream),
            lambda chunks: self._finish(request, request_body, response.status_code, chunks),
        )
        return httpx2.Response(
            response.status_code,
            headers=response.headers,
            stream=stream,
            extensions=response.extensions,
            request=request,
        )

    async def aclose(self) -> None:
        """Close the wrapped transport."""
        await self.transport.aclose()

    def _finish(
        self,
        request: httpx2.Request,
        request_body: bytes,
        status: int,
        chunks: tuple[bytes, ...],
    ) -> None:
        self.interactions.append(
            RecordedInteraction(request.method, str(request.url), status, request_body, chunks)
        )


class _HttpxRecordingStream(httpx.AsyncByteStream):
    def __init__(
        self,
        stream: httpx.AsyncByteStream,
        finish: Callable[[tuple[bytes, ...]], None],
    ) -> None:
        self.stream = stream
        self.finish = finish
        self.chunks: list[bytes] = []
        self.is_finished = False

    async def __aiter__(self) -> AsyncIterator[bytes]:
        try:
            async for chunk in self.stream:
                self.chunks.append(chunk)
                yield chunk
        finally:
            self._finish()

    async def aclose(self) -> None:
        await self.stream.aclose()
        self._finish()

    def _finish(self) -> None:
        if not self.is_finished:
            self.is_finished = True
            self.finish(tuple(self.chunks))


class _Httpx2RecordingStream(httpx2.AsyncByteStream):
    def __init__(
        self,
        stream: httpx2.AsyncByteStream,
        finish: Callable[[tuple[bytes, ...]], None],
    ) -> None:
        self.stream = stream
        self.finish = finish
        self.chunks: list[bytes] = []
        self.is_finished = False

    async def __aiter__(self) -> AsyncIterator[bytes]:
        try:
            async for chunk in self.stream:
                self.chunks.append(chunk)
                yield chunk
        finally:
            self._finish()

    async def aclose(self) -> None:
        await self.stream.aclose()
        self._finish()

    def _finish(self) -> None:
        if not self.is_finished:
            self.is_finished = True
            self.finish(tuple(self.chunks))