"""
This module provides a client for data collectors to publish data to a specified topic.
"""

from typing import TypeVar
import os
import logging
from .component_io import ChannelDefinition, ComponentIO
from .typing import (
    AudioPublishEnvelope,
    ChannelType,
    PublishEnvelope,
    TimestreamPublishEnvelope,
)

logging.basicConfig(level=os.environ.get("LOG_LEVEL", "ERROR"))
logger = logging.getLogger(__name__)

P = TypeVar("P", bound=ChannelType)


class DataCollectorClient(ComponentIO[None, PublishEnvelope[P]]):
    """Client for publishing data to a data collector service.

    This component is designed to be part of a larger pipeline system. It
    automatically constructs a unique NATS topic for publishing results based on the
    pipeline and module IDs inherited from the `ComponentIO` base class.

    This class is generic, allowing it to handle and publish any specified data type.

    Args:
        Generic (P): The type of data that the client will publish.

    Attributes:
        publish_topic (str): The NATS topic where data will be published, formatted as
            'data-collector/{component_id}/{component_version_id}/result'.
        publish_audio_topic (str): The NATS topic where audio data will be published, formatted as
            'data-collector/{component_id}/{component_version_id}/audio/result'.
    """

    channels: list[ChannelDefinition]

    def __init__(self) -> None:
        """Initializes the DataCollectorClient and sets the publish topic.

        Asynchronously publishes a data payload to the pre-configured topic.

        Args:
            data (T): The data to be published. The type is determined by the
                generic type `T` provided during the class instantiation.
        """
        super().__init__()

        self.channels = super()._get_environment_variable_channel_list("CHANNELS")

        self.publish_topic = f"{self.asset_id}/data-collector/{self.component_id}/{self.component_version_id}/result"
        self.subscribe_context_engine_topic = f"{self.asset_id}/context-engine/result"

    def publish_timestream(self, data: P) -> None:
        """Asynchronously publishes data to the predefined topic.

        This method is a convenience wrapper around the parent's `publish` method,
        sending the data to the topic specified in `self.publish_topic`.

        Args:
            data (T): The data payload to publish. The type is generic and
                should be compatible with the configured message broker's serializer.

        Raises:
            ValueError: If any keys in the data do not correspond to existing channels.
        """
        keys = list(data.keys())
        channel_names = [channel.name for channel in self.channels]
        missing_keys = set(channel_names) - set(keys)
        if missing_keys:
            logger.error("Missing channels for keys: %s", missing_keys)
            raise ValueError(f"Missing channels for keys: {missing_keys}")
        message: TimestreamPublishEnvelope[P] = {
            "type": "TIMESTREAM",
            "data": data,
        }
        self._publish(self.publish_topic, message)

    def publish_audio(self, data: P) -> None:
        """Asynchronously publishes audio to the predefined topic.

        This method is a convenience wrapper around the parent's `publish` method,
        sending the audio to the topic specified in `self.publish_audio_topic`.

        Args:
            data (T): The data payload to publish. The type is generic and
                should be compatible with the configured message broker's serializer.
            meta (AudioMeta): Metadata payload.

        Raises:
            ValueError: If any keys in the data do not correspond to existing channels.
        """
        keys = list(data.keys())
        channel_names = [channel.name for channel in self.channels]
        missing_keys = set(channel_names) - set(keys)
        if missing_keys:
            logger.error("Missing channels for keys: %s", missing_keys)
            raise ValueError(f"Missing channels for keys: {missing_keys}")
        message: AudioPublishEnvelope[P] = {
            "type": "AUDIO",
            "data": data,
        }
        self._publish(self.publish_topic, message)
