#!/usr/bin/env python3

import struct

import smbus2

import rclpy
from rclpy.node import Node
from leo_msgs.msg import ChargingMonitorData


CONFIG_VOLTAGE_RANGE_32V = 0b1  # set bus voltage range to 32V (default)
CONFIG_GAIN_DIV_8_320MV = 0b11  # shunt prog. gain set to /8, 320 mV range
CONFIG_ADCRES_12BIT_128S = 0b1111  # 12bit, 128 samples, 68.10ms
CONFIG_SANDBVOLT_CONTINUOUS = 0x07  # shunt and bus voltage continuous
CONFIG_REG_VALUE = (
    (CONFIG_VOLTAGE_RANGE_32V << 13)
    | (CONFIG_GAIN_DIV_8_320MV << 11)
    | (CONFIG_ADCRES_12BIT_128S << 7)
    | (CONFIG_ADCRES_12BIT_128S << 3)
    | CONFIG_SANDBVOLT_CONTINUOUS
)  # Value to write to config register

# Define LSB values
SHUNT_VOLTAGE_LSB = 0.00001  # 10uV per bit
BUS_VOLTAGE_LSB = 0.004  # 4mV per bit

# Calibration value calculation
RSHUNT = 0.1  # (Resistor value in ohms)
# VBUS_MAX = 32V             (Assumes 32V)
# VSHUNT_MAX = 0.32          (Assumes Gain 8, 320mV)
# MaxPossible_I = VSHUNT_MAX / RSHUNT
# MaxPossible_I = 3.2 A
# MaxExpected_I = 3 A
# MaximumLSB = MaxPossible_I/32768
# MaximumLSB = 0.000098              (98uA per bit)
# MinimumLSB = MaxExpected_I/32768
# MinimumLSB = 0.000091              (91uA per bit)
# We choose LSB between the min and max
# We chose 95uA per bit
# Target LSB (~95uA/bit), then force calibration register's LSB to 0 (value must be even)
CURRENT_LSB_TARGET = 0.000095
CALIBRATION_REG_VALUE = int(0.04096 / (CURRENT_LSB_TARGET * RSHUNT)) & 0xFFFE

# Recompute actual LSBs from the final calibration value
CURRENT_LSB = 0.04096 / (CALIBRATION_REG_VALUE * RSHUNT)
POWER_LSB = CURRENT_LSB * 20

REG_CONFIG = 0x00
REG_SHUNTVOLTAGE = 0x01
REG_BUSVOLTAGE = 0x02
REG_POWER = 0x03
REG_CURRENT = 0x04
REG_CALIBRATION = 0x05


class ChargingMonitor(Node):
    def __init__(self):
        super().__init__("charging_monitor")

        self.declare_parameter("i2c_bus_nr", 1)
        self.declare_parameter("sensor_address", 0x40)
        self.declare_parameter("data_publish_rate", 10.0)

        self.i2c_bus_nr = (
            self.get_parameter("i2c_bus_nr").get_parameter_value().integer_value
        )
        self.sensor_address = (
            self.get_parameter("sensor_address").get_parameter_value().integer_value
        )
        self.data_publish_rate = (
            self.get_parameter("data_publish_rate").get_parameter_value().double_value
        )

        if self.data_publish_rate <= 0:
            raise ValueError("data_publish_rate must be positive")

        self.bus = smbus2.SMBus()

        self.publisher = self.create_publisher(ChargingMonitorData, "~/data", 10)

    def _write_reg16(self, register: int, value: int) -> None:
        # INA219 expects big-endian (MSB first). Use block write to control byte order.
        data = struct.pack(">H", value & 0xFFFF)
        self.bus.write_i2c_block_data(self.sensor_address, register, list(data))

    def _read_reg16_u(self, register: int) -> int:
        data = self.bus.read_i2c_block_data(self.sensor_address, register, 2)
        return struct.unpack(">H", bytes(data))[0]

    def _read_reg16_s(self, register: int) -> int:
        data = self.bus.read_i2c_block_data(self.sensor_address, register, 2)
        return struct.unpack(">h", bytes(data))[0]

    def initialize_sensor(self):
        # Connect to I2C bus
        self.bus.open(self.i2c_bus_nr)

        # Write configuration to INA219
        self._write_reg16(REG_CONFIG, CONFIG_REG_VALUE)

        # Write calibration value to INA219
        self._write_reg16(REG_CALIBRATION, CALIBRATION_REG_VALUE)

        self.get_logger().info(
            "Charging monitor sensor initialized, starting data publishing..."
        )

        self.publish_timer = self.create_timer(
            1.0 / self.data_publish_rate, self.publish_data
        )

    def ensure_sensor_ready(self) -> bool:
        try:
            config = self._read_reg16_u(REG_CONFIG)
            calib = self._read_reg16_u(REG_CALIBRATION)
        except OSError as e:
            self.get_logger().warning(f"I2C not responding, will retry: {e}")
            return False

        needs_reinit = (config != CONFIG_REG_VALUE) or (calib != CALIBRATION_REG_VALUE)
        if needs_reinit:
            self.get_logger().warning(
                "INA219 configuration mismatch detected (possible reset or corruption). "
            )
            try:
                self._write_reg16(REG_CONFIG, CONFIG_REG_VALUE)
                self._write_reg16(REG_CALIBRATION, CALIBRATION_REG_VALUE)
            except OSError as e:
                self.get_logger().error(
                    f"I2C error while reapplying config/calibration: {e}"
                )
                return False

        return True

    def publish_data(self):
        if not self.ensure_sensor_ready():
            return

        msg = ChargingMonitorData()
        msg.header.stamp = self.get_clock().now().to_msg()

        try:
            # Read shunt voltage
            shunt_voltage_raw = self._read_reg16_s(REG_SHUNTVOLTAGE)
            msg.shunt_voltage = shunt_voltage_raw * SHUNT_VOLTAGE_LSB

            # Read bus voltage
            bus_voltage_reg = self._read_reg16_u(REG_BUSVOLTAGE)
            # Bus voltage register uses bits 15..3; shift right by 3 before scaling
            msg.bus_voltage = (bus_voltage_reg >> 3) * BUS_VOLTAGE_LSB

            # Read current
            current_raw = self._read_reg16_s(REG_CURRENT)
            msg.current = current_raw * CURRENT_LSB

            # Read power
            power_raw = self._read_reg16_u(REG_POWER)
            msg.power = power_raw * POWER_LSB
        except OSError as e:
            self.get_logger().error(
                f"I2C communication error while reading charging monitor data: {e}"
            )
            return
        except Exception as e:
            self.get_logger().error(
                f"Unexpected error while reading charging monitor data: {e}"
            )
            return

        self.publisher.publish(msg)


if __name__ == "__main__":
    rclpy.init()
    charging_monitor = ChargingMonitor()

    try:
        charging_monitor.initialize_sensor()
    except Exception as e:
        charging_monitor.get_logger().error(
            f"Failed to initialize Charging monitor: {e}"
        )
        charging_monitor.bus.close()
        rclpy.shutdown()
        exit(1)

    try:
        rclpy.spin(charging_monitor)
    except KeyboardInterrupt:
        pass
    finally:
        charging_monitor.bus.close()
