"""Wrapper for async getstream client class."""

import logging
from typing import TYPE_CHECKING, Collection, Type

from ddtrace import tracer
from stream.client import AsyncStreamClient as BaseStreamClient

from .. import config
from ..feed import AsyncFeed
from ..utils import datadog, retry
from .base import Activity, BaseExtendedClient


if TYPE_CHECKING:
    from aiohttp import ClientResponse
    from stream.feed.base import BaseFeed

__all__ = ["AsyncStreamClient"]


logger = logging.getLogger(__name__)


class AsyncStreamClient(BaseExtendedClient, BaseStreamClient):
    """Extended Stream (getstream.io) Client.

    It does several things:
    1. Save ratelimit info from response headers
    2. Add DataDog traces to GetStream API calls
    3. Add trace context to stream messages
    4. Add correlation id to stream messages
    5. Add throttling mechanism to prevent API rate limiting from being reached or exceeded.
    6. Add retry mechanism to API calls in case of API rate limit been reached.
    """

    def _get_feed_cls(self) -> Type["BaseFeed"]:
        return AsyncFeed

    @tracer.wrap("stream.add_to_many", service=config.SERVICE_NAME)
    async def add_to_many(self, activity: Activity, feeds: Collection[str]):
        """Override to add trace context to activity."""
        self.inject_trace_data(activity)

        return await super().add_to_many(activity, feeds)

    @retry.on_ratelimit_reached
    async def _make_request(
        self,
        method,
        relative_url,
        signature,
        service_name="api",
        params=None,
        data=None,
    ):
        """Override to add retry and throttling functionality."""
        await self._throttler.async_delay(self)
        return await super()._make_request(method, relative_url, signature, service_name, params, data)

    async def _parse_response(self, response: "ClientResponse"):
        # Save last call ratelimit values
        self._ratelimit_info.set_from_headers(response.headers)

        request_info = response.request_info
        datadog.set_tags(response.status, str(request_info.url), request_info.method, self._ratelimit_info.to_dict())

        return await super()._parse_response(response)