# TODO: this file contains too much implementation -> restructure.
import abc
import enum
import logging
import pathlib
import traceback
from typing import (  # type: ignore[attr-defined]  # Chains uses Python >=3.9.
    Any,
    AsyncIterator,
    Callable,
    ClassVar,
    Generic,
    GenericAlias,  # This causes above type error.
    Iterable,
    Literal,
    Mapping,
    Optional,
    Protocol,
    Type,
    TypeVar,
    Union,
    cast,
    get_args,
    get_origin,
    runtime_checkable,
)

import pydantic
from truss.base import truss_config
from truss.base.constants import PRODUCTION_ENVIRONMENT_NAME

from truss_chains import baseten_llm

BASETEN_API_SECRET_NAME = "baseten_chain_api_key"
SECRET_DUMMY = "***"
TRUSS_CONFIG_CHAINS_KEY = "chains_metadata"
GENERATED_CODE_DIR = ".chains_generated"
DYNAMIC_CHAINLET_CONFIG_KEY = "dynamic_chainlet_config"
OTEL_TRACE_PARENT_HEADER_KEY = "traceparent"
RUN_REMOTE_METHOD_NAME = "run_remote"  # Chainlet method name exposed as endpoint.
MODEL_ENDPOINT_METHOD_NAME = "predict"  # Model method name exposed as endpoint.
HEALTH_CHECK_METHOD_NAME = "is_healthy"
# Below arg names must correspond to `definitions.ABCChainlet`.
CONTEXT_ARG_NAME = "context"  # Referring to Chainlets `__init__` signature.
SELF_ARG_NAME = "self"
REMOTE_CONFIG_NAME = "remote_config"

K = TypeVar("K", contravariant=True)
V = TypeVar("V", covariant=True)


C = TypeVar("C")


class _classproperty(Generic[C, V]):
    def __init__(self, fget: Callable[[Type[C]], V]) -> None:
        self._fget = fget

    def __get__(self, instance: object, owner: Type[C]) -> V:
        return self._fget.__get__(None, owner)()


def classproperty(fget: Callable[[Type[C]], V]) -> _classproperty[C, V]:
    return _classproperty(fget)


@runtime_checkable
class MappingNoIter(Protocol[K, V]):
    def __getitem__(self, key: K) -> V: ...

    def __len__(self) -> int: ...

    def __contains__(self, key: K) -> bool: ...


class SafeModel(pydantic.BaseModel):
    """Pydantic base model with reasonable config."""

    model_config = pydantic.ConfigDict(
        arbitrary_types_allowed=False, strict=True, validate_assignment=True
    )


class SafeModelNonSerializable(pydantic.BaseModel):
    """Pydantic base model with reasonable config - allowing arbitrary types."""

    model_config = pydantic.ConfigDict(
        arbitrary_types_allowed=True,
        strict=True,
        validate_assignment=True,
        extra="forbid",
    )


class ChainsUsageError(TypeError):
    """Raised when user-defined Chainlets do not adhere to API constraints."""


class MissingDependencyError(TypeError):
    """Raised when a needed resource could not be found or is not defined."""


class ChainsRuntimeError(Exception):
    """Raised when components are not used the expected way at runtime."""


class ChainsDeploymentError(Exception):
    """Raised when interaction with a Chain deployment are not possible."""


class AbsPath:
    _abs_file_path: str
    _creating_module: str
    _original_path: str

    def __init__(
        self, abs_file_path: str, creating_module: str, original_path: str
    ) -> None:
        self._abs_file_path = abs_file_path
        self._creating_module = creating_module
        self._original_path = original_path

    def _raise_if_not_exists(self, abs_path: str) -> None:
        path = pathlib.Path(abs_path)
        if not (path.is_file() or (path.is_dir() and any(path.iterdir()))):
            raise MissingDependencyError(
                f"With the file path `{self._original_path}` an absolute path relative "
                f"to the calling module `{self._creating_module}` was created, "
                f"resulting `{self._abs_file_path}` - but no file was found."
            )

    @property
    def abs_path(self) -> str:
        if self._abs_file_path != self._original_path:
            logging.debug(
                f"Using abs path `{self._abs_file_path}` for path specified as "
                f"`{self._original_path}` (in `{self._creating_module}`)."
            )
        abs_path = self._abs_file_path
        self._raise_if_not_exists(abs_path)
        return abs_path


