# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# pyformat: disable

"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""

from __future__ import annotations

import functools
import json
from datetime import timedelta
from types import TracebackType
from typing import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Dict,
    Generic,
    Iterator,
    Literal,
    Mapping,
    NoReturn,
    Optional,
    ParamSpec,
    Tuple,
    Type,
    TypeVar,
    Union,
    cast,
    get_args,
    get_origin,
    overload,
)

import httpx
from pydantic import BaseModel
from typing_extensions import TypeAliasType

from .._version import __response_mode_header__
from .._hooks.types import AfterParseErrorContext
from .eventstreaming import Stream, AsyncStream
from .unmarshal_json_response import unmarshal_json_response

P = ParamSpec("P")
T = TypeVar("T")
U = TypeVar("U")
ResponseT = TypeVar("ResponseT", bound="_APIResponseBase[Any]")
AsyncResponseT = TypeVar("AsyncResponseT", bound="_AsyncAPIResponseBase[Any]")

ResponseMode = Literal["buffered", "raw_stream", "event_stream"]
ResponseModeHeader = Literal["parsed", "raw", "streaming"]

_DEFAULT_PARSE_KEY = object()
_PARSED_RESPONSE_MODE: ResponseModeHeader = "parsed"
_RAW_RESPONSE_MODE: ResponseModeHeader = "raw"
_STREAMING_RESPONSE_MODE: ResponseModeHeader = "streaming"

# Internal marker header set by the raw/streaming response helpers and removed
# before the request is sent, so it never reaches the wire.
_RESPONSE_MODE_HEADER = __response_mode_header__.lower()


def raise_parse_error(
    hooks: Optional[Any],
    hook_ctx: Optional[AfterParseErrorContext],
    response: httpx.Response,
    exc: Exception,
) -> NoReturn:
    """Raise the hook-substituted parse error.

    When hooks return the same exception instance, preserve its existing cause
    chain instead of overwriting it with `raise exc from exc`.
    """
    parse_error = exc
    if hooks is not None and hook_ctx is not None:
        parse_error = hooks.after_parse_error(hook_ctx, response, exc)
    if parse_error is exc:
        raise exc.with_traceback(exc.__traceback__)
    raise parse_error from exc


async def raise_parse_error_async(
    async_hooks: Optional[Any],
    hooks: Optional[Any],
    hook_ctx: Optional[AfterParseErrorContext],
    response: httpx.Response,
    exc: Exception,
) -> NoReturn:
    """Raise the async hook-substituted parse error."""
    parse_error = exc
    if async_hooks is not None and hook_ctx is not None:
        parse_error = await async_hooks.after_parse_error(hook_ctx, response, exc)
    elif hooks is not None and hook_ctx is not None:
        parse_error = hooks.after_parse_error(hook_ctx, response, exc)
    if parse_error is exc:
        raise exc.with_traceback(exc.__traceback__)
    raise parse_error from exc


def _unwrap_type(t: Any) -> Any:
    """Strip ``TypeAliasType('Name', T)`` and ``Annotated[T, ...]`` wrappers."""
    if isinstance(t, TypeAliasType):
        t = t.__value__
    if get_origin(t) is not None and hasattr(t, "__metadata__"):
        t = t.__origin__
    return t


def _event_stream_origin(
    to: Any,
) -> Optional[
    Union[
        Type[Stream[Any]],
        Type[AsyncStream[Any]],
    ]
]:
    """Resolve ``to`` to its underlying ``Stream`` / ``AsyncStream``
    subclass, or ``None`` if it isn't one. Accepts subclasses, generic
    parametrizations (``Stream[T]``), ``Annotated[...]``, and
    ``TypeAliasType``."""
    unwrapped = _unwrap_type(to)
    origin = get_origin(unwrapped) or unwrapped
    if isinstance(origin, type) and issubclass(origin, (Stream, AsyncStream)):
        return origin
    return None


