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

import logging
from abc import abstractmethod
from collections.abc import Callable
from functools import cached_property
from typing import TYPE_CHECKING, Any, ClassVar

import betterproto

import aqpxlib.message as pxmsg
from aqpxlib.abstract.abstract_device import PxAbstractDevice
from aqpxlib.abstract.abstract_instance import PxAbstractInstance
from aqpxlib.exceptions import PxError

if TYPE_CHECKING:
    from aqpxlib.toplevel import AqProtocolExerciser

_LOGGER = logging.getLogger(__name__)


class PxAbstractBus(PxAbstractInstance):
    """Abstract Bus Class for all bus types."""

    _available_devices: ClassVar[list[type[PxAbstractDevice]]] = []

    controller_class: type[PxAbstractDevice]

    @cached_property
    def available_devices_by_protobuf_key(self):
        """A mapping of device protobuf field name to device classes."""
        return {cls.protobuf_key: cls for cls in self._available_devices}

    @cached_property
    def available_devices_by_identifier(self):
        """A mapping of device identifier (from abstract instance) to device classes."""
        return {cls.identifier: cls for cls in self._available_devices}

    def __init__(self, exerciser: "AqProtocolExerciser") -> None:
        """Initialize an abstract bus."""
        super().__init__(exerciser)
        self._handle_index: int = -1

    def _get_electrical_config(self) -> pxmsg.PxElectricalConfig | None:
        """Get the electrical configuration of the bus."""
        assert self._handle_index != -1, "Bus handle index is not set"

        _LOGGER.info("Getting electrical config for bus %d", self._handle_index)
        msg = pxmsg.PxProtobufMsg(
            get_electrical_config_request=pxmsg.PxElectricalConfigRequest(
                handle_idx=self._handle_index,
            ),
        )

        response = self.exerciser.send_request(msg)
        assert response is not None
        return response.get_electrical_config_request.config

    def _set_electrical_config(self, config: pxmsg.PxElectricalConfig) -> bool:
        """Set the electrical configuration of the bus."""
        assert self._handle_index != -1, "Bus handle index is not set"

        msg = pxmsg.PxProtobufMsg(
            set_electrical_config_request=pxmsg.PxElectricalConfigRequest(
                handle_idx=self._handle_index,
                config=config,
            ),
        )

        try:
            _ = self.exerciser.send_request(msg)
        except PxError as e:
            _LOGGER.exception("Failed to set electrical config: %s", e)
            return False

        return True

    @property
    def handle_index(self) -> int:
        """The index of the bus interface."""
        return self._handle_index

    @handle_index.setter
    def handle_index(self, handle_index: int):
        """Set the index of the bus interface."""
        self._handle_index = handle_index

    @property
    def voltage(self) -> int | None:
        """The voltage of the bus."""
        electrical_config = self._get_electrical_config()
        if electrical_config is None:
            return None

        return electrical_config.voltage

    @voltage.setter
    def voltage(self, voltage: int):
        """Set the voltage of the bus."""
        if voltage < 0:
            err_msg = f"Invalid voltage value: {voltage}"
            raise ValueError(err_msg)

        config = pxmsg.PxElectricalConfig(
            voltage=voltage,
        )

        if not self._set_electrical_config(config):
            err_msg = f"Failed to set voltage: {voltage}"
            raise ValueError(err_msg)

    @property
    def la_threshold_voltage(self) -> int | None:
        """The la threshold voltage of the bus."""
        electrical_config = self._get_electrical_config()
        if electrical_config is None:
            return None

        return electrical_config.la_threshold_voltage

    @la_threshold_voltage.setter
    def la_threshold_voltage(self, voltage: int):
        """Set the la threshold voltage of the bus."""
        if voltage < 0:
            err_msg = f"Invalid voltage value: {voltage}"
            raise ValueError(err_msg)

        config = pxmsg.PxElectricalConfig(
            la_threshold_voltage=voltage,
        )

        if not self._set_electrical_config(config):
            err_msg = f"Failed to set la threshold voltage: {voltage}"
            raise ValueError(err_msg)

    @property
    def pullup_resistance(self) -> int | None:
        """The pullup resistance of the bus."""
        electrical_config = self._get_electrical_config()
        if electrical_config is None:
            return None

        return electrical_config.pullup_resistance

    @pullup_resistance.setter
    def pullup_resistance(self, resistance: int):
        """Set the pullup resistance of the bus."""
        if resistance < 0:
            err_msg = f"Invalid resistance value: {resistance}"
            raise ValueError(err_msg)

        config = pxmsg.PxElectricalConfig(
            pullup_resistance=resistance,
        )

        if not self._set_electrical_config(config):
            err_msg = f"Failed to set pullup resistance: {resistance}"
            raise ValueError(err_msg)

    @property
    def pulldown_resistance(self) -> int | None:
        """The pulldown resistance of the bus."""
        electrical_config = self._get_electrical_config()
        if electrical_config is None:
            return None

        return electrical_config.pulldown_resistance

    @pulldown_resistance.setter
    def pulldown_resistance(self, resistance: int):
        """Set the pulldown resistance of the bus."""
        if resistance < 0:
            err_msg = f"Invalid resistance value: {resistance}"
            raise ValueError(err_msg)

        config = pxmsg.PxElectricalConfig(
            pulldown_resistance=resistance,
        )

        if not self._set_electrical_config(config):
            err_msg = f"Failed to set pulldown resistance: {resistance}"
            raise ValueError(err_msg)

    @property
    def bus_hold_enable(self) -> bool | None:
        """The bus hold enable of the bus."""
        electrical_config = self._get_electrical_config()
        if electrical_config is None:
            return None

        return electrical_config.bus_hold_enable

    @bus_hold_enable.setter
    def bus_hold_enable(self, enable: bool):
        """Set the bus hold enable of the bus."""
        config = pxmsg.PxElectricalConfig(
            bus_hold_enable=enable,
        )

        if not self._set_electrical_config(config):
            err_msg = f"Failed to set bus hold enable: {enable}"
            raise ValueError(err_msg)

    def get_config(self) -> pxmsg.PxHandle | None:
        """Get the configuration of the bus."""
        assert self._handle_index != -1, "Bus handle index is not set"

        handle_config = self.exerciser.get_handle(self._handle_index)

        bus_config = handle_config
        bus_key, bus_config = betterproto.which_one_of(bus_config, "protocol")
        if bus_key == self.protobuf_key:
            return handle_config
        _LOGGER.exception("Bus configuration is not supported")
        return None

    def get_target_pb_msg(
        self,
        target: PxAbstractDevice,
    ) -> betterproto.Message | None:
        """Get the configuration of the target from the bus."""
        config = self.get_config()
        if config is None:
            _LOGGER.exception("Failed to get bus configuration")
            return None

        for target_config in config.targets:
            if target_config.id == target.device_id:
                return target_config
        return None

    def get_controller_pb_msg(
        self,
        controller: PxAbstractDevice,
    ) -> betterproto.Message | None:
        """Get the configuration of the controller from the bus."""
        controllers_config = self.get_config().controllers
        for controller_config in controllers_config:
            if controller_config.id == controller.device_id:
                return controller_config
        return None

    def remove_controller(self, controller: PxAbstractDevice):
        """Remove a controller from the bus."""
        config = self.get_config()
        if config is None:
            _LOGGER.exception("Failed to get bus configuration")
            return

        config.controllers = [
            controller_pb_msg
            for controller_pb_msg in config.controllers
            if controller_pb_msg.id != controller.device_id
        ]
        self.set_config(config)

    def remove_target(self, target: PxAbstractDevice):
        """Remove a target from the bus."""
        config = self.get_config()
        if config is None:
            _LOGGER.exception("Failed to get bus configuration")
            return

        config.targets = [target_pb_msg for target_pb_msg in config.targets if target_pb_msg.id != target.device_id]

        self.set_config(config)

    def remove_device(self, device: PxAbstractDevice):
        """Remove a device from the bus."""
        if self.controller_class is not None and isinstance(
            device,
            self.controller_class,
        ):
            self.remove_controller(device)
        else:
            self.remove_target(device)

    def get_device_state(
        self,
        device: PxAbstractDevice,
    ) -> betterproto.Message | None:
        """Get the state of the device from the bus."""
        if type(device) not in self._available_devices:
            err_msg = f"Device {device} is not supported on {self} bus."
            raise ValueError(err_msg)

        msg = pxmsg.PxProtobufMsg(
            px_device_operation=pxmsg.PxDeviceOperationRequest(
                handle_idx=self.handle_index,
                device_id=device.device_id,
                op=pxmsg.PxDeviceOperation(
                    px_device_get_state=pxmsg.PxDeviceState(),
                ),
            ),
        )

        response = self.exerciser.send_request(
            msg,
        )
        assert response is not None
        dev_state = response.px_device_operation.op.px_device_get_state
        _, state = betterproto.which_one_of(dev_state, "state")
        return state

    def set_device_state(self, device: PxAbstractDevice, state: betterproto.Message):
        """Set the state of the device on the bus."""
        if type(device) not in self._available_devices:
            err_msg = f"Device {device} is not supported on {self} bus."
            raise ValueError(err_msg)

        msg = pxmsg.PxProtobufMsg(
            px_device_operation=pxmsg.PxDeviceOperationRequest(
                handle_idx=self.handle_index,
                device_id=device.device_id,
                op=pxmsg.PxDeviceOperation(
                    px_device_set_state=pxmsg.PxDeviceState(
                        **{
                            device.protobuf_key: state,
                        },
                    ),
                ),
            ),
        )

        _ = self.exerciser.send_request(msg)

    def get_device_pb_msg(
        self,
        device: PxAbstractDevice,
    ) -> betterproto.Message | None:
        """Get the configuration of the device from the bus."""
        if isinstance(device, self.controller_class):
            return self.get_controller_pb_msg(device)
        return self.get_target_pb_msg(device)

    def set_config(self, config: pxmsg.PxHandle) -> pxmsg.PxHandle:
        """Set the configuration of the bus."""
        assert self._handle_index != -1, "Bus handle index is not set"

        return self.exerciser.set_handle(self._handle_index, config)

    def get_controllers(self) -> list[PxAbstractDevice]:
        """Get the controllers of the bus."""
        config = self.get_config()
        if config is None:
            _LOGGER.exception("Failed to get bus configuration")
            return []

        if hasattr(config, "controllers"):
            # TODO: Handle buses that have more than 1 controller
            for controller_config in config.controllers:
                if self.controller_class is not None:
                    controller = self.controller_class(
                        self.exerciser,
                        bus=self,
                        device_id=controller_config.id,
                    )
                    return [controller]

        _LOGGER.info("Can not find any controllers")
        return []

    def attach_controller(self, controller: PxAbstractDevice):
        """Attach a controller to the bus."""
        if type(controller) not in self._available_devices:
            err_msg = f"Device {controller} is not supported on {self} bus."
            raise ValueError(err_msg)

        config = self.get_config()
        if config is None:
            _LOGGER.exception("Failed to get bus configuration")
            return
        if hasattr(config, "controller") and betterproto.serialized_on_wire(
            config.controller,
        ):
            err_msg = f"Bus already has a controller, config = {config}"
            raise ValueError(err_msg)

        config.controllers.append(
            pxmsg.PxControllerDevice.from_dict(
                {
                    "id": 0,
                    "name": controller.name,
                    controller.protobuf_key: controller.config.to_dict(),
                },
            ),
        )
        controller.device_id = self.set_config(config).controllers[len(config.controllers) - 1].id
        controller.bus = self

    def attach_target(self, target: PxAbstractDevice):
        """Attach a target to the bus."""
        if type(target) not in self._available_devices:
            err_msg = f"Device {target} is not supported on {self} bus."
            raise ValueError(err_msg)

        config = self.get_config()
        if config is None:
            _LOGGER.exception("Failed to get bus configuration")
            return
        config.targets.append(
            pxmsg.PxTargetDevice.from_dict(
                {
                    "id": 0,
                    "name": target.name,
                    target.protobuf_key: target.config.to_dict(),
                },
            ),
        )

        target.device_id = self.set_config(config).targets[len(config.targets) - 1].id
        target.bus = self

    def get_targets(self) -> list[PxAbstractDevice]:
        """Get the targets of the bus."""
        config = self.get_config()
        if config is None:
            _LOGGER.exception("Failed to get bus configuration")
            return []

        targets = []
        for target_pb_msg in config.targets:
            pb_key, _ = betterproto.which_one_of(target_pb_msg, "config")
            target_class = self.available_devices_by_protobuf_key.get(pb_key, None)
            if target_class is not None:
                target = target_class(
                    self.exerciser,
                    bus=self,
                    device_id=target_pb_msg.id,
                )
                targets.append(target)
        return targets

    def get_devices(self) -> list[PxAbstractDevice]:
        """Get the devices of the bus."""
        return self.get_controllers() + self.get_targets()

    def create_device(
        self,
        exerciser: "AqProtocolExerciser",
        device_identifier: str,
        *args,
        **kwargs,
    ) -> PxAbstractDevice:
        """Create a new device on the bus.

        A new device interface will be created and attached to the bus.

        Args:
            exerciser (AqProtocolExerciser): An exerciser instance which is used to send request to server
            device_identifier (str): The unique identifier of the device to create
            *args: Additional arguments to pass to the device constructor
            **kwargs: Additional keyword arguments to pass to the device constructor

        Returns:
            PxAbstractDevice: The interface of the created device

        Raises:
            ValueError: If the `device_identifier` is not supported

        """
        if device_identifier not in self.available_devices_by_identifier:
            err_msg = f"Device {device_identifier} not supported"
            raise ValueError(err_msg)

        dev = self.available_devices_by_identifier[device_identifier](
            exerciser,
            *args,
            **kwargs,
        )
        dev.attach_to_bus(self)

        return dev

    @abstractmethod
    def _to_protobuf_handle_config(self, *args, **kwargs) -> pxmsg.PxHandle:
        """Convert the bus to a protobuf handle configuration.

        Args:
            *args: Additional arguments to pass to the bus constructor
            **kwargs: Additional keyword arguments to pass to the bus constructor

        Returns:
            pxmsg.PxHandle: The protobuf structure of the bus

        Raises:
            NotImplementedError: If the bus subclass does not implement this method

        """
        err_msg = "Bus subclass must implement this method"
        raise NotImplementedError(err_msg)

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

        This method is used to check if the bus is equal to another bus.

        Args:
            other (Any): The other bus to compare with or a protobuf structure of the bus

        Raises:
            NotImplementedError: If the bus subclass does not implement this method

        """
        err_msg = "Bus subclass must implement this method"
        raise NotImplementedError(err_msg)

    def _find_or_create_bus(self, *args, **kwargs) -> int:
        """Find or create a bus interface.

        This method is used to find an available handle slot and create a new bus interface.
        It the bus handle does not exist, a new bus interface will be created.
        Otherwise, the existing bus interface will be returned.

        Args:
            *args: Additional arguments to pass to the bus constructor
            **kwargs: Additional keyword arguments to pass to the bus constructor

        Returns:
            int: The index of the bus interface if found, otherwise -1

        Raises:
            ValueError: If no empty slot is found

        """
        handles = self.exerciser.get_handles()
        if handles is None:
            _LOGGER.exception("Failed to get handles")
            return -1

        available_handle_slot = -1
        for idx, handle in enumerate(handles):
            name, current_handle = betterproto.which_one_of(handle, "protocol")
            if name == self.protobuf_key:
                if self == current_handle:
                    return idx
            elif name in ["resv", ""]:
                # Compare with reserved handle and unset handles
                available_handle_slot = idx if available_handle_slot == -1 else available_handle_slot

        _LOGGER.info("No existing bus found, creating a new bus")
        if available_handle_slot == -1:
            _LOGGER.exception("No empty slot found, cannot create a new bus")
            return -1

        # Update topology
        new_bus = self._to_protobuf_handle_config(*args, **kwargs)
        self.exerciser.set_handle(available_handle_slot, new_bus)

        _LOGGER.info("New bus created at index %d", available_handle_slot)

        return available_handle_slot

    def add_event_handler(self, event: str, handler: Callable[[Any], Any]):
        """Add an event handler to the current bus instance."""
        err_msg = "Bus subclass must implement this method"
        raise NotImplementedError(err_msg)
