#!/usr/bin/env python3
"""
I2C Doctor — USB/I2C Bridge host script.

Version: V1.2.0

Requires: pip install pyserial

Usage:
Connect the I2C Doctor to your computer via USB and put it in "USB/I2C BRIDGE" mode.
Then run the bridge.py script as either:
- Interactive CLI:
    python3 bridge.py --port /dev/tty.usbserial-* [--baud 115200]
or
- Library:
    from bridge import I2CBridge
    with I2CBridge('/dev/tty.usbserial-*') as b:
        print(b.scan())
        data = b.read_reg(0x48, 0x00, 2)
        b.write_reg(0x48, 0x01, bytes([0x60]))

Protocol (ASCII, newline-terminated):
    DIAGNOSE             -> full I2C bus diagnosis
    SCAN                 -> OK:3C 48\r\n (space-separated hex addresses, empty if none)
    PROBE XX             -> OK\r\n / ERR:NACK\r\n
    PROBE_READ XX        -> OK\r\n / ERR:NACK\r\n
    PROBE_READ XX RR     -> OK\r\n / ERR:NACK\r\n
    READ XX N            -> OK:AA BB\r\n (addr hex, count decimal — no register)
    READ XX RR N         -> OK:AA BB\r\n (addr hex, reg hex, count decimal)
    WRITE XX YY..        -> OK\r\n       (addr hex, data hex bytes — no register)
    READ_REG XX RR N     -> OK:AA BB\r\n (addr hex, reg hex, count decimal)
    WRITE_REG XX RR YY.. -> OK\r\n       (addr hex, reg hex, data hex bytes)
    SPEED                -> OK:SPEED=AAA kHz\r\n (speed decimal)
    SPEED N              -> OK\r\n       (speed decimal)
    RECOVER              -> OK\r\n / ERR:BUS_STUCK\r\n
    MONITOR              -> OK:MONITORING\r\n
    STRESS XX N          -> OK:TOTAL=AAA ACK=BBB NACK=CCC\r\n (addr hex, count decimal — no register)
    STRESS XX RR N       -> OK:TOTAL=AAA ACK=BBB NACK=CCC\r\n (addr hex, reg hex, count decimal)
    LONG_CABLE N         -> OK\r\n       (level 0-5)
    NOISE N              -> OK\r\n       (level 0-5)
    SDA_HOLD N           -> OK\r\n       (level 0-5)
    CLOCK_STRETCH N      -> OK\r\n       (level 0-5)
"""

import argparse
import time
import serial
import json
import re
from typing import Iterator