class BasetenImage(enum.Enum):
    """Default images, curated by baseten, for different python versions. If a Chainlet
    uses GPUs, drivers will be included in the image."""

    # Enum values correspond to truss canonical python versions.
    PY39 = "py39"
    PY310 = "py310"
    PY311 = "py311"


class CustomImage(SafeModel):
    """Configures the usage of a custom image hosted on dockerhub."""

    image: str
    python_executable_path: Optional[str] = None
    docker_auth: Optional[truss_config.DockerAuthSettings] = None


class DockerImage(SafeModelNonSerializable):
    """Configures the docker image in which a remoted chainlet is deployed.

    Note:
        Any paths are relative to the source file where ``DockerImage`` is
        defined and must be created with the helper function ``make_abs_path_here``.
        This allows you for example organize chainlets in different (potentially nested)
        modules and keep their requirement files right next their python source files.

    Args:
        base_image: The base image used by the chainlet. Other dependencies and
          assets are included as additional layers on top of that image. You can choose
          a Baseten default image for a supported python version (e.g.
          ``BasetenImage.PY311``), this will also include GPU drivers if needed, or
          provide a custom image (e.g. ``CustomImage(image="python:3.11-slim")``)..
        pip_requirements_file: Path to a file containing pip requirements. The file
          content is naively concatenated with ``pip_requirements``.
        pip_requirements: A list of pip requirements to install.  The items are
          naively concatenated with the content of the ``pip_requirements_file``.
        apt_requirements: A list of apt requirements to install.
        data_dir: Data from this directory is copied into the docker image and
          accessible to the remote chainlet at runtime.
        external_package_dirs: A list of directories containing additional python
          packages outside the chain's workspace dir, e.g. a shared library. This code
          is copied into the docker image and importable at runtime.
    """

    base_image: Union[BasetenImage, CustomImage] = BasetenImage.PY311
    pip_requirements_file: Optional[AbsPath] = None
    pip_requirements: list[str] = []
    apt_requirements: list[str] = []
    data_dir: Optional[AbsPath] = None
    external_package_dirs: Optional[list[AbsPath]] = None

    @pydantic.root_validator(pre=True)
    def migrate_fields(cls, values):
        if "base_image" in values:
            base_image = values["base_image"]
            if isinstance(base_image, str):
                doc_link = "https://docs.baseten.co/chains-reference/sdk#class-truss-chains-dockerimage"
                raise ChainsUsageError(
                    "`DockerImage.base_image` as string is deprecated. Specify as "
                    f"`BasetenImage` or `CustomImage` (see docs: {doc_link})."
                )
        return values


class ComputeSpec(pydantic.BaseModel):
    """Parsed and validated compute.  See ``Compute`` for more information."""

    # TODO[rcano] add node count
    cpu_count: int = 1
    predict_concurrency: int = 1
    memory: str = "2Gi"
    accelerator: truss_config.AcceleratorSpec = truss_config.AcceleratorSpec()


CpuCountT = Literal["cpu_count"]
CPU_COUNT: CpuCountT = "cpu_count"


