"""Host library for the ATmega8U2 USB-I2C bridge + TCA9539 helper.  pip install pyserial"""
import struct
import serial

MAGIC = 0xA5
OP_XFER, OP_SPEED, OP_PING, OP_RECOVER = 1, 2, 3, 4
MAXLEN = 32
ERRORS = {1: "NACK on address", 2: "NACK on data", 3: "bus error / timeout",
          0x10: "unknown op", 0x11: "bad length"}


class I2CError(IOError):
    def __init__(self, status):
        super().__init__(f"I2C error {status}: {ERRORS.get(status, '?')}")
        self.status = status


class Bridge:
    def __init__(self, port, timeout=0.5):
        self.s = serial.Serial(port, 115200, timeout=timeout)  # baud is ignored by CDC
        self.s.reset_input_buffer()
        self.ping()

    def close(self):
        self.s.close()

    def _read(self, n):
        d = self.s.read(n)
        if len(d) != n:
            raise TimeoutError("no reply from bridge")
        return d

    def ping(self):
        self.s.write(bytes([MAGIC, OP_PING]))
        st, proto, maxlen = self._read(3)
        if st:
            raise I2CError(st)
        return proto, maxlen

    def set_delay(self, n):
        """SCL delay; 0 = 400 kHz (default), each step adds 0.375 us per period."""
        self.s.write(bytes([MAGIC, OP_SPEED, n]))
        st, = self._read(1)
        if st:
            raise I2CError(st)

    def recover(self):
        self.s.write(bytes([MAGIC, OP_RECOVER]))
        st, = self._read(1)
        if st:
            raise I2CError(st)

    @staticmethod
    def _frame(addr, w, rlen):
        w = bytes(w)
        if len(w) > MAXLEN or rlen > MAXLEN:
            raise ValueError("max 32 bytes per phase")
        return bytes([MAGIC, OP_XFER, addr, len(w), rlen]) + w

    def xfer(self, addr, w=b"", rlen=0):
        """One transaction: write w, then (repeated START) read rlen bytes."""
        return self.batch([(addr, w, rlen)])[0]

    def batch(self, ops, raise_on_error=True):
        """Pipeline many (addr, wdata, rlen) transactions in as few USB packets as possible.
        This is the fast path: one USB round trip (~1 ms) for dozens of transactions."""
        ops = list(ops)
        self.s.write(b"".join(self._frame(*o) for o in ops))
        out, first_err = [], 0
        for _, _, rlen in ops:                    # always drain every reply to stay in sync
            st, = self._read(1)
            if st:
                first_err = first_err or st
                out.append(None)
            else:
                out.append(self._read(rlen) if rlen else b"")
        if first_err and raise_on_error:
            raise I2CError(first_err)
        return out

    def scan(self):
        res = self.batch([(a, b"", 0) for a in range(8, 120)], raise_on_error=False)
        return [a for a, r in zip(range(8, 120), res) if r is not None]


class TCA9539:
    IN0, OUT0, POL0, CFG0 = 0x00, 0x02, 0x04, 0x06

    def __init__(self, bridge, addr=0x74):       # 0x74..0x77 from A1/A0
        self.b, self.addr = bridge, addr

    def _w16(self, reg, v):
        self.b.xfer(self.addr, [reg, v & 0xFF, v >> 8])

    def _r16(self, reg):
        lo, hi = self.b.xfer(self.addr, [reg], 2)
        return lo | hi << 8

    def set_direction(self, input_mask):          # bit=1 -> input (power-on default 0xFFFF)
        self._w16(self.CFG0, input_mask)

    def write(self, value):
        self._w16(self.OUT0, value)

    def read(self):
        return self._r16(self.IN0)

    def set_polarity(self, mask):
        self._w16(self.POL0, mask)


if __name__ == "__main__":
    import sys, time
    b = Bridge(sys.argv[1])
    print("proto/maxlen:", b.ping())
    found = b.scan()
    print("devices:", [hex(a) for a in found])
    if found:
        io = TCA9539(b, found[0])
        io.set_direction(0xFF00)                   # port0 outputs, port1 inputs
        for i in range(8):
            io.write(1 << i)
            time.sleep(0.1)
        io.write(0)
        print("port1 inputs: %02X" % (io.read() >> 8))
        t = time.time()
        for _ in range(200):
            io.write(0x55)                          # unbatched: bound by USB 1 ms polling
        print("single writes/s: %.0f" % (200 / (time.time() - t)))
        t = time.time()
        b.batch([(found[0], [TCA9539.OUT0, 0x55, 0], 0)] * 400)
        print("batched writes/s: %.0f" % (400 / (time.time() - t)))