"""I2C master over a CH347 USB adapter (Waveshare USB to UART/I2C/SPI/JTAG).

Knows nothing about any particular I2C device -- it only moves bytes to and
from a 7-bit address.

The CH347 in Mode 1 (UART1+I2C+SPI) presents three interfaces: two CDC-ACM
interfaces that the kernel binds for the UART, and interface 2, a vendor-class
interface carrying I2C/SPI/JTAG that no kernel driver claims. We talk to
interface 2 directly over libusb.

Framing follows the CH347 stream protocol as implemented by the
aystarik/ch347-i2c-spi-gpio Linux driver, whose USB id table entry
USB_DEVICE_INTERFACE_NUMBER(0x1a86, 0x55db, 0x02) matches this device.
"""

import usb.core
import usb.util

VENDOR_ID = 0x1A86
PRODUCT_ID = 0x55DB
INTERFACE = 2

EP_OUT = 0x06
EP_IN = 0x86

TIMEOUT_MS = 1000

# Stream protocol opcodes.
CMD_STREAM = 0xAA  # start of an I2C command stream
CMD_END = 0x00  # end of stream
CMD_STA = 0x74  # emit START
CMD_STO = 0x75  # emit STOP
CMD_OUT = 0x80  # write; low 6 bits are the byte count
CMD_IN = 0xC0  # read; low 6 bits are the byte count
CMD_SET = 0x60  # set bus speed; low bits select the rate

# Bus speeds, as the low nibble of CMD_SET.
SPEED_20KHZ = 0
SPEED_100KHZ = 1
SPEED_400KHZ = 2
SPEED_750KHZ = 3

# The count field is 6 bits, and a write also spends one byte on the address.
MAX_XFER = 0x3F

ACK = 1


class CH347I2C:
    """I2C master on a CH347 adapter.

    Every byte the adapter clocks out produces one status byte in the reply,
    where 1 means the slave ACKed. A read reply is one address-ACK byte
    followed by the payload.
    """

    def __init__(self, speed=SPEED_100KHZ):
        self._dev = usb.core.find(idVendor=VENDOR_ID, idProduct=PRODUCT_ID)
        if self._dev is None:
            raise IOError(
                f"No CH347 adapter found ({VENDOR_ID:04x}:{PRODUCT_ID:04x}). "
                "Is it plugged in and in Mode 1?"
            )

        # Deliberately no set_configuration(): the device is already configured,
        # and re-setting it would reset the CDC-ACM interfaces the kernel is
        # using for /dev/ttyACM*.
        try:
            usb.util.claim_interface(self._dev, INTERFACE)
        except usb.core.USBError as exc:
            raise IOError(
                f"Cannot claim CH347 interface {INTERFACE}: {exc}. "
                "If this is a permission error, run ./setup.sh."
            ) from exc

        self._claimed = True
        self.set_speed(speed)

    def close(self):
        if getattr(self, "_claimed", False):
            usb.util.release_interface(self._dev, INTERFACE)
            usb.util.dispose_resources(self._dev)
            self._claimed = False

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.close()
        return False

    def _xfer(self, out_bytes, in_len):
        """Send one command stream, then read in_len status/data bytes."""
        written = self._dev.write(EP_OUT, bytes(out_bytes), TIMEOUT_MS)
        if written != len(out_bytes):
            raise IOError(
                f"Short USB write to CH347: sent {written} of {len(out_bytes)} bytes"
            )
        if in_len == 0:
            return b""

        reply = bytes(self._dev.read(EP_IN, in_len, TIMEOUT_MS))
        if len(reply) != in_len:
            raise IOError(
                f"Short USB read from CH347: got {len(reply)} of {in_len} bytes"
            )
        return reply

    def set_speed(self, speed):
        """Select the I2C clock rate (one of the SPEED_* constants)."""
        if speed not in (SPEED_20KHZ, SPEED_100KHZ, SPEED_400KHZ, SPEED_750KHZ):
            raise ValueError(f"Invalid I2C speed {speed}, expected 0-3")
        self._xfer([CMD_STREAM, CMD_SET | speed, CMD_END], 0)

    def write(self, addr, data):
        """Write data (bytes) to a 7-bit address. Raises IOError on NACK."""
        data = bytes(data)
        # One byte of the transfer budget goes to the address.
        if len(data) > MAX_XFER - 1:
            raise ValueError(
                f"Write of {len(data)} bytes exceeds the CH347 limit of {MAX_XFER - 1}"
            )

        packet = [CMD_STREAM, CMD_STA, CMD_OUT | (len(data) + 1), addr << 1]
        packet += data
        packet += [CMD_STO, CMD_END]

        # One status byte per clocked-out byte: the address plus the payload.
        reply = self._xfer(packet, len(data) + 1)
        if reply[0] != ACK:
            raise IOError(f"No ACK from I2C address 0x{addr:02x} on write")
        if any(byte != ACK for byte in reply[1:]):
            raise IOError(
                f"I2C address 0x{addr:02x} NACKed a data byte "
                f"(status {reply.hex(' ')})"
            )

    def read(self, addr, count):
        """Read count bytes from a 7-bit address. Raises IOError on NACK."""
        if not 1 <= count <= MAX_XFER:
            raise ValueError(f"Read of {count} bytes outside 1..{MAX_XFER}")

        packet = [CMD_STREAM, CMD_STA, CMD_OUT | 1, (addr << 1) | 1]
        # All but the last byte are ACKed by us; the final bare CMD_IN NACKs to
        # tell the slave to stop.
        if count > 1:
            packet.append(CMD_IN | (count - 1))
        packet += [CMD_IN, CMD_STO, CMD_END]

        reply = self._xfer(packet, count + 1)
        if reply[0] != ACK:
            raise IOError(f"No ACK from I2C address 0x{addr:02x} on read")
        return reply[1:]

    def probe(self, addr):
        """Return True if a device ACKs its address. A NACK is not an error."""
        packet = [CMD_STREAM, CMD_STA, CMD_OUT | 1, addr << 1, CMD_STO, CMD_END]
        try:
            return self._xfer(packet, 1)[0] == ACK
        except IOError:
            return False

    def scan(self, first=0x08, last=0x77):
        """Return the addresses in [first, last] that respond."""
        return [addr for addr in range(first, last + 1) if self.probe(addr)]