class I2CBridge:
    """USB to I2C bridge for the I2C Doctor device."""

    def __init__(self, port: str, baud: int = 115200, timeout: float = 5.0):
        self.ser = serial.Serial(port, baud, timeout=timeout)
        time.sleep(0.2)
        self.ser.reset_input_buffer()

    def _readline(self) -> str:
        """Reads a single line from serial, stripped of trailing whitespace."""
        return self.ser.readline().decode(errors='replace').strip()

    def _cmd(self, cmd: str) -> str:
        self.ser.reset_input_buffer()
        self.ser.write((cmd + '\n').encode())
        return self._readline()

    def scan(self) -> list:
        """Scan the I2C bus. Returns list of found 7-bit addresses (int)."""
        resp = self._cmd('SCAN')
        if not resp.startswith('OK:'):
            raise IOError(f'SCAN failed: {resp!r}')
        part = resp[3:].strip()
        return [int(x, 16) for x in part.split()] if part else []

    def probe(self, addr: int) -> bool:
        """Probe address `addr` (7-bit int). Returns True if ACK received."""
        resp = self._cmd(f'PROBE {addr:02X}')
        return resp == 'OK'

    def probe_read(self, addr: int, reg: int = None) -> bool:
        """Probe read address `addr` (7-bit int) and `reg` if present. Returns True if ACK received."""
        cmd_str = f"PROBE_READ {addr:02X} {reg:02X}" if reg is not None else f"PROBE_READ {addr:02X}"
        resp = self._cmd(cmd_str)
        return resp == 'OK'

    def read(self, addr: int, length: int) -> bytes:
        """Read `length` bytes from device at `addr` (7-bit) without register address."""
        resp = self._cmd(f'READ {addr:02X} {length}')
        if not resp.startswith('OK:'):
            raise IOError(f'READ 0x{addr:02X} failed: {resp!r}')
        return bytes(int(x, 16) for x in resp[3:].split())

    def write(self, addr: int, data: bytes) -> None:
        """Write `data` to device at `addr` (7-bit) without register address."""
        hex_bytes = ' '.join(f'{b:02X}' for b in data)
        resp = self._cmd(f'WRITE {addr:02X} {hex_bytes}')
        if resp != 'OK':
            raise IOError(f'WRITE 0x{addr:02X} failed: {resp!r}')

    def read_reg(self, addr: int, reg: int, length: int) -> bytes:
        """Read `length` bytes from register `reg` of device at `addr` (7-bit)."""
        resp = self._cmd(f'READ_REG {addr:02X} {reg:02X} {length}')
        if not resp.startswith('OK:'):
            raise IOError(f'READ_REG 0x{addr:02X} reg 0x{reg:02X} failed: {resp!r}')
        return bytes(int(x, 16) for x in resp[3:].split())

    def write_reg(self, addr: int, reg: int, data: bytes) -> None:
        """Write `data` to register `reg` of device at `addr` (7-bit)."""
        hex_bytes = ' '.join(f'{b:02X}' for b in data)
        resp = self._cmd(f'WRITE_REG {addr:02X} {reg:02X} {hex_bytes}')
        if resp != 'OK':
            raise IOError(f'WRITE_REG 0x{addr:02X} reg 0x{reg:02X} failed: {resp!r}')

    def speed(self, speed: int = None) -> int:
        """Get or set I2C bus speed in kHz."""
        cmd_str = f"SPEED {speed}" if speed is not None else "SPEED"
        resp = self._cmd(cmd_str)
        if not resp.startswith('OK'):
            raise IOError(f'SPEED command failed: {resp!r}')
        if resp.startswith('OK:'):
            raw_val = resp[3:].strip().replace('SPEED=', '').replace(' kHz', '')
            return int(raw_val)
        return speed

    def diagnose(self) -> dict:
        """
        Run full hardware diagnosis. 
        Returns a dict containing VDD, voltages, currents, pullups, rise times, 
        max speed, shorts detection, score, status, and issues.
        """
        resp = self._cmd('DIAGNOSE')
        if not resp.startswith('OK:'):
            raise IOError(f'DIAGNOSE failed: {resp!r}')
        raw_data = resp[3:].strip()
        result = {}
        pattern = re.compile(r'(\w+)=(?:"([^"]*)"|(\S+))')
        for match in pattern.finditer(raw_data):
            key = match.group(1).lower()
            val = match.group(2) if match.group(2) is not None else match.group(3)
            if val.upper() in ('TRUE', 'DETECTED'):
                val = True
            elif val.upper() in ('FALSE', 'NONE'):
                val = False if val.upper() == 'FALSE' else val
            else:
                try:
                    val = int(val)
                except ValueError:
                    try:
                        val = float(val.rstrip('kKVvAms'))
                    except ValueError:
                        pass
            result[key] = val
        return result

    def diagnose_json(self) -> dict:
        """Runs diagnosis and returns result directly parsed via JSON."""
        resp = self._cmd('DIAGNOSE_JSON')
        if not resp.startswith('OK:'):
            raise IOError(f'DIAGNOSE_JSON failed: {resp!r}')
        json_raw = resp[3:].strip()
        return json.loads(json_raw)

    def bus_recovery(self) -> bool:
        """Sends 9 to 36 clock pulses on SCL to recovery hanging I2C slaves."""
        resp = self._cmd('BUS_RECOVERY')
        if resp.startswith('OK'):
            return True
        elif resp.startswith('ERR:'):
            return False
        raise IOError(f'BUS_RECOVERY failed: {resp!r}')

    def monitor(self) -> Iterator[str]:
        """
        Starts live I2C bus monitoring.
        Yields incoming status/data lines as strings continuously.
        Use stop_monitor() or send a byte to break out of monitoring mode on the hardware.
        """
        header = self._cmd('MONITOR')
        if not header.startswith("OK:MONITORING"):
            raise IOError(f"MONITOR failed to start: {header!r}")
        yield header
        while True:
            line = self._readline()
            if line:
                yield line
                if line.startswith("OK:STOPPED"):
                    break

    def stop_monitor(self) -> None:
        """Sends a newline to break the monitor loop on the hardware and flushes buffer."""
        self.ser.write(b'\n')
        time.sleep(0.1)
        self.ser.reset_input_buffer()

    def stress_test(self, addr: int, count: int, reg: int = None) -> dict:
        """
        Runs a high-speed stress test on the target address.
        Optionally reads from a specific register.
        Returns a dict with 'total', 'ack', 'nack', and 'pass_rate'.
        """
        if reg is not None:
            cmd_str = f"STRESS_TEST {addr:02X} {reg:02X} {count}"
        else:
            cmd_str = f"STRESS_TEST {addr:02X} {count}"
        dt = round(time.time() * 1000)
        calc_timeout = 5.0 * (count / 1000.0) + 5.0
        old_timeout = self.ser.timeout
        self.ser.timeout = calc_timeout
        try:
            resp = self._cmd(cmd_str)
        finally:
            self.ser.timeout = old_timeout
        if not resp.startswith('OK:'):
            raise IOError(f'STRESS_TEST failed: {resp!r}')
        raw_data = resp[3:].strip()
        parts = dict(pair.split('=') for pair in raw_data.split())
        total = int(parts.get('TOTAL', count))
        ack = int(parts.get('ACK', 0))
        nack = int(parts.get('NACK', 0))
        pass_rate = (ack / total * 100.0) if total > 0 else 0.0
        return {
            'total': total,
            'ack': ack,
            'nack': nack,
            'pass_rate': pass_rate,
            'dt': (round(time.time() * 1000) - dt)
        }

    def long_cable(self, level: int) -> str:
        """Sets long cable fault injection level (0-5)."""
        return self._cmd(f"LONG_CABLE {level}")

    def noise(self, level: int) -> str:
        """Sets noise fault injection level (0-5)."""
        return self._cmd(f"NOISE {level}")

    def sda_hold(self, level: int) -> str:
        """Sets sda_hold fault injection level (0-5)."""
        return self._cmd(f"SDA_HOLD {level}")

    def clock_stretch(self, level: int) -> str:
        """Sets clock_stretch fault injection level (0-5)."""
        return self._cmd(f"CLOCK_STRETCH {level}")

    def close(self) -> None:
        self.ser.close()

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()

