"""Defines the client for a pre-processor component in a data processing pipeline.

This module contains the `PreProcessorClient` class, which handles the communication
(publishing and subscribing) for a pre-processing stage. It is designed to
receive data from a data collector and publish the processed results for
subsequent components in the pipeline.
"""

import logging
import os
from typing import Any, TypeVar, Callable, cast
from .component_io import ComponentIO
from .typing import ChannelType, FeatureType

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


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


class PreProcessorClient(ComponentIO[S, P]):
    """A client for a pre-processing component in a data pipeline.

    This class facilitates communication for a pre-processor. It subscribes to data
    from a 'data-collector' component and publishes the processed results. The specific
    data type it handles is defined by the generic type `T`.

    It inherits from `ComponentIO` to leverage common pipeline communication logic.

    Attributes:
        publish_topic (str): The NATS topic for publishing processed data.
        subscribe_topic (str): The NATS topic for subscribing to raw data.
    """

    def __init__(self) -> None:
        """Initializes the PreProcessorClient instance.

        Sets up the NATS topics for publishing and subscribing based on the
        pipeline and module IDs inherited from the `ComponentIO` base class.
        """
        super().__init__()

        self.channels = super()._get_environment_variable_channel_list("CHANNELS")
        self.features = super()._get_environment_variable_feature_list("FEATURES")
        self.pipeline_id = super()._get_environment_variable("PIPELINE_ID")
        self.data_collector_component_id = super()._get_environment_variable(
            "DATA_COLLECTOR_COMPONENT_ID"
        )
        self.data_collector_component_version_id = super()._get_environment_variable(
            "DATA_COLLECTOR_COMPONENT_VERSION_ID"
        )

        self.publish_topic = f"{self.asset_id}/pre-processor/{self.pipeline_id}/{self.component_id}/{self.component_version_id}/result"
        self.subscribe_topic = f"{self.asset_id}/data-collector/{self.data_collector_component_id}/{self.data_collector_component_version_id}/result"

    def publish_data(self, data: P) -> None:
        """Asynchronously publishes processed data.
        Args:
            data (P): The processed data to be published to the pre-processor's
                     result topic.
        """
        keys = data.keys()
        missing_keys = {feature.name for feature in self.features} - set(keys)
        if missing_keys:
            logger.error("Missing features for keys: %s", missing_keys)
            raise ValueError(f"Missing features for keys: {missing_keys}")

        for feature in self.features:
            shape = self._shape_from_object(data[feature.name])
            if list(shape) != feature.shape:
                logger.error(
                    "Feature '%s' has incorrect shape. Expected %s, got %s",
                    feature.name,
                    feature.shape,
                    list(shape),
                )
                error_message = (
                    f"Feature '{feature.name}' has incorrect shape. "
                    f"Expected {feature.shape}, got {list(shape)}"
                )
                raise ValueError(error_message)
            self._validate_list_type(data[feature.name], (float, int))
        self._publish(self.publish_topic, data)

    def subscribe(self, handler: Callable[[S], None]) -> Callable[[], None]:
        """Subscribes to raw data from a data collector.

        Sets up a subscription to the corresponding data collector's result topic.
        When a message is received, the provided handler is called with the data.

        Args:
            handler (Callable[[S], None]): A callback function to be invoked
                with the received data.

        Returns:
            Callable[[], None]: A function that can be called to unsubscribe
                from the topic.
        """
        return self._subscribe(self.subscribe_topic, handler)

    def subscribe_timestream(self, handler: Callable[[S], None]) -> Callable[[], None]:
        """Subscribes to timestream topic and returns only the payload data object.

        Expects messages in envelope format:
        {"type": "TIMESTREAM", "data": {...}}.
        """

        def _handle_message(message: ChannelType) -> None:
            payload = self._extract_data_payload(message, expected_type="TIMESTREAM")
            if payload is not None:
                handler(cast(S, payload))

        return self._subscribe(self.subscribe_topic, _handle_message)

    def subscribe_audio(self, handler: Callable[[S], None]) -> Callable[[], None]:
        """Subscribes to data collector topic and returns only audio payload data.

        Expects messages in envelope format:
        {"type": "AUDIO", "data": {...}}.
        """

        def _handle_message(message: ChannelType) -> None:
            payload = self._extract_data_payload(message, expected_type="AUDIO")
            if payload is not None:
                handler(cast(S, payload))

        return self._subscribe(self.subscribe_topic, _handle_message)

    def _extract_data_payload(
        self, message: ChannelType, expected_type: str
    ) -> ChannelType | None:
        message_type = message.get("type")
        if message_type != expected_type:
            return None
        if "data" not in message:
            raise ValueError("Expected envelope with 'data' field")
        payload = message["data"]
        if not isinstance(payload, dict):
            raise ValueError("Expected envelope 'data' field to be an object")
        return payload

    def _shape_from_object(self, obj: list) -> tuple:

        shape = []

        def _shape_from_object_r(element: list, axis: int):
            try:
                for i, e in enumerate(element):
                    _shape_from_object_r(e, axis + 1)
                    while len(shape) <= axis:
                        shape.append(0)
                    l = i + 1
                    s = shape[axis]
                    if l > s:
                        shape[axis] = l
            except TypeError:
                pass

        _shape_from_object_r(obj, 0)
        return tuple(shape)

    def _validate_list_type(self, obj: Any, expected_type: tuple[type, ...]) -> None:
        if not isinstance(obj, list):
            raise TypeError(f"Expected list, got {type(obj)}")
        for item in obj:
            if isinstance(item, list):
                self._validate_list_type(item, expected_type)
            else:
                if not isinstance(item, expected_type):
                    raise TypeError(
                        f"Expected list of {expected_type}, got {type(item)}"
                    )
