# Copyright 2025 Acute Technology Inc. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""Abstract Device Class for all device types."""

from abc import abstractmethod
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, ClassVar, Optional

import betterproto

import aqpxlib.message as pxmsg
from aqpxlib.abstract.abstract_instance import PxAbstractInstance

if TYPE_CHECKING:
    from aqpxlib.toplevel import AqProtocolExerciser

    from .abstract_bus import PxAbstractBus


class PxAbstractDevice(PxAbstractInstance):
    """Abstract Device Class for all device types.

    This class defines methods that are common to all device types.

    """

    valid_buses: ClassVar[list[str]] = []

    # Protobuf class for all device event types
    _protobuf_event_cls: type[betterproto.Message] = pxmsg.PxDeviceEvent
    _protobuf_event_cls_tag: int = 2

    def __init__(
        self,
        exerciser: "AqProtocolExerciser",
        bus: Optional["PxAbstractBus"] = None,
        name: str = "",
        device_id: int = 0,
        config: betterproto.Message | None = None,
    ) -> None:
        """Initialize an abstract device."""
        super().__init__(exerciser)

        # Bus that the device is connected to
        self.bus: PxAbstractBus | None = None
        self.device_id = 0

        self._local_name = name
        self._local_config = config

        if bus is not None:
            if bus.identifier not in self.valid_buses:
                err_msg = f"Invalid bus of type '{bus.identifier}' given to device '{type(self).__name__}', "
                err_msg += f"expected one of {self.valid_buses}"
                raise TypeError(err_msg)
            self.device_id = device_id
            self.bus = bus

    @property
    def state(self) -> betterproto.Message | None:
        """Get the state of the device."""
        if self.bus is None:
            return None
        return self.bus.get_device_state(self)

    def _set_state(self, state: betterproto.Message):
        """Set the state of the device."""
        if self.bus is None:
            return
        self.bus.set_device_state(self, state)

    @property
    def name(self) -> str:
        """Get the name of the device."""
        if self.bus is None:
            return self._local_name

        return self.bus.get_device_pb_msg(self).name

    @property
    def config(self) -> betterproto.Message | None:
        """Get the config of the device."""
        if self.bus is None:
            return self._local_config

        pb_msg = self.bus.get_device_pb_msg(self)
        if pb_msg is not None and hasattr(pb_msg, self.protobuf_key):
            return getattr(pb_msg, self.protobuf_key)

        return None

    @config.setter
    @abstractmethod
    def config(self, config: betterproto.Message) -> None:
        """Set the config of the device."""
        err_msg = "Setting config is not supported for this device"
        raise NotImplementedError(err_msg)

    def __eq__(self, other: object) -> bool:
        """Check if the device is equal to another device.

        This method is used to check if the device is equal to another device.
        """
        if not isinstance(other, PxAbstractDevice):
            return False
        return self.device_id == other.device_id and self.bus == other.bus and self.name == other.name

    def detach_from_bus(self):
        """Detach the device from the bus.

        The opposite of the attach_to_bus method.
        """
        self._local_config = self.config
        self._local_name = self.name

        self.bus.remove_device(self)
        self.bus = None
        self.device_id = 0

    def perform_operation(self, operation: pxmsg.PxDeviceOperation):
        """Send an operation."""
        if self.bus is None:
            err_msg = "Device is not attached to a bus, attach first"
            raise ValueError(err_msg)

        msg = pxmsg.PxProtobufMsg(
            px_device_operation=pxmsg.PxDeviceOperationRequest(
                handle_idx=self.bus.handle_index,
                device_id=self.device_id,
                op=operation,
            ),
        )

        response = self.bus.exerciser.send_request(msg)
        assert response is not None
        return response.px_device_operation.op

    def add_event_handler(self, event: str, handler: Callable[[Any], Any]):
        """Add an event handler to the current device instance."""
        event_id = self.available_events_map.get(event)
        if event_id is None:
            err_msg = f"Event {event} is not available for this instance"
            raise ValueError(err_msg)

        event_id = f"{self._protobuf_event_cls_tag:08X}-{self.device_id:08X}-{event_id}"
        self.exerciser.add_event_handler(event_id, handler)

    def get_event_handlers(self) -> dict[str, Callable]:
        """Retreive handlers registerd on this device."""
        event_prefix = [
            f"{self._protobuf_event_cls_tag:08X}-{self.device_id:08X}-{self.available_events_map[e]}"
            for e in self._protobuf_available_event_fields
        ]
        event_maps = self.exerciser.get_event_handlers(
            lambda x: any(x.startswith(prefix) for prefix in event_prefix),
        )

        return_map = {}
        for event_id, handler in event_maps.items():
            event_name = self._protobuf_event_cls._betterproto.field_name_by_number.get(  # noqa: SLF001
                int(event_id.split("-")[-1], 16),
            )
            return_map[event_name] = handler

        return return_map