class Compute:
    """Specifies which compute resources a chainlet has in the *remote* deployment.

    Note:
        Not all combinations can be exactly satisfied by available hardware, in some
        cases more powerful machine types are chosen to make sure requirements are met
        or over-provisioned. Refer to the
        `baseten instance reference <https://docs.baseten.co/performance/instances>`_.
    """

    # Builder to create ComputeSpec.
    # This extra layer around `ComputeSpec` is needed to parse the accelerator options.

    _spec: ComputeSpec

    def __init__(
        self,
        cpu_count: int = 1,
        memory: str = "2Gi",
        gpu: Union[str, truss_config.Accelerator, None] = None,
        gpu_count: int = 1,
        predict_concurrency: Union[int, CpuCountT] = 1,
    ) -> None:
        """
        Args:
            cpu_count: Minimum number of CPUs to allocate.
            memory: Minimum memory to allocate, e.g. "2Gi" (2 gibibytes).
            gpu: GPU accelerator type, e.g. "A10G", "A100", refer to the
              `truss config <https://docs.baseten.co/reference/config#resources-accelerator>`_
              for more choices.
            gpu_count: Number of GPUs to allocate.
            predict_concurrency: Number of concurrent requests a single replica of a
              deployed chainlet handles.

        Concurrency concepts are explained in `this guide <https://docs.baseten.co/deploy/guides/concurrency#predict-concurrency>`_. # noqa: E501
        It is important to understand the difference between `predict_concurrency` and
        the concurrency target (used for autoscaling, i.e. adding or removing replicas).
        Furthermore, the ``predict_concurrency`` of a single instance is implemented in
        two ways:

        - Via python's ``asyncio``, if ``run_remote`` is an async def. This
          requires that ``run_remote`` yields to the event loop.

        - With a threadpool if it's a synchronous function. This requires
          that the threads don't have significant CPU load (due to the GIL).
        """
        accelerator = truss_config.AcceleratorSpec()
        if gpu:
            accelerator.accelerator = truss_config.Accelerator(gpu)
            accelerator.count = gpu_count
            accelerator = truss_config.AcceleratorSpec(
                accelerator=truss_config.Accelerator(gpu), count=gpu_count
            )
        if predict_concurrency == CPU_COUNT:
            predict_concurrency_int = cpu_count
        else:
            assert isinstance(predict_concurrency, int)
            predict_concurrency_int = predict_concurrency

        self._spec = ComputeSpec(
            cpu_count=cpu_count,
            memory=memory,
            accelerator=accelerator,
            predict_concurrency=predict_concurrency_int,
        )

    def get_spec(self) -> ComputeSpec:
        return self._spec.model_copy(deep=True)


class AssetSpec(SafeModel):
    """Parsed and validated assets. See ``Assets`` for more information."""

    secrets: dict[str, str] = pydantic.Field({})
    cached: list[truss_config.ModelRepo] = []
    external_data: list[truss_config.ExternalDataItem] = []


class Assets:
    """Specifies which assets a chainlet can access in the remote deployment.

    For example, model weight caching can be used like this::

        import truss_chains as chains
        from truss.base import truss_config

        mistral_cache = truss_config.ModelRepo(
            repo_id="mistralai/Mistral-7B-Instruct-v0.2",
            allow_patterns=["*.json", "*.safetensors", ".model"]
        )
        chains.Assets(cached=[mistral_cache], ...)

    See `truss caching guide <https://docs.baseten.co/deploy/guides/model-cache#enabling-caching-for-a-model>`_
    for more details on caching.
    """

    # Builder to create asset spec.
    # This extra layer around `AssetSpec` is needed to add secret_keys.
    _spec: AssetSpec

    def __init__(
        self,
        cached: Iterable[truss_config.ModelRepo] = (),
        secret_keys: Iterable[str] = (),
        external_data: Iterable[truss_config.ExternalDataItem] = (),
    ) -> None:
        """
        Args:
            cached: One or more ``truss_config.ModelRepo`` objects.
            secret_keys: Names of secrets stored on baseten, that the
              chainlet should have access to. You can manage secrets on baseten
              `here <https://app.baseten.co/settings/secrets>`_.
            external_data: Data to be downloaded from public URLs and made available
              in the deployment (via ``context.data_dir``). See
              `here <https://docs.baseten.co/reference/config#external-data>`_ for
              more details.
        """
        self._spec = AssetSpec(
            cached=list(cached),
            secrets={k: SECRET_DUMMY for k in secret_keys},
            external_data=list(external_data),
        )

    def get_spec(self) -> AssetSpec:
        """Returns parsed and validated assets."""
        return self._spec.model_copy(deep=True)


class ChainletOptions(SafeModelNonSerializable):
    """
    Args:
        enable_b10_tracing: enables baseten-internal trace data collection. This
          helps baseten engineers better analyze chain performance in case of issues.
          It is independent of a potentially user-configured tracing instrumentation.
          Turning this on, could add performance overhead.
        env_variables: static environment variables available to the deployed chainlet.
        health_checks: Configures health checks for the chainlet.
    """

    enable_b10_tracing: bool = False
    env_variables: Mapping[str, str] = {}
    health_checks: truss_config.HealthChecks = truss_config.HealthChecks()


class ChainletMetadata(SafeModelNonSerializable):
    is_entrypoint: bool = False
    chain_name: Optional[str] = None
    init_is_patched: bool = False


class EntityType(str, enum.Enum):
    CHAINLET = enum.auto()
    MODEL = enum.auto()
    BASETEN_LLM = enum.auto()


