"""
Provides a client for inter-component communication within AWS Greengrass.

This module defines the `ComponentIO` class, a generic client that simplifies
publishing and subscribing to topics using the AWS IoT Greengrass Core IPC service.
It is designed to be used by Greengrass components to communicate with each other.

The client requires the `DEPLOYMENT_ID`, `COMPONENT_ID` and `COMPONENT_VERSION_ID` environment variables
to be set, which are typically provided by the Greengrass environment.

Example:
    ```python

    # Define the type of data you expect to send/receive
    io = ComponentIO[dict]()
    io.connect()

    def on_message(message: dict):
        print(f"Received message: {message}")

    async def main():
        # Subscribe to a topic
        unsubscribe = io._subscribe("my/topic", on_message)

        # Publish a message
        await io._publish("my/topic", {"key": "value"})

        # Wait for a bit to receive the message
        await asyncio.sleep(2)

        # Clean up the subscription
        unsubscribe()

    if __name__ == "__main__":
        asyncio.run(main())
    ```

"""

from dataclasses import dataclass
from typing import Any, Generic, Mapping, TypeVar, Optional, Callable

import os
import sys
import traceback
import json
import logging
from enum import Enum
from awsiot.greengrasscoreipc.clientv2 import GreengrassCoreIPCClientV2
from awsiot.greengrasscoreipc.model import (
    BinaryMessage,
    PublishMessage,
    JsonMessage,
)

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


class DataType(Enum):
    """Enumeration of supported data types."""

    FLOAT32 = "FLOAT32"
    INT32 = "INT32"
    STRING = "STRING"
    BOOLEAN = "BOOLEAN"
    BYTES = "BYTES"


@dataclass
class ContextDefinition:
    """Represents a data context with a name and index."""

    name: str
    index: int


@dataclass
class ChannelDefinition:
    """Represents a data channel with a name and index."""

    name: str
    index: int


@dataclass
class FeatureDefinition:
    """Represents a feature with a name and index."""

    name: str
    index: int
    # pylint: disable=invalid-name
    dataType: DataType
    shape: list[int]
    description: Optional[str] = None


S = TypeVar("S", bound=Mapping[str, str] | None)
P = TypeVar("P", bound=Mapping[str, Any] | None)


