#!/usr/bin/env python3
"""Bench-test an INA219 I2C voltage/current sensor: live readout loop.

Talks to the chip directly over smbus2 (already a project dependency for
the BME280) -- no extra libraries. Current is computed from the shunt
voltage drop and the known shunt resistance, skipping the chip's internal
calibration-register machinery: at bench-monitoring precision the simple
division is equivalent and much harder to get wrong.

Typical breakout boards carry a 0.1 ohm shunt (marked "R100") giving a
+/-3.2A range. If yours is marked R050 (0.05) or R010 (0.01), pass
--shunt-ohms accordingly or every current reading will be off by that
factor.

Run:
  python3 ina219_test.py                 # address 0x40, 0.1 ohm shunt
  python3 ina219_test.py --addr 0x41     # if A0 jumper is bridged
  python3 ina219_test.py --interval 0.2  # faster readout
"""

from __future__ import annotations

import argparse
import sys
import time

# INA219 register map
_REG_CONFIG = 0x00
_REG_SHUNT_VOLTAGE = 0x01  # signed, LSB = 10 uV
_REG_BUS_VOLTAGE = 0x02    # bits 15..3, LSB = 4 mV

# 32V bus range, +/-320mV shunt gain, 12-bit ADCs, continuous shunt+bus.
# (This is also the chip's power-on default; written explicitly so a
# previously misconfigured chip can't skew readings.)
_CONFIG_VALUE = 0x399F


class INA219:
    def __init__(self, bus: int = 1, address: int = 0x40, shunt_ohms: float = 0.1):
        import smbus2

        self._bus = smbus2.SMBus(bus)
        self._address = address
        self._shunt_ohms = shunt_ohms
        self._write_register(_REG_CONFIG, _CONFIG_VALUE)

    def _write_register(self, reg: int, value: int) -> None:
        # INA219 is big-endian; SMBus word ops are little-endian, so use
        # explicit byte blocks instead of read/write_word_data.
        self._bus.write_i2c_block_data(self._address, reg, [(value >> 8) & 0xFF, value & 0xFF])

    def _read_register(self, reg: int) -> int:
        hi, lo = self._bus.read_i2c_block_data(self._address, reg, 2)
        return (hi << 8) | lo

    def read(self) -> dict:
        raw_shunt = self._read_register(_REG_SHUNT_VOLTAGE)
        if raw_shunt > 0x7FFF:  # sign-extend
            raw_shunt -= 0x10000
        shunt_v = raw_shunt * 10e-6

        raw_bus = self._read_register(_REG_BUS_VOLTAGE)
        overflow = bool(raw_bus & 0x1)
        bus_v = (raw_bus >> 3) * 0.004

        current_a = shunt_v / self._shunt_ohms
        return {
            "bus_v": bus_v,                       # at VIN- relative to GND
            "supply_v": bus_v + shunt_v,          # at VIN+ (adds shunt drop back)
            "shunt_mv": shunt_v * 1000,
            "current_a": current_a,
            "power_w": bus_v * current_a,
            "overflow": overflow,
        }


def main() -> int:
    parser = argparse.ArgumentParser(description="INA219 live voltage/current readout")
    parser.add_argument("--bus", type=int, default=1, help="I2C bus (default 1)")
    parser.add_argument("--addr", type=lambda s: int(s, 0), default=0x40,
                        help="I2C address (default 0x40; 0x41/0x44/0x45 via A0/A1 jumpers)")
    parser.add_argument("--shunt-ohms", type=float, default=0.1,
                        help="shunt resistance (default 0.1 = the usual R100 part)")
    parser.add_argument("--interval", type=float, default=1.0, help="seconds between readings")
    args = parser.parse_args()

    try:
        sensor = INA219(bus=args.bus, address=args.addr, shunt_ohms=args.shunt_ohms)
    except Exception as exc:
        print(f"INA219 not reachable at {args.addr:#x} on bus {args.bus}: {exc}")
        print("Check wiring and run: sudo i2cdetect -y 1   (expect '40' in the grid)")
        return 1

    print(f"INA219 at {args.addr:#x}, shunt {args.shunt_ohms} ohm -- Ctrl+C to stop")
    print(f"{'bus V':>8} {'shunt mV':>10} {'current A':>10} {'power W':>9}")
    try:
        while True:
            r = sensor.read()
            flag = "  OVERFLOW" if r["overflow"] else ""
            print(f"{r['bus_v']:8.3f} {r['shunt_mv']:10.2f} {r['current_a']:10.3f} "
                  f"{r['power_w']:9.2f}{flag}")
            time.sleep(args.interval)
    except KeyboardInterrupt:
        print()
    return 0


if __name__ == "__main__":
    sys.exit(main())