class FrameworkConfig(SafeModelNonSerializable):
    entity_type: EntityType
    supports_dependencies: bool
    endpoint_method_name: str


class RemoteConfig(SafeModelNonSerializable):
    """Bundles config values needed to deploy a chainlet remotely.

    This is specified as a class variable for each chainlet class, e.g.::

            import truss_chains as chains


            class MyChainlet(chains.ChainletBase):
                remote_config = chains.RemoteConfig(
                    docker_image=chains.DockerImage(
                        pip_requirements=["torch==2.0.1", ...]
                    ),
                    compute=chains.Compute(cpu_count=2, gpu="A10G", ...),
                    assets=chains.Assets(secret_keys=["hf_access_token"], ...),
                )

    """

    docker_image: DockerImage = DockerImage()
    compute: Compute = Compute()
    assets: Assets = Assets()
    name: Optional[str] = None
    options: ChainletOptions = ChainletOptions()

    def get_compute_spec(self) -> ComputeSpec:
        return self.compute.get_spec()

    def get_asset_spec(self) -> AssetSpec:
        return self.assets.get_spec()


DEFAULT_TIMEOUT_SEC = 600.0


class RPCOptions(SafeModel):
    """Options to customize RPCs to dependency chainlets.

    Args:
        retries: The number of times to retry the remote chainlet in case of failures
          (e.g. due to transient network issues). For streaming, retries are only made
          if the request fails before streaming any results back. Failures mid-stream
          not retried.
        timeout_sec: Timeout for the HTTP request to this chainlet.
        use_binary: Whether to send data in binary format. This can give a parsing
         speedup and message size reduction (~25%) for numpy arrays. Use
         ``NumpyArrayField`` as a field type on pydantic models for integration and set
         this option to ``True``. For simple text data, there is no significant benefit.
    """

    retries: int = 1
    timeout_sec: float = DEFAULT_TIMEOUT_SEC
    use_binary: bool = False


class ServiceDescriptor(SafeModel):
    """Bundles values to establish an RPC session to a dependency chainlet,
    specifically with ``StubBase``."""

    name: str
    display_name: str
    options: RPCOptions


class DeployedServiceDescriptor(ServiceDescriptor):
    predict_url: str


class Environment(SafeModel):
    """The environment the chainlet is deployed in.

    Args:
        name: The name of the environment.
    """

    name: str
    # can add more fields here as we add them to dynamic_config configmap


class DeploymentContext(SafeModelNonSerializable):
    """Bundles config values and resources needed to instantiate Chainlets.

    The context can optionally added as a trailing argument in a Chainlet's
    ``__init__`` method and then used to set up the chainlet (e.g. using a secret as
    an access token for downloading model weights).

    Args:
        data_dir: The directory where the chainlet can store and access data,
          e.g. for downloading model weights.
        chainlet_to_service: A mapping from chainlet names to service descriptors.
          This is used to create RPC sessions to dependency chainlets. It contains only
          the chainlet services that are dependencies of the current chainlet.
        secrets: A mapping from secret names to secret values. It contains only the
          secrets that are listed in ``remote_config.assets.secret_keys`` of the
          current chainlet.
        environment: The environment that the chainlet is deployed in.
          None if the chainlet is not associated with an environment.
    """

    data_dir: Optional[pathlib.Path] = None
    chainlet_to_service: Mapping[str, DeployedServiceDescriptor]
    secrets: MappingNoIter[str, str]
    environment: Optional[Environment] = None

    def get_service_descriptor(self, chainlet_name: str) -> DeployedServiceDescriptor:
        if chainlet_name not in self.chainlet_to_service:
            raise MissingDependencyError(f"{chainlet_name}")
        return self.chainlet_to_service[chainlet_name]

    def get_baseten_api_key(self) -> str:
        if self.secrets is None:
            raise ChainsRuntimeError(
                f"Secrets not set in `{self.__class__.__name__}` object."
            )
        error_msg = (
            "For using chains, it is required to setup a an API key with name "
            f"`{BASETEN_API_SECRET_NAME}` on Baseten to allow chain Chainlet to "
            "call other Chainlets. For local execution, secrets can be provided "
            "to `run_local`."
        )
        if BASETEN_API_SECRET_NAME not in self.secrets:
            raise MissingDependencyError(error_msg)

        api_key = self.secrets[BASETEN_API_SECRET_NAME]
        if api_key == SECRET_DUMMY:
            raise MissingDependencyError(
                f"{error_msg}. Retrieved dummy value of `{api_key}`."
            )
        return api_key