def _extract_stream_chunk_type(to: Any) -> Optional[type]:
    """Recover the chunk type ``T`` from ``Stream[T]`` or from a
    pre-bound subclass ``class MyStream(Stream[T]): ...``."""
    unwrapped = _unwrap_type(to)
    args = get_args(unwrapped)
    if args:
        return args[0]
    for base in getattr(unwrapped, "__orig_bases__", ()):
        base_args = get_args(base)
        if base_args:
            return base_args[0]
    return None


def _can_honor_event_stream(to: Any, response: httpx.Response) -> bool:
    """True when spec-declared ``event_stream`` is confirmed by caller intent
    or wire Content-Type. ``to=None`` defers to the spec-driven default parser
    and must not be demoted — preserves streaming semantics when the wire
    Content-Type is unreliable (proxy strip, server misconfig)."""
    if to is None:
        return True
    if _event_stream_origin(to) is not None:
        return True
    content_type = response.headers.get("content-type", "")
    media_type = content_type.partition(";")[0].strip().lower()
    return media_type == "text/event-stream"


def _build_event_stream(
    *,
    to: Any,
    decoder: Optional[Callable[[str], Any]],
    sentinel: Optional[str],
    response: httpx.Response,
    client_ref: Optional[object],
    async_: bool,
) -> Any:
    """Dispatch ``parse(to=Stream)`` to a typed SSE iterator."""

    base_class = AsyncStream if async_ else Stream
    origin = _event_stream_origin(to)
    if origin is None or not issubclass(origin, base_class):
        iter_methods = (
            "aiter_lines() / aiter_bytes()" if async_ else "iter_lines() / iter_bytes()"
        )
        raise TypeError(
            f"parse(to=...) on SSE responses must use {base_class.__name__}[T]; got {to!r}. Iterate via {iter_methods} for raw access."
        )

    if decoder is not None:
        resolved_decoder: Callable[[str], Any] = decoder
    else:
        chunk_t = _extract_stream_chunk_type(to)
        if chunk_t is None:
            raise TypeError(
                f"parse(to={base_class.__name__}[T]) requires a type parameter."
            )
        if not (isinstance(chunk_t, type) and issubclass(chunk_t, BaseModel)):
            raise TypeError(
                f"parse(to={base_class.__name__}[T]) requires T to be a BaseModel when decoder= is not supplied; got {chunk_t!r}."
            )

        def _synthesized_decoder(raw: str, _t: Any = chunk_t) -> Any:
            envelope = json.loads(raw)
            if not isinstance(envelope, dict) or "data" not in envelope:
                raise ValueError(
                    f"Synthesized SSE decoder expected an envelope of shape {{'data': ...}}, got {envelope!r}. Pass decoder=<fn> to parse(...) to handle non-standard envelopes."
                )
            return _t.model_validate(envelope["data"])

        resolved_decoder = _synthesized_decoder

    return origin(
        response,
        resolved_decoder,
        sentinel=sentinel,
        client_ref=client_ref,
    )


