"""Wrapper for async getstream feed class."""

import logging
from typing import TYPE_CHECKING, Collection

from ddtrace import tracer
from stream.feed import AsyncFeed as BaseFeed

from .. import config
from .base import BaseExtendedFeed


if TYPE_CHECKING:
    from ..client.base import Activity


__all__ = ["AsyncFeed"]


logger = logging.getLogger(__name__)


class AsyncFeed(BaseExtendedFeed, BaseFeed):
    """Extended Feed class for Stream (getstream.io) Client."""

    @tracer.wrap("stream.feed.add_activity", service=config.SERVICE_NAME)
    async def add_activity(self, activity_data: "Activity"):
        """Override to add trace context and correlation id to activity."""
        self.client.inject_trace_data(activity_data)
        return await super().add_activity(activity_data)

    @tracer.wrap("stream.feed.add_activities", service=config.SERVICE_NAME)
    async def add_activities(self, activity_list: Collection["Activity"]):
        """Override to add trace context and correlation id to activity."""
        self.client.inject_trace_data(activity_list)
        return await super().add_activities(activity_list)