class TrussMetadata(SafeModel):
    """Plugin for the truss config (in config["model_metadata"]["chains_metadata"])."""

    chainlet_to_service: Mapping[str, ServiceDescriptor]


class ABCChainlet(abc.ABC):
    remote_config: ClassVar[RemoteConfig] = RemoteConfig()
    # `meta_data` is not shared between subclasses, each has an isolated copy.
    meta_data: ClassVar[ChainletMetadata] = ChainletMetadata()
    _framework_config: ClassVar[FrameworkConfig]

    @classmethod
    def has_custom_init(cls) -> bool:
        return cls.__init__ is not object.__init__

    @classproperty
    @classmethod
    def name(cls) -> str:
        return cls.__name__

    @classproperty
    @classmethod
    def display_name(cls) -> str:
        return cls.remote_config.name or cls.name

    @classproperty
    @classmethod
    def supports_dependencies(cls) -> bool:
        return cls._framework_config.supports_dependencies

    @classproperty
    @classmethod
    def entity_type(cls) -> EntityType:
        return cls._framework_config.entity_type

    @classproperty
    @classmethod
    def endpoint_method_name(cls) -> str:
        return cls._framework_config.endpoint_method_name

    # Cannot add this abstract method to API, because we want to allow arbitrary
    # arg/kwarg names and specifying any function signature here would give type errors
    # @abc.abstractmethod
    # def run_remote(self, *args, **kwargs) -> Any:
    #     ...


class TypeDescriptor(SafeModelNonSerializable):
    """For describing I/O types of Chainlets."""

    raw: Any  # The raw type annotation object (could be a type or GenericAlias).

    @property
    def is_pydantic(self) -> bool:
        return (
            isinstance(self.raw, type)
            and not isinstance(self.raw, GenericAlias)
            and issubclass(self.raw, pydantic.BaseModel)
        )

    @property
    def has_pydantic_args(self):
        origin = get_origin(self.raw)
        if not origin:
            return False
        args = get_args(self.raw)
        return any(
            isinstance(arg, type) and issubclass(arg, pydantic.BaseModel)
            for arg in args
        )

    @property
    def is_websocket(self) -> bool:
        return self.raw == WebSocketProtocol


class StreamingTypeDescriptor(TypeDescriptor):
    origin_type: type
    arg_type: type

    @property
    def is_string(self) -> bool:
        return self.arg_type is str

    @property
    def is_pydantic(self) -> bool:
        return False


class InputArg(SafeModelNonSerializable):
    name: str
    type: TypeDescriptor
    is_optional: bool


class EndpointAPIDescriptor(SafeModelNonSerializable):
    name: str = RUN_REMOTE_METHOD_NAME
    input_args: list[InputArg]
    output_types: list[TypeDescriptor]
    is_async: bool
    is_streaming: bool

    @property
    def streaming_type(self) -> StreamingTypeDescriptor:
        if (
            not self.is_streaming
            or len(self.output_types) != 1
            or not isinstance(self.output_types[0], StreamingTypeDescriptor)
        ):
            raise ValueError(f"{self} is not a streaming endpoint.")
        return cast(StreamingTypeDescriptor, self.output_types[0])

    @property
    def is_websocket(self):
        return any(arg.type.is_websocket for arg in self.input_args)

    @property
    def has_pydantic_input(self) -> bool:
        return not self.is_websocket

    @property
    def has_pydantic_output(self) -> bool:
        return not (self.is_streaming or self.is_websocket)

    @property
    def is_baseten_llm(self) -> bool:
        return any(arg.type.raw == baseten_llm.ModelInput for arg in self.input_args)


class DependencyDescriptor(SafeModelNonSerializable):
    chainlet_cls: Type[ABCChainlet]
    options: RPCOptions

    @property
    def name(self) -> str:
        return self.chainlet_cls.name

    @property
    def display_name(self) -> str:
        return self.chainlet_cls.display_name


class HealthCheckAPIDescriptor(SafeModelNonSerializable):
    name: str = HEALTH_CHECK_METHOD_NAME
    is_async: bool