class _APIResponseBase(Generic[T]):
    def __init__(
        self,
        *,
        raw: httpx.Response,
        parser: Callable[[httpx.Response], T],
        mode: ResponseMode = "buffered",
        client_ref: Optional[object] = None,
        hook_ctx: Optional[Any] = None,
        hooks: Optional[Any] = None,
    ) -> None:
        self.http_response = raw
        self._parser = parser
        self._mode = mode
        self._client_ref = client_ref
        self._hook_ctx = hook_ctx
        self._hooks = hooks
        self._parsed_by_type: dict[Any, Any] = {}

    def _raise_parse_error(self, exc: Exception) -> NoReturn:
        if self._hook_ctx is None:
            raise_parse_error(None, None, self.http_response, exc)
        raise_parse_error(self._hooks, self._hook_ctx, self.http_response, exc)

    @property
    def headers(self) -> httpx.Headers:
        return self.http_response.headers

    @property
    def http_request(self) -> httpx.Request:
        return self.http_response.request

    @property
    def status_code(self) -> int:
        return self.http_response.status_code

    @property
    def request_id(self) -> Optional[str]:
        return self.http_response.headers.get("x-request-id")

    @property
    def url(self) -> httpx.URL:
        return self.http_response.url

    @property
    def method(self) -> str:
        return self.http_request.method

    @property
    def http_version(self) -> str:
        return self.http_response.http_version

    @property
    def elapsed(self) -> timedelta:
        return self.http_response.elapsed

    @property
    def is_closed(self) -> bool:
        return self.http_response.is_closed

    def _parse(
        self,
        *,
        to: Optional[Type[U]] = None,
        decoder: Optional[Callable[[str], Any]] = None,
        sentinel: Optional[str] = None,
    ) -> T | U:
        mode = self._mode
        if mode == "event_stream" and not _can_honor_event_stream(
            to, self.http_response
        ):
            mode = "buffered"
        if to is not None and mode == "event_stream":
            return _build_event_stream(
                to=to,
                decoder=decoder,
                sentinel=sentinel,
                response=self.http_response,
                client_ref=self._client_ref,
                async_=False,
            )
        parse_key = to or _DEFAULT_PARSE_KEY
        if parse_key not in self._parsed_by_type:
            # raw_stream / event_stream wrappers leave the body unread so the
            # caller can iter_bytes/iter_lines; buffered mode eagerly reads.
            body = None
            try:
                if to is None:
                    if mode == "buffered":
                        self.read()
                    self._parsed_by_type[parse_key] = self._parser(self.http_response)
                else:
                    body = self.text()
                    self._parsed_by_type[parse_key] = unmarshal_json_response(
                        to, self.http_response, body, validate=False
                    )
            except Exception as exc:
                self._raise_parse_error(exc)
        return self._parsed_by_type[parse_key]

    def read(self) -> bytes:
        return self.http_response.read()

    def text(self) -> str:
        self.read()
        return self.http_response.text

    def json(self) -> Any:
        self.read()
        return self.http_response.json()

    def close(self) -> None:
        self.http_response.close()

    def __enter__(self: ResponseT) -> ResponseT:
        return self

    def __exit__(
        self,
        exc_type: Optional[type[BaseException]],
        exc: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        self.close()

    def iter_bytes(self, chunk_size: Optional[int] = None) -> Iterator[bytes]:
        yield from self.http_response.iter_bytes(chunk_size)

    def iter_text(self, chunk_size: Optional[int] = None) -> Iterator[str]:
        yield from self.http_response.iter_text(chunk_size)

    def iter_lines(self) -> Iterator[str]:
        yield from self.http_response.iter_lines()


class APIResponse(_APIResponseBase[T]):
    @overload
    def parse(self) -> T: ...

    @overload
    def parse(self, *, to: Type[U]) -> U: ...

    @overload
    def parse(self, *, to: Type[U], decoder: Callable[[str], Any]) -> U: ...

    @overload
    def parse(self, *, to: Type[U], sentinel: str) -> U: ...

    @overload
    def parse(
        self,
        *,
        to: Type[U],
        decoder: Callable[[str], Any],
        sentinel: str,
    ) -> U: ...

    def parse(
        self,
        *,
        to: Optional[Type[U]] = None,
        decoder: Optional[Callable[[str], Any]] = None,
        sentinel: Optional[str] = None,
    ) -> T | U:
        """Decode the response body into a typed value.
        Three modes, depending on the value provided for ``to``:

        1. ``to=None`` (default) — runs the operation's generator-emitted
           parser and returns ``T``.

        2. ``to=<BaseModel>`` on a non-SSE response — reads the body as
           JSON and validates against the given Pydantic model. Returns an
           instance of that model.

        3. ``to=Stream[T]`` on an SSE response — constructs a typed
           iterator over Server-Sent Events frames. Dispatch is driven by
           the spec-derived ``mode="event_stream"`` set at wrapper construction,
           not by the wire ``Content-Type`` header.

        SSE frame decoding
        ------------------
        Each SSE frame is wrapped into a JSON envelope of shape
        ``{"id"?, "event"?, "data"?, "retry"?}`` (None-valued fields
        omitted) before reaching the decoder. The decoder receives the
        serialized envelope as a single string per frame.

        * ``decoder=None`` — synthesizes a decoder that JSON-parses the
          envelope, takes ``["data"]``, and validates as ``T`` via
          ``T.model_validate(...)``. ``T`` must subclass
          ``pydantic.BaseModel``. Correct when the wire ``data:`` payload
          is itself an instance of ``T``.

        * ``decoder=<Callable[[str], T]>`` — bypasses synthesis. The
          callable receives the JSON-encoded envelope and returns ``T``.
          Use when the wire payload has an outer wrapper that must be
          peeled, when ``T`` is a discriminated union, or whenever the
          synthesized mapping does not match the spec.

        * ``sentinel=<str>`` — when an SSE ``data:`` line equals this
          string verbatim, iteration terminates cleanly. Use for
          protocols that mark end-of-stream with a literal token.

        Errors
        ------
        * ``TypeError`` if ``to`` is supplied on an SSE response
          (``mode="event_stream"``) and is neither ``Stream[T]``
          nor ``AsyncStream[T]``.
        * ``TypeError`` if ``to`` is the right stream class but ``T`` is
          not a ``BaseModel`` and ``decoder`` was not supplied.
        * ``pydantic.ValidationError`` from the (synthesized or
          user-supplied) decoder when the wire shape does not match
          ``T``.
        """
        return self._parse(to=to, decoder=decoder, sentinel=sentinel)


class StreamedAPIResponse(_APIResponseBase[T]):
    def parse(self) -> T:
        return self._parse()


class _AsyncAPIResponseBase(Generic[T]):
    def __init__(
        self,
        *,
        raw: httpx.Response,
        parser: Callable[[httpx.Response], Awaitable[T]],
        mode: ResponseMode = "buffered",
        client_ref: Optional[object] = None,
        hook_ctx: Optional[Any] = None,
        hooks: Optional[Any] = None,
        async_hooks: Optional[Any] = None,
    ) -> None:
        self.http_response = raw
        self._parser = parser
        self._mode = mode
        self._client_ref = client_ref
        self._hook_ctx = hook_ctx
        self._hooks = hooks
        self._async_hooks = async_hooks
        self._parsed_by_type: dict[Any, Any] = {}

    async def _raise_parse_error(self, exc: Exception) -> NoReturn:
        if self._hook_ctx is None:
            raise_parse_error(None, None, self.http_response, exc)
        await raise_parse_error_async(
            self._async_hooks, self._hooks, self._hook_ctx, self.http_response, exc
        )

    @property
    def headers(self) -> httpx.Headers:
        return self.http_response.headers

    @property
    def http_request(self) -> httpx.Request:
        return self.http_response.request

    @property
    def status_code(self) -> int:
        return self.http_response.status_code

    @property
    def request_id(self) -> Optional[str]:
        return self.http_response.headers.get("x-request-id")

    @property
    def url(self) -> httpx.URL:
        return self.http_response.url

    @property
    def method(self) -> str:
        return self.http_request.method

    @property
    def http_version(self) -> str:
        return self.http_response.http_version

    @property
    def elapsed(self) -> timedelta:
        return self.http_response.elapsed

    @property
    def is_closed(self) -> bool:
        return self.http_response.is_closed

    async def _parse(
        self,
        *,
        to: Optional[Type[U]] = None,
        decoder: Optional[Callable[[str], Any]] = None,
        sentinel: Optional[str] = None,
    ) -> T | U:
        mode = self._mode
        if mode == "event_stream" and not _can_honor_event_stream(
            to, self.http_response
        ):
            mode = "buffered"
        if to is not None and mode == "event_stream":
            return _build_event_stream(
                to=to,
                decoder=decoder,
                sentinel=sentinel,
                response=self.http_response,
                client_ref=self._client_ref,
                async_=True,
            )
        parse_key = to or _DEFAULT_PARSE_KEY
        if parse_key not in self._parsed_by_type:
            # raw_stream / event_stream wrappers leave the body unread so the
            # caller can iter_bytes/iter_lines; buffered mode eagerly reads.
            body = None
            try:
                if to is None:
                    if mode == "buffered":
                        await self.read()
                    self._parsed_by_type[parse_key] = await self._parser(
                        self.http_response
                    )
                else:
                    body = await self.text()
                    self._parsed_by_type[parse_key] = unmarshal_json_response(
                        to, self.http_response, body, validate=False
                    )
            except Exception as exc:
                await self._raise_parse_error(exc)
        return self._parsed_by_type[parse_key]

    async def read(self) -> bytes:
        return await self.http_response.aread()

    async def text(self) -> str:
        await self.read()
        return self.http_response.text

    async def json(self) -> Any:
        await self.read()
        return self.http_response.json()

    async def close(self) -> None:
        await self.http_response.aclose()

    async def __aenter__(self: AsyncResponseT) -> AsyncResponseT:
        return self

    async def __aexit__(
        self,
        exc_type: Optional[type[BaseException]],
        exc: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        await self.close()

    async def iter_bytes(
        self, chunk_size: Optional[int] = None
    ) -> AsyncIterator[bytes]:
        async for chunk in self.http_response.aiter_bytes(chunk_size):
            yield chunk

    async def iter_text(self, chunk_size: Optional[int] = None) -> AsyncIterator[str]:
        async for chunk in self.http_response.aiter_text(chunk_size):
            yield chunk

    async def iter_lines(self) -> AsyncIterator[str]:
        async for chunk in self.http_response.aiter_lines():
            yield chunk


class AsyncAPIResponse(_AsyncAPIResponseBase[T]):
    @overload
    async def parse(self) -> T: ...

    @overload
    async def parse(self, *, to: Type[U]) -> U: ...

    @overload
    async def parse(self, *, to: Type[U], decoder: Callable[[str], Any]) -> U: ...

    @overload
    async def parse(self, *, to: Type[U], sentinel: str) -> U: ...

    @overload
    async def parse(
        self,
        *,
        to: Type[U],
        decoder: Callable[[str], Any],
        sentinel: str,
    ) -> U: ...

    async def parse(
        self,
        *,
        to: Optional[Type[U]] = None,
        decoder: Optional[Callable[[str], Any]] = None,
        sentinel: Optional[str] = None,
    ) -> T | U:
        """Decode the response body into a typed value.

        Async sibling of ``APIResponse.parse``. Modes, decoder synthesis,
        and ``sentinel`` semantics follow the same contract. For SSE
        responses pass ``to=AsyncStream[T]`` (not ``Stream[T]``)
        and iterate with ``async for``.

        Errors
        ------
        * ``TypeError`` if ``to`` is supplied on an SSE response
          (``mode="event_stream"``) and is neither ``Stream[T]``
          nor ``AsyncStream[T]``.
        * ``TypeError`` if ``to`` is the right stream class but ``T`` is
          not a ``BaseModel`` and ``decoder`` was not supplied.
        * ``pydantic.ValidationError`` from the (synthesized or
          user-supplied) decoder when the wire shape does not match
          ``T``.
        """
        return await self._parse(to=to, decoder=decoder, sentinel=sentinel)


class AsyncStreamedAPIResponse(_AsyncAPIResponseBase[T]):
    async def parse(self) -> T:
        return await self._parse()


class ResponseContextManager(Generic[ResponseT]):
    def __init__(self, request_func: Callable[[], ResponseT]) -> None:
        self._request_func = request_func
        self._response: Optional[ResponseT] = None

    def __enter__(self) -> ResponseT:
        self._response = self._request_func()
        return self._response

    def __exit__(
        self,
        exc_type: Optional[type[BaseException]],
        exc: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        if self._response is not None:
            self._response.close()


class AsyncResponseContextManager(Generic[AsyncResponseT]):
    def __init__(self, request_func: Callable[[], Awaitable[AsyncResponseT]]) -> None:
        self._request_func = request_func
        self._response: Optional[AsyncResponseT] = None

    async def __aenter__(self) -> AsyncResponseT:
        self._response = await self._request_func()
        return self._response

    async def __aexit__(
        self,
        exc_type: Optional[type[BaseException]],
        exc: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        if self._response is not None:
            await self._response.close()


def _set_response_mode(kwargs: Dict[str, Any], mode: str, header_kwarg: str) -> None:
    headers = dict(kwargs.get(header_kwarg) or {})
    headers[_RESPONSE_MODE_HEADER] = mode
    kwargs[header_kwarg] = headers


def consume_response_mode(
    http_headers: Optional[Mapping[str, str]],
) -> Tuple[ResponseModeHeader, Optional[Mapping[str, str]]]:
    """Read and remove the marker header, returning the response mode and the
    remaining headers."""
    if http_headers is None:
        return _PARSED_RESPONSE_MODE, None
    mode: ResponseModeHeader = _PARSED_RESPONSE_MODE
    remaining: Dict[str, str] = {}
    for name, value in http_headers.items():
        if name.lower() == _RESPONSE_MODE_HEADER:
            if value in (_RAW_RESPONSE_MODE, _STREAMING_RESPONSE_MODE):
                mode = cast(ResponseModeHeader, value)
        else:
            remaining[name] = value
    return mode, (remaining or None)


def to_raw_response_wrapper(
    func: Callable[P, T], header_kwarg: str
) -> Callable[P, APIResponse[T]]:
    @functools.wraps(func)
    def wrapped(*args: P.args, **kwargs: P.kwargs) -> APIResponse[T]:
        _set_response_mode(
            cast(Dict[str, Any], kwargs), _RAW_RESPONSE_MODE, header_kwarg
        )
        return cast(APIResponse[T], func(*args, **kwargs))

    return wrapped


def async_to_raw_response_wrapper(
    func: Callable[P, Awaitable[T]], header_kwarg: str
) -> Callable[P, Awaitable[AsyncAPIResponse[T]]]:
    @functools.wraps(func)
    async def wrapped(*args: P.args, **kwargs: P.kwargs) -> AsyncAPIResponse[T]:
        _set_response_mode(
            cast(Dict[str, Any], kwargs), _RAW_RESPONSE_MODE, header_kwarg
        )
        return cast(AsyncAPIResponse[T], await func(*args, **kwargs))

    return wrapped


def to_streamed_response_wrapper(
    func: Callable[P, T], header_kwarg: str
) -> Callable[P, ResponseContextManager[StreamedAPIResponse[T]]]:
    @functools.wraps(func)
    def wrapped(
        *args: P.args, **kwargs: P.kwargs
    ) -> ResponseContextManager[StreamedAPIResponse[T]]:
        _set_response_mode(
            cast(Dict[str, Any], kwargs), _STREAMING_RESPONSE_MODE, header_kwarg
        )
        request_func = functools.partial(func, *args, **kwargs)
        return ResponseContextManager(
            cast(Callable[[], StreamedAPIResponse[T]], request_func)
        )

    return wrapped


def async_to_streamed_response_wrapper(
    func: Callable[P, Awaitable[T]], header_kwarg: str
) -> Callable[P, AsyncResponseContextManager[AsyncStreamedAPIResponse[T]]]:
    @functools.wraps(func)
    def wrapped(
        *args: P.args, **kwargs: P.kwargs
    ) -> AsyncResponseContextManager[AsyncStreamedAPIResponse[T]]:
        _set_response_mode(
            cast(Dict[str, Any], kwargs), _STREAMING_RESPONSE_MODE, header_kwarg
        )
        request_func = functools.partial(func, *args, **kwargs)
        return AsyncResponseContextManager(
            cast(Callable[[], Awaitable[AsyncStreamedAPIResponse[T]]], request_func)
        )

    return wrapped