def _interactive(bridge: I2CBridge) -> None:
    try:
        import readline  # enables arrow-key history on macOS/Linux
        readline.set_history_length(100)
    except ImportError:
        pass
    print("I2C Doctor Bridge — ready. Commands: diagnose, scan, probe, read, write, read_reg, write_reg, monitor, help, quit")
    while True:
        try:
            line = input('i2c> ').strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not line:
            continue
        parts = line.split()
        cmd = parts[0].lower()
        try:
            if cmd in ('quit', 'exit', 'q'):
                break
            elif cmd == 'help':
                print('  diagnose                                run full I2C bus diagnosis (key-value text)')
                print('  diagnose_json                           run full I2C bus diagnosis (raw JSON format)')
                print('  scan                                    scan bus, list found addresses')
                print('  probe         <addr_hex>                check if address ACKs')
                print('  probe_read    <addr_hex>                check if reading address ACKs')
                print('  probe_read    <addr_hex> <reg_hex>      check if reading register ACKs')
                print('  read          <addr_hex> <n>            read n bytes (no register addr)')
                print('  read          <addr_hex> <reg_hex> <n>  read n bytes from register')
                print('  read_reg      <addr_hex> <reg_hex> <n>  read n bytes from register')
                print('  write         <addr_hex> <..>           write bytes (no register addr)')
                print('  write_reg     <addr_hex> <reg_hex> <..> write bytes to register')
                print('  speed                                   get speed in kbit/s (10, 50, 100, 200, 400, 1000)')
                print('  speed         <kbit/s>                  set speed in kbit/s (10, 50, 100, 200, 400, 1000)')
                print('  recover                                 clock SCL 9 to 36 times to free stuck SDA line')
                print('  monitor                                 logs bus activity in real-time')
                print('  stress        <addr_hex> <n>            run high-speed read stress test (no register addr)')
                print('  stress        <addr_hex> <reg_hex> <n>  run high-speed read stress test')
                print('  long_cable    <n>                       set long cable fault injection level (0-5)')
                print('  noise         <n>                       set noise fault injection level (0-5)')
                print('  sda_hold      <n>                       set SDA hold fault injection level (0-5)')
                print('  clock_stretch <n>                       set clock stretch fault injection level (0-5)')
            elif cmd == 'scan':
                addrs = bridge.scan()
                if addrs:
                    print('Found:', ', '.join(f'0x{a:02X}' for a in addrs))
                else:
                    print('No devices found')
            elif cmd == 'probe' and len(parts) == 2:
                if bridge.probe(int(parts[1], 16)):
                    print("ACK")
                else:
                    print("ERR:NACK")
            elif cmd == 'probe_read' and len(parts) == 2:
                if bridge.probe_read(int(parts[1], 16)):
                    print("ACK")
                else:
                    print("ERR:NACK")
            elif cmd == 'probe_read' and len(parts) == 3:
                if bridge.probe_read(int(parts[1], 16), int(parts[2], 16)):
                    print("ACK")
                else:
                    print("ERR:NACK")
            elif cmd == 'read' and len(parts) == 4:
                data = bridge.read_reg(int(parts[1], 16), int(parts[2], 16), int(parts[3]))
                print(' '.join(f'{b:02X}' for b in data))
            elif cmd == 'read' and len(parts) == 3:
                data = bridge.read(int(parts[1], 16), int(parts[2]))
                print(' '.join(f'{b:02X}' for b in data))
            elif cmd == 'read_reg' and len(parts) == 4:
                data = bridge.read_reg(int(parts[1], 16), int(parts[2], 16), int(parts[3]))
                print(' '.join(f'{b:02X}' for b in data))
            elif cmd == 'write' and len(parts) >= 3:
                bridge.write(int(parts[1], 16), bytes(int(x, 16) for x in parts[2:]))
                print('OK')
            elif cmd == 'write_reg' and len(parts) >= 4:
                bridge.write_reg(int(parts[1], 16), int(parts[2], 16), bytes(int(x, 16) for x in parts[3:]))
                print('OK')
            elif cmd == 'speed' and len(parts) == 2:
                bridge.speed(int(parts[1], 10))
                print('OK')
            elif cmd == 'speed' and len(parts) == 1:
                current_spd = bridge.speed()
                print(f"{current_spd} kHz")
            elif cmd == 'diagnose':
                diag = bridge.diagnose()
                print("=== I2C DOCTOR DIAGNOSTICS ===")
                print(f"  VDD User:     {diag.get('vdd', 'N/A')} V")
                print(f"  Bus Voltage:  SDA: {diag.get('v_sda', 'N/A')} V | SCL: {diag.get('v_scl', 'N/A')} V")
                print(f"  Current:      SDA: {diag.get('i_sda', 'N/A')} mA | SCL: {diag.get('i_scl', 'N/A')} mA")
                print(f"  Pullups:      SDA: {diag.get('r_sda', 'N/A')} kΩ | SCL: {diag.get('r_scl', 'N/A')} kΩ ({diag.get('r_stat', '')})")
                print(f"  Rise Time:    SDA: {diag.get('tr_sda', 'N/A')} ns | SCL: {diag.get('tr_scl', 'N/A')} ns")
                print(f"  Speed:        SDA: {diag.get('f_sda', 'N/A')} kHz | SCL: {diag.get('f_scl', 'N/A')} kHz")
                print(f"  Max Speed:    {diag.get('max_spd', 'N/A')} kHz")
                print(f"  Shorts:       {'DETECTED!' if diag.get('shorts') is True else 'None'}")
                print(f"  Found Devs:   {diag.get('devs', 0)}")
                print(f"  Bus Score:    {diag.get('score', 'N/A')}% [{diag.get('status', 'N/A')}]")
                if diag.get('issues') and diag.get('issues') != 'None':
                    print(f"  Issues:       {diag.get('issues')}")
                if diag.get('recmdt') and diag.get('recmdt') != 'None':
                    print(f"  Recs:         {diag.get('recmdt')}")
                print("==============================\n")
            elif cmd == 'diagnose_json':
                diag = bridge.diagnose_json()
                print(json.dumps(diag, indent=2))
            elif cmd == 'bus_recovery' or cmd == 'recover':
                if bridge.bus_recovery():
                    print("OK")
                else:
                    print("ERR:BUS_STUCK")
            elif cmd == 'monitor':
                print("Press Ctrl+C to stop monitoring...\n")
                try:
                    for line in bridge.monitor():
                        print(line)
                except KeyboardInterrupt:
                    print("\nStopping monitor...")
                finally:
                    bridge.stop_monitor()
            elif ((cmd == 'stress_test' or cmd == 'stress') and 3 <= len(parts) <= 4):
                addr = int(parts[1], 16)
                if len(parts) == 4:
                    reg = int(parts[2], 16)
                    count = int(parts[3])
                else:
                    reg = None
                    count = int(parts[2])
                print(f"Running stress test on 0x{addr:02X} ({count} cycles)...")
                res = bridge.stress_test(addr, count, reg)
                print("=== STRESS TEST RESULTS ===")
                print(f"  Total Cycles: {res['total']}")
                print(f"  ACK (Pass):   {res['ack']}")
                print(f"  NACK (Fail):  {res['nack']}")
                print(f"  Pass Rate:    {res['pass_rate']:.2f}%")
                print(f"  Delta Time:   {res['dt']} ms")
                print("===========================\n")
            elif cmd == 'long_cable' and len(parts) == 2:
                print(bridge.long_cable(int(parts[1])))
            elif cmd == 'noise' and len(parts) == 2:
                print(bridge.noise(int(parts[1])))
            elif cmd == 'sda_hold' and len(parts) == 2:
                print(bridge.sda_hold(int(parts[1])))
            elif cmd == 'clock_stretch' and len(parts) == 2:
                print(bridge.clock_stretch(int(parts[1])))
            else:
                print(f"Unknown or malformed command: '{line}'. Type 'help'.")
        except (IOError, ValueError) as e:
            print(f'Error: {e}')

def main() -> None:
    parser = argparse.ArgumentParser(description='I2C Doctor USB Bridge')
    parser.add_argument('--port', required=True, help='Serial port (e.g. /dev/tty.usbserial-*, /dev/ttyUSB0, COM3)')
    parser.add_argument('--baud', type=int, default=115200, help='Baud rate (default: 115200)')
    args = parser.parse_args()
    with I2CBridge(args.port, args.baud) as bridge:
        _interactive(bridge)

if __name__ == '__main__':
    main()