# 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.data_generator import PxDataGeneratorBus
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]]] = [
        PxDataGeneratorBus,
        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
