1"""
  2Functions for accessing SMBus.
  3
  4# (c) 2020-2024 TurnAround Factor, Inc.
  5#
  6# CUI DISTRIBUTION CONTROL
  7# Controlled by: DLA J68 R&D SBIP
  8# CUI Category: Small Business Research and Technology
  9# Distribution/Dissemination Controls: PROTECTED BY SBIR DATA RIGHTS
 10# POC: GOV SBIP Program Manager Denise Price, 571-767-0111
 11# Distribution authorized to U.S. Government Agencies only, to protect information not owned by the
 12# U.S. Government and protected by a contractor’s ‘limited rights’ statement, or received with the understanding that
 13# it is not routinely transmitted outside the U.S. Government (determination made September 14, 2024). Other requests
 14# for this document shall be referred to ACOR DLA Logistics Operations (J-68), 8725 John J. Kingman Rd., Suite 4317,
 15# Fort Belvoir, VA 22060-6221
 16#
 17# SBIR DATA RIGHTS
 18# Contract No.:SP4701-23-C-0083
 19# Contractor Name: TurnAround Factor, Inc.
 20# Contractor Address: 10365 Wood Park Ct. Suite 313 / Ashland, VA 23005
 21# Expiration of SBIR Data Rights Period: September 24, 2029
 22# The Government's rights to use, modify, reproduce, release, perform, display, or disclose technical data or computer
 23# software marked with this legend are restricted during the period shown as provided in paragraph (b)(4) of the Rights
 24# in Noncommercial Technical Data and Computer Software--Small Business Innovative Research (SBIR) Program clause
 25# contained in the above identified contract. No restrictions apply after the expiration date shown above. Any
 26# reproduction of technical data, computer software, or portions thereof marked with this legend must also reproduce
 27# the markings.
 28"""
 29
 30from __future__ import annotations
 31
 32import errno
 33from string import Template
 34from typing import ClassVar, cast
 35
 36import pytest
 37from smbus2 import SMBus as PySMBus
 38from typing_extensions import Self
 39
 40from hitl_tester.modules.bms_types import BMSFlags
 41from hitl_tester.modules.bms.smbus_types import SMBusReg, SMBusRegType, BatteryMode, SMBusFunctionType, SMBusError
 42
 43SMBUS_ADDRESS = 0xB
 44
 45
 46class SMBus:
 47    """SMBus object"""
 48
 49    instance: ClassVar[Self | None] = None
 50
 51    def __new__(cls):
 52        """Make SMBus a singleton."""
 53        if cls.instance is None:
 54            cls.instance = super().__new__(cls)
 55        return cls.instance
 56
 57    def read_register(self, register: SMBusReg) -> tuple[SMBusRegType | None, bytes]:
 58        """Read from an SMBus register."""
 59        assert hasattr(pytest, "flags") and isinstance(pytest.flags, BMSFlags)
 60        channel = cast(str, pytest.flags.config["smbus_i2c_channel"])
 61
 62        size = register.value.size
 63        try:
 64            with PySMBus(channel) as bus:
 65                if block_read := size is None:  # Simulate an SMBus block read as the Pi driver doesn't support it
 66                    size = min(32, bus.read_i2c_block_data(SMBUS_ADDRESS, register.value.address, 1)[0] + 1)
 67                data = bytes(bus.read_i2c_block_data(SMBUS_ADDRESS, register.value.address, size))[block_read:]
 68        except OSError as e:
 69            raise SMBusError(
 70                f"{errno.errorcode[e.errno]} ({e.strerror}) when reading "
 71                f"{register.value.size} byte{'' if register.value.size==1 else 's'} from register "
 72                f"0x{register.value.address:02X} ({register.name}), channel {channel}, "
 73                f"SMBus address 0x{SMBUS_ADDRESS:02X}"
 74            ) from e
 75
 76        try:
 77            return register.value.type.value(data), data
 78        except (ValueError, IndexError, KeyError):
 79            return None, data
 80
 81    def write_register(self, register: SMBusReg, word: int):
 82        """Write to an SMBus register."""
 83        if not register.value.writable:
 84            raise RuntimeError("Attempted to write to read-only register.")
 85
 86        assert hasattr(pytest, "flags") and isinstance(pytest.flags, BMSFlags)
 87        channel = cast(str, pytest.flags.config["smbus_i2c_channel"])
 88
 89        try:
 90            with PySMBus(channel) as bus:
 91                bus.write_word_data(SMBUS_ADDRESS, register.value.address, word)
 92        except OSError as e:
 93            raise SMBusError(
 94                f"{errno.errorcode[e.errno]} ({e.strerror}) when writing "
 95                f"0x{word:04X} to register 0x{register.value.address:02X} ({register.name}), channel {channel}, "
 96                f"SMBus address 0x{SMBUS_ADDRESS:02X})"
 97            ) from e
 98
 99    def parse_smbus_data(self) -> dict[str, SMBusRegType | None]:
100        """Return a dict with all smbus data."""
101
102        def to_title(text: str) -> str:
103            """Format titles."""
104            return text.title().replace("_", " ")
105
106        try:
107            parsed_battery_mode = cast(BatteryMode, self.read_register(SMBusReg.BATTERY_MODE)[0])
108        except SMBusError:
109            parsed_battery_mode = BatteryMode(bytes())
110        capacity_unit = "cW" if parsed_battery_mode.capacity_mode else "mA"
111
112        smbus_data: dict[str, SMBusRegType | None] = {}
113        for register in SMBusReg:
114            unit = Template(register.value.unit).safe_substitute(capacity=capacity_unit)
115            title = f"{to_title(register.name)}{f' ({unit})' if unit else ''}"
116            try:
117                register_value, register_raw = self.read_register(register)
118            except SMBusError:
119                register_value, register_raw = None, bytes()
120
121            if register in (SMBusReg.SPECIFICATION_INFO, SMBusReg.BATTERY_STATUS, SMBusReg.BATTERY_MODE):
122                register_value = register_value or register.value.type.value(bytes([0x31]))  # Create dummy object
123                for attribute in register_value.__dict__:
124                    smbus_data[f"{title} {to_title(attribute)}"] = str(getattr(register_value, attribute))
125                if register is SMBusReg.BATTERY_MODE:
126                    smbus_data[f"{title} (Raw)"] = f"0x{register_raw.hex().upper()}"
127            else:
128                smbus_data[title] = register_value
129                if register.value.type not in (
130                    SMBusFunctionType.INT,
131                    SMBusFunctionType.UNSIGNED_INT,
132                    SMBusFunctionType.HEX,
133                ):
134                    smbus_data[f"{title} (Raw)"] = f"0x{register_raw.hex().upper()}"
135
136        return smbus_data