# Copyright 2025 Acute Technology Inc. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""Acute Protocol Exerciser Top Level Class."""

from collections.abc import Callable
from typing import Any, ClassVar

import aqpxlib.message as pxmsg
from aqpxlib.abstract import PxAbstractBus
from aqpxlib.i2c import PxI2CBus
from aqpxlib.i3c import PxI3CBus
from aqpxlib.sessions import PxSession
from aqpxlib.spi import PxSPIBus
from aqpxlib.uart import PxUartBus


class AqProtocolExerciser:
    """AqProtocolExerciser class."""

    # A mapping of bus identifiers to bus classes
    _available_buses: ClassVar[list[type[PxAbstractBus]]] = [
        PxI2CBus,
        PxI3CBus,
        PxSPIBus,
        PxUartBus,
    ]

    _available_buses_by_protobuf_key: ClassVar[dict[str, type[PxAbstractBus]]] = {
        cls.protobuf_key: cls for cls in _available_buses
    }

    _available_buses_by_identifier: ClassVar[dict[str, type[PxAbstractBus]]] = {
        cls.identifier: cls for cls in _available_buses
    }

    def __init__(self, generation: str = "PX2K", host: str = "localhost", port: int = 60600):
        """Top level class that represents a Protocol Exerciser session."""
        self.generation: str = generation
        self._session: PxSession = PxSession()
        self._session.connect(host=host, port=port)

    def __enter__(self):
        """Create and connect to a Protocol Exerciser session."""
        return self

    def __exit__(self, *args):
        """Disconnect to an Protocol Exerciser session."""
        self._session.disconnect()

    @classmethod
    def connect(cls, host: str = "localhost", port: int = 60600) -> "AqProtocolExerciser":
        """Class method of creating and connecting to a Protocol Exerciser session.

        Args:
            host (str, optional): The hostname or IP address to connect to. Defaults to "localhost".
            port (int, optional): The port to connect to. Defaults to 60600.

        """
        return cls(host=host, port=port)

    def disconnect(self) -> None:
        """Disconnect from the exerciser session."""
        self._session.disconnect()

    def send_request(self, msg: pxmsg.PxProtobufMsg) -> pxmsg.PxProtobufMsg | None:
        """Send a request to the exerciser.

        Args:
            msg (pxmsg.PxProtobufMsg): The message to send

        Returns:
            pxmsg.PxProtobufMsg: The response from the exerciser

        """
        return self._session.request(msg)

    def subscribe(self, event_type: str):
        """Subscribe to an event."""
        self._session.subscribe(event_type)

    def get_handle(self, handle_index: int) -> pxmsg.PxHandle:
        """Get the handle."""
        msg = pxmsg.PxProtobufMsg(
            get_handle_request=pxmsg.TopologyPxHandleRequest(
                handle_idx=handle_index,
            ),
        )

        response = self._session.request(msg)
        assert response is not None
        return response.get_handle_request.px_handle

    def set_handle(self, handle_index: int, handle: pxmsg.PxHandle) -> pxmsg.PxHandle:
        """Set the handle."""
        msg = pxmsg.PxProtobufMsg(
            set_handle_request=pxmsg.TopologyPxHandleRequest(
                handle_idx=handle_index,
                px_handle=handle,
            ),
        )

        response = self._session.request(msg)
        assert response is not None
        return response.set_handle_request.px_handle

    def get_handles(self) -> list[pxmsg.PxHandle]:
        """Get the handles."""
        msg = pxmsg.PxProtobufMsg(
            get_handles_request=pxmsg.TopologyPxHandlesRequest(),
        )

        response = self._session.request(msg)
        assert response is not None
        return response.get_handles_request.px_handles

    def set_handles(self, handles: list[pxmsg.PxHandle]):
        """Set the handles."""
        msg = pxmsg.PxProtobufMsg(
            set_handles_request=pxmsg.TopologyPxHandlesRequest(px_handles=handles),
        )

        response = self._session.request(msg)
        assert response is not None
        return response.set_handles_request.px_handles

    def create_bus(self, bus_identifier: str, *args, **kwargs) -> PxAbstractBus:
        """Create an bus instance from current Exerciser instance."""
        bus_class = self._available_buses_by_identifier.get(bus_identifier, None)
        if bus_class is None:
            err_msg = f"Bus identifier: {bus_identifier} is not supported!"
            raise ValueError(err_msg)

        return bus_class(self, *args, **kwargs)

    def add_event_handler(self, event: str, handler: Callable[[Any], Any]):
        """Add an event handler."""
        self._session.add_event_handler(event, handler)

    def remove_event_handler(self, event: str):
        """Remove an event handler."""
        self._session.remove_event_handler(event)

    def get_event_handlers(
        self,
        filter_by: Callable[[str], bool] | None = None,
    ) -> dict[str, Callable]:
        """Get all available event handlers."""
        return self._session.get_event_handlers(filter_by)

    def set_pod_io_configs(self, configs: list[pxmsg.PxPodIoElectricalConfig]):
        """Set electrical configuration on PXPOD."""
        msg = pxmsg.PxProtobufMsg(
            set_pod_io_electrical_configs_request=pxmsg.PxPodIoElectricalConfigsRequest(
                configs=configs,
            ),
        )
        _ = self._session.request(msg)

    def get_pod_io_configs(self) -> list[pxmsg.PxPodIoElectricalConfig]:
        """Get current electrical configuration on PXPOD."""
        msg = pxmsg.PxProtobufMsg(
            get_pod_io_electrical_configs_request=pxmsg.PxPodIoElectricalConfigsRequest(),
        )

        response = self._session.request(msg)
        assert response is not None
        return response.get_pod_io_electrical_configs_request.configs
