# Copyright 2025 Acute Technology Inc. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""Base I2C Controller Class."""

from typing import (
    TYPE_CHECKING,
    ClassVar,
)

import aqpxlib.message as pxmsg
from aqpxlib.abstract import PxAbstractController
from aqpxlib.message import I2CControllerConfig as PxI2CControllerConfig
from aqpxlib.message import I2CControllerState as PxI2CControllerState

if TYPE_CHECKING:
    from aqpxlib.toplevel import AqProtocolExerciser


class PxI2CController(PxAbstractController):
    """An I2C controller device class.

    This class provides a high-level interface to an I2C controller with base I2C fuctionality.
    It does not automatically create an actual instance on the Protocol Exerciser until user call
    the `attach_to_bus` method.

    Typical usage example:

    ```python
    i2c_bus = PxI2CBus(exerciser_session)
    i2c_controller = PxI2CController(exerciser_session)
    i2c_controller.attach_to_bus(i2c_bus)

    ... # Do any operation from this controller

    i2c_controller.detach_from_bus()
    ```
    """

    # Used to identify the class
    identifier: str = "i2c_controller"
    valid_buses: ClassVar[list[str]] = ["i2c"]

    protobuf_key: str = "i2_c_controller"

    # All inherited devices shall add the event fields to the list they acccept
    _protobuf_available_event_fields: ClassVar[list[str]] = []

    def __init__(
        self,
        exerciser: "AqProtocolExerciser",
        *args,
        config: PxI2CControllerConfig | None = None,
        **kwargs,
    ) -> None:
        """Initialize an I2C controller."""
        if config is None:
            config = PxI2CControllerConfig(
                enabled=1,
            )

        super().__init__(exerciser, config=config, *args, **kwargs)

    def send_transactions(
        self,
        transactions: list[pxmsg.I2CTransaction],
    ) -> list[pxmsg.I2CTransaction]:
        """Send a sequence of I2C transactions.

        This method enqueue a sequence of I2C transactions to be sent to the I2C bus,
        starting from a `START` pattern end with `STOP`.
        Each element in `transactions` is an `pxmsg.I2CTransaction` object, which contains one of the
        supported transaction type listed below:

        * pxmsg.I2CWriteTransaction
        * pxmsg.I2CReadTransaction

        The transactions are sent in the order they are provided.

        Args:
            transactions: A list of `pxmsg.I2CTransaction` objects to be sent.

        Returns:
            PxError: If the transaction fails.

        """
        return self.perform_operation(
            operation=pxmsg.PxDeviceOperation(
                i2_c_controller_send_transaction=pxmsg.I2CSendTransaction(
                    transactions=transactions,
                ),
            ),
        ).i2_c_controller_send_transaction.transactions

    def i2c_write(
        self,
        address: int,
        data: list[int],
    ) -> int | None:
        """Perform a write opteration to an I2C device.

        This method sends data to an I2C device with a given `address`.

        Args:
            address: The address of the I2C device to write to.
            data: The data to be written to the I2C device.

        Returns:
            The number of bytes written. Returns None if the target NACKs.

        Raises:
            ValueError: `data` must be a list of integers within (0, 256).

        """
        transactions = [
            pxmsg.I2CTransaction(
                i2_c_write=pxmsg.I2CWriteTransaction(
                    address=address,
                    data=bytes(data),
                ),
            ),
        ]

        resp = self.send_transactions(
            transactions,
        )

        if not resp[0].i2_c_write.is_ack:
            return None
        return resp[0].i2_c_write.bytes_written

    def i2c_write_addr(
        self,
        address: int,
        sub_address: int,
        data: list[int],
    ) -> int | None:
        """Perform a write operation to an I2C device with sub-address.

        This method sends data to an I2C device with a given `address` and `sub_address`.

        Args:
            address: The address of the I2C device to write to.
            sub_address: The sub-address of the I2c device to write to.
            data: The data to be written to the I2C device.

        Returns:
            The number of bytes written. Returns None if the target NACKs.

        Raises:
            OverflowError: If the `sub_address` is not within (0, 256).
            ValueError: bytes must be in range(0, 256)

        """
        data_to_send = sub_address.to_bytes(1, "big") + bytes(data)
        transactions = [
            pxmsg.I2CTransaction(
                i2_c_write=pxmsg.I2CWriteTransaction(
                    address=address,
                    data=data_to_send,
                ),
            ),
        ]

        resp = self.send_transactions(
            transactions,
        )

        if not resp[0].i2_c_write.is_ack:
            return None
        return resp[0].i2_c_write.bytes_written

    def i2c_read(
        self,
        address: int,
        count: int,
    ) -> list[int]:
        """Perform a read operation from an I2C device.

        This method tries to read data from an I2C device with a given `address`
        It may return less than `count` bytes if the target abort the read operation early.

        Args:
            address: The address of the I2C device to read from.
            count: The number of bytes to read from the I2C device.

        Returns:
            The data read from the I2C device. Returns an empty list if the target NACKs.

        Raises:
            ValueError: Cannot perform read operation with count 0.

        """
        if count == 0:
            err_msg = "Cannot perform read operation with count 0"
            raise ValueError(err_msg)

        transactions = [
            pxmsg.I2CTransaction(
                i2_c_read=pxmsg.I2CReadTransaction(
                    address=address,
                    count=count,
                ),
            ),
        ]

        resp = self.send_transactions(
            transactions,
        )

        if not resp[0].i2_c_read.is_ack:
            return []
        return list(resp[0].i2_c_read.data)

    def i2c_read_addr(
        self,
        address: int,
        sub_address: int,
        count: int,
    ) -> list[int]:
        """Perform a read operation from an I2C device with sub-address.

        This method tries to read data from an I2C device with a given `address` and `sub_address`.
        It may return an amount of data which is less than `count` if the target aborts the read operation earlier.

        Args:
            address: The address of the I2C device to read from.
            sub_address: The sub-address of the I2C device to read from.
            count: The numebr of bytes to read from the I2C device.

        Retunrs:
            The data read from I2C device. Returns None if the target NACKs.

        Raises:
            OverflowError: If the `sub_address` is not within (0, 256).

        """
        if count == 0:
            err_msg = "Cannot perform read operation with count 0"
            raise ValueError(err_msg)

        transactions = [
            pxmsg.I2CTransaction(
                i2_c_write=pxmsg.I2CWriteTransaction(
                    address=address,
                    data=sub_address.to_bytes(1, "big"),
                ),
            ),
            pxmsg.I2CTransaction(
                i2_c_read=pxmsg.I2CReadTransaction(
                    address=address,
                    count=count,
                ),
            ),
        ]

        resp = self.send_transactions(
            transactions,
        )

        if not resp[0].i2_c_write.is_ack:
            return []

        return list(resp[1].i2_c_read.data)

    def i2c_scan_all(self) -> list[bool]:
        """Perform scan operation that scanning all available targets.

        This method scan all the available targets by sending write to all the targets.
        Then check each target return ACK or NACK, and set the corresponding element in returned list.

        Returns:
            The list contains the ACK or NACK. Can check whether the target was available \
            by taking the address as index to access the list.

        """
        available_target_list = [False for x in range(0x80)]
        for addr in range(0x08, 0x78, 1):
            transactions = [
                pxmsg.I2CTransaction(
                    i2_c_write=pxmsg.I2CWriteTransaction(
                        address=addr,
                        data=b"",
                    ),
                ),
            ]

            resp = self.send_transactions(
                transactions,
            )

            if resp[0].i2_c_write.is_ack:
                available_target_list[addr] = True

        return available_target_list


__all__ = [
    "PxI2CController",
    "PxI2CControllerConfig",
    "PxI2CControllerState",
]