class ChainletAPIDescriptor(SafeModelNonSerializable):
    chainlet_cls: Type[ABCChainlet]
    src_path: str
    has_context: bool
    dependencies: Mapping[str, DependencyDescriptor]
    endpoint: EndpointAPIDescriptor
    health_check: Optional[HealthCheckAPIDescriptor]

    def __hash__(self) -> int:
        return hash(self.chainlet_cls)

    @property
    def name(self) -> str:
        return self.chainlet_cls.name

    @property
    def display_name(self) -> str:
        return self.chainlet_cls.display_name


class StackFrame(SafeModel):
    filename: str
    lineno: Optional[int]
    name: str
    line: Optional[str]

    @classmethod
    def from_frame_summary(cls, frame: traceback.FrameSummary):
        return cls(
            filename=frame.filename,
            lineno=frame.lineno,
            name=frame.name,
            line=frame.line,
        )

    def to_frame_summary(self) -> traceback.FrameSummary:
        return traceback.FrameSummary(
            filename=self.filename, lineno=self.lineno, name=self.name, line=self.line
        )


class RemoteErrorDetail(SafeModel):
    """When a remote chainlet raises an exception, this pydantic model contains
    information about the error and stack trace and is included in JSON form in the
    error response.
    """

    exception_cls_name: str
    exception_module_name: Optional[str]
    exception_message: str
    user_stack_trace: list[StackFrame]

    def _to_stack_summary(self) -> traceback.StackSummary:
        return traceback.StackSummary.from_list(
            frame.to_frame_summary() for frame in self.user_stack_trace
        )

    def format(self) -> str:
        """Format the error for printing, similar to how Python formats exceptions
        with stack traces."""
        stack = "".join(traceback.format_list(self._to_stack_summary()))
        exc_info = (
            f"\n(Exception class defined in `{self.exception_module_name}`.)"
            if self.exception_module_name
            else ""
        )
        error = (
            f"Chainlet-Traceback (most recent call last):\n"
            f"{stack}{self.exception_cls_name}: {self.exception_message}{exc_info}"
        )
        return error


class GenericRemoteException(Exception): ...


########################################################################################


class PushOptions(SafeModelNonSerializable):
    chain_name: str
    only_generate_trusses: bool = False


class PushOptionsBaseten(PushOptions):
    remote: str
    publish: bool
    environment: Optional[str]

    @classmethod
    def create(
        cls,
        chain_name: str,
        publish: bool,
        promote: Optional[bool],
        only_generate_trusses: bool,
        remote: str,
        environment: Optional[str] = None,
    ) -> "PushOptionsBaseten":
        if promote and not environment:
            environment = PRODUCTION_ENVIRONMENT_NAME
        if environment:
            publish = True
        return PushOptionsBaseten(
            remote=remote,
            chain_name=chain_name,
            publish=publish,
            only_generate_trusses=only_generate_trusses,
            environment=environment,
        )


class PushOptionsLocalDocker(PushOptions):
    # Local docker-to-docker requests don't need auth, but we need to set a
    # value different from `SECRET_DUMMY` to not trigger the check that the secret
    # is unset. Additionally, if local docker containers make calls to models deployed
    # on baseten, a real API key must be provided (i.e. the default must be overridden).
    baseten_chain_api_key: str = "docker_dummy_key"
    # If enabled, chains code is copied from the local package into `/app/truss_chains`
    # in the docker image (which takes precedence over potential pip/site-packages).
    # This should be used for integration tests or quick local dev loops.
    use_local_chains_src: bool = False


class WebSocketProtocol(Protocol):
    """Describes subset of starlette/fastAPIs websocket interface that we expose."""

    headers: Mapping[str, str]

    async def close(self, code: int = 1000, reason: Optional[str] = None) -> None: ...

    async def receive_text(self) -> str: ...
    async def receive_bytes(self) -> bytes: ...
    async def receive_json(self) -> Any: ...

    async def send_text(self, data: str) -> None: ...
    async def send_bytes(self, data: bytes) -> None: ...
    async def send_json(self, data: Any) -> None: ...

    def iter_text(self) -> AsyncIterator[str]: ...
    def iter_bytes(self) -> AsyncIterator[bytes]: ...
    def iter_json(self) -> AsyncIterator[Any]: ...