class PxAbstractController(PxAbstractDevice):
    """Abstract Controller Class for all controller types.

    This is a generic controller class for convenience that can be used to create a controller device,
    or a device that acts as a transmitter, such as UART Tx.

    Attributes:
        is_controller (bool): Whether the device is a controller or act as an transmitter (TX).

    """

    is_controller = True

    def attach_to_bus(self, bus: "PxAbstractBus"):
        """Attach the controller to a bus.

        This method is used to make the device aware of the bus it is connected to.
        It is used to set the bus attribute of the device.

        Note:
            This method will not automatically create an actual instance on
            the Protocol Exerciser until user call the `attach_to_bus` method.

        Args:
            bus (PxAbstractBus): The bus to attach the device to

        Raises:
            ValueError: If the device is already attached to a bus

        """
        if self.bus is not None:
            err_msg = "Device already attached to a bus, detach first"
            raise ValueError(err_msg)

        bus.attach_controller(self)

    @PxAbstractDevice.config.setter
    def config(self, config: betterproto.Message) -> None:
        """Set the config of the controller."""
        if self.bus is None:
            return

        bus_config = self.bus.get_config()
        if bus_config is None:
            err_msg = "Config is not supported for this device"
            raise ValueError(err_msg)

        for idx, controller_config in enumerate(bus_config.controllers):
            if controller_config.id == self.device_id:
                bus_config.controllers[idx] = pxmsg.PxControllerDevice.from_dict(
                    {
                        "id": self.device_id,
                        "name": self.name,
                        self.protobuf_key: config.to_dict(),
                    },
                )

                self.bus.set_config(bus_config)
                break

        self._local_config = config


class PxAbstractTarget(PxAbstractDevice):
    """Abstract Target Class for all target types.

    This is a generic target class for convenience that can be used to create a target device,
    or a device that acts as a receiver, such as UART Rx.

    Attributes:
        is_controller (bool): Whether the device is a controller or act as an transmitter (TX).

    """

    is_controller = False

    def attach_to_bus(self, bus: "PxAbstractBus"):
        """Attach the target to a bus.

        This method is used to make the device aware of the bus it is connected to.
        It is used to set the bus attribute of the device.

        Note:
            This method will not automatically create an actual instance on
            the Protocol Exerciser until user call the `attach_to_bus` method.

        Args:
            bus (PxAbstractBus): The bus to attach the device to

        Raises:
            ValueError: If the device is already attached to a bus

        """
        if self.bus is not None:
            err_msg = "Device already attached to a bus, detach first"
            raise ValueError(err_msg)

        bus.attach_target(self)

    @PxAbstractDevice.config.setter
    def config(self, config: betterproto.Message) -> None:
        """Set the config of the target."""
        if self.bus is None:
            return

        bus_config = self.bus.get_config()
        if bus_config is None:
            err_msg = "Config is not supported for this device"
            raise ValueError(err_msg)

        for idx, target_config in enumerate(bus_config.targets):
            if target_config.id == self.device_id:
                bus_config.targets[idx] = pxmsg.PxTargetDevice.from_dict(
                    {
                        "id": self.device_id,
                        "name": self.name,
                        self.protobuf_key: config.to_dict(),
                    },
                )

                self.bus.set_config(bus_config)
                break

        self._local_config = config