class ComponentIO(Generic[S, P]):  # pylint: disable=too-few-public-methods
    """
     A client for communication between Greengrass components via IPC.

    This class wraps the AWS Greengrass Core IPC client to provide a simplified
    interface for publishing and subscribing to topics. It automatically handles
    message serialization (JSON) and deserialization.

    The `DEPLOYMENT_ID`, `COMPONENT_ID` and `COMPONENT_VERSION_ID` environment variables must be set before
    instantiating this class.

    The `connect()` method must be called before any publish or subscribe
    operations.

    This is a generic class, and the type variable `T` represents the expected
    type of the message payload.

    This is a generic class with two type variables:
    - S: Represents the expected type of incoming message payloads when subscribing to topics
    - P: Represents the type of outgoing message payloads when publishing to topics

    Attributes:
        deployment_id (str): The ID of the deployment, read from the
                           `DEPLOYMENT_ID` environment variable.
        component_id (str): The ID of the component, read from the
                            `COMPONENT_ID` environment variable.
        component_version_id (str): The ID of the component, read from the
                         `COMPONENT_VERSION_ID` environment variable.
        asset_id (str): The ID of the asset, read from the
                        `ASSET_ID` environment variable.

        Initializes the ComponentIO client.

        Reads deployment and component version IDs from environment variables.

        Raises:
            ValueError: If `DEPLOYMENT_ID`, `COMPONENT_ID`, and `COMPONENT_VERSION_ID` environment
                        variables are not set.
    """

    _ipc: Optional[GreengrassCoreIPCClientV2]
    deployment_id: str
    component_id: str
    component_version_id: str

    def __init__(self) -> None:
        """Initializes the ComponentIO client."""
        self.deployment_id = self._get_environment_variable("DEPLOYMENT_ID")
        self.component_id = self._get_environment_variable("COMPONENT_ID")
        self.component_version_id = self._get_environment_variable(
            "COMPONENT_VERSION_ID"
        )
        self.asset_id = self._get_environment_variable("ASSET_ID")
        self._ipc = None

    def connect(self) -> None:
        """Connects to the Greengrass Core IPC service.

        This method initializes the connection to the underlying message broker.
        It must be called before using `_publish` or `_subscribe`. If already
        connected, this method does nothing.
        """
        if self._ipc is None:
            self._ipc = GreengrassCoreIPCClientV2()

    def _publish(self, topic: str, message: P) -> None:
        """Publishes a message to a specified topic.

        The message is serialized based on its type. If it's `bytes` or
        `bytearray`, it's sent as a binary message. Otherwise, it's serialized
        to JSON.

        All payloads are validated against ``MAX_PUBLISH_SIZE_BYTES`` (2MB)
        before publishing.

            topic: The topic to publish the message to.
            message: The message payload to publish.

        Raises:
            ValueError: If the payload exceeds the 2MB limit.
            RuntimeError: If the client is not connected. `connect()` must be
                          called first.
        """
        if isinstance(message, (bytes, bytearray)):
            payload_bytes = bytes(message)
            size = len(payload_bytes)
            if size > MAX_PUBLISH_SIZE_BYTES:
                raise ValueError(
                    f"Payload size {size} bytes exceeds the 2MB limit ({MAX_PUBLISH_SIZE_BYTES} bytes)."
                )
            pm = PublishMessage(
                binary_message=BinaryMessage(message=payload_bytes))
        else:
            # Let IPC SDK do JSON serialization of Python primitive/dict/list
            # Convert Mapping to Dict if needed
            dict_message = dict(message) if message is not None else None
            size = len(json.dumps(dict_message).encode("utf-8"))
            if size > MAX_PUBLISH_SIZE_BYTES:
                raise ValueError(
                    f"Payload size {size} bytes exceeds the 2MB limit ({MAX_PUBLISH_SIZE_BYTES} bytes)."
                )
            pm = PublishMessage(
                json_message=JsonMessage(message=dict_message)
            )  # type: ignore

        if self._ipc is None:
            raise RuntimeError(
                "Client is not connected. Call connect() before publishing."
            )
        self._ipc.publish_to_topic(topic=topic, publish_message=pm)

    def _subscribe(
        self, topic: str, handler: Callable[[S], None]
    ) -> Callable[[], None]:
        """Subscribes to a topic and registers a handler for incoming messages.

        When a message is received on the topic, the `handler` function is
        called with the message payload. The method automatically handles
        deserialization. For binary messages, it attempts to decode them as
        JSON first, falling back to raw bytes if decoding fails.

            topic: The topic to subscribe to.
            handler: A callback function that will be invoked with the
                     message payload when a message is received.

        Raises:
            RuntimeError: If the client is not connected. `connect()` must be
                          called first.

        Returns:
            A callable function that, when invoked, will close the subscription
            and stop receiving messages.
        """
        if self._ipc is None:
            raise RuntimeError(
                "Client is not connected. Call connect() before subscribing."
            )

        def _on_event(event) -> None:
            try:
                if not event:
                    return
                if (
                    getattr(event, "json_message", None)
                    and event.json_message.message is not None
                ):
                    handler(event.json_message.message)
                elif (
                    getattr(event, "binary_message", None)
                    and event.binary_message.message is not None
                ):
                    raw = event.binary_message.message
                    # Attempt JSON decode; fallback to raw bytes
                    try:
                        handler(json.loads(raw.decode("utf-8")))
                    except (json.JSONDecodeError, UnicodeDecodeError):
                        handler(raw)
            except (TypeError, ValueError) as e:
                print(
                    f"Failed to process subscription event: {e}", file=sys.stderr)
                traceback.print_exc()

        def _on_error(error: Exception) -> bool:
            print(f"Error in subscription stream: {error}", file=sys.stderr)
            return False

        def _on_closed() -> None:
            print("Subscription stream closed")

        op = self._ipc.subscribe_to_topic(
            topic=topic,
            on_stream_event=_on_event,
            on_stream_error=_on_error,
            on_stream_closed=_on_closed,
        )

        def unsubscribe() -> None:
            op[1].close()

        return unsubscribe

    def _get_environment_variable_context_list(
        self, var_name: str | None
    ) -> list[ContextDefinition]:
        value = self._get_environment_variable(var_name)

        try:
            return list(map(lambda c: ContextDefinition(**c), json.loads(value)))
        except (json.JSONDecodeError, TypeError):
            logger.error("Failed to parse: '%s'. Shutting down.", value)
            sys.exit(3)

    def _get_environment_variable_channel_list(
        self, var_name: str | None
    ) -> list[ChannelDefinition]:
        value = self._get_environment_variable(var_name)

        try:
            return list(map(lambda c: ChannelDefinition(**c), json.loads(value)))
        except (json.JSONDecodeError, TypeError):
            logger.error("Failed to parse: '%s'. Shutting down.", value)
            sys.exit(3)

    def _get_environment_variable_feature_list(
        self, var_name: str | None
    ) -> list[FeatureDefinition]:
        value = self._get_environment_variable(var_name)

        try:
            return list(map(lambda c: FeatureDefinition(**c), json.loads(value)))
        except (json.JSONDecodeError, TypeError):
            logger.error("Failed to parse: '%s'. Shutting down.", value)
            sys.exit(3)

    def _get_environment_variable(self, var_name: str | None) -> str:
        if var_name is None:
            logger.error(
                "%s environment variable is not set. Shutting down.", var_name)
            sys.exit(1)

        value = os.getenv(var_name)
        if value is None:
            logger.error(
                "%s environment variable is not set. Shutting down.", var_name)
            sys.exit(3)
        return value
