#!/usr/bin/python3 -u
# Publish Lithionics Li3 BMS battery telemetry to Home Assistant via MQTT
# discovery. Protocol reverse-engineered from the com.lithionics.bms Android
# app (BLEMaster / MainBmsCsParameters): classic HM-10 BLE-UART (service
# ffe0, characteristic ffe1, notify+write, no pairing). On connect we send
# "$traceon" then "$info"; the device then streams CRLF-terminated CSV
# telemetry lines forever.
#
# BLE connect failures are retried internally (MQTT session stays up across
# them) rather than exiting, because bluetoothd on this Pi's onboard adapter
# occasionally wedges its discovery state after a run of failed connection
# attempts (Discovering stays "yes" forever, and every subsequent connect
# fails with le-connection-abort-by-local) — observed happening from our own
# repeated scan/connect cycling, not anything external. We restart
# bluetooth.service ourselves as soon as a scan fails to clear that, rate-
# limited (BLUETOOTH_RESTART_COOLDOWN_S) so we don't do it so often it
# disrupts the other BLE gear on this hub (motion sensors, IR remote, etc.
# also served by bluetoothd here).
#
# Only exits (letting systemd Restart=always give us a fresh process) on
# MQTT-level failure, which shouldn't happen in normal operation.

import asyncio
import json
import subprocess
import sys
import time

import paho.mqtt.client as mqtt
from bleak import BleakClient, BleakScanner

BLE_RETRY_DELAY_S = 8
STUCK_DISCOVERY_THRESHOLD = 1       # consecutive scan failures before we intervene
BLUETOOTH_RESTART_COOLDOWN_S = 300  # don't restart bluetooth.service more than this often

config_path = sys.argv[1] if len(sys.argv) > 1 else "/etc/van-li3/config.json"
with open(config_path) as f:
    cfg = json.load(f)

ADDR = cfg["ble_address"]
FFE1 = "0000ffe1-0000-1000-8000-00805f9b34fb"

broker = cfg["broker"]
mqtt_port = cfg.get("port", 1883)
mq_user = cfg["username"]
mq_pw = cfg["password"]
client_id = cfg.get("client_id", "van-li3-battery")

device_id = cfg.get("device_id", "li3_battery")
device_name = cfg.get("device_name", "Li3 Battery")

discovery_prefix = cfg.get("discovery_prefix", "homeassistant")
state_topic = cfg.get("state_topic", f"van/{device_id}/state")
publish_interval_s = cfg.get("publish_interval_s", 15)

SENSORS = [
    # (key, name, unit, device_class, display_precision) -- precision None means
    # a non-numeric value (hex code, version string, serial number): published
    # as-is with no unit, no rounding, no state_class.
    ("voltage", "Pack Voltage", "V", "voltage", 2),
    ("cell1_voltage", "Cell 1 Voltage", "V", "voltage", 2),
    ("cell2_voltage", "Cell 2 Voltage", "V", "voltage", 2),
    ("cell3_voltage", "Cell 3 Voltage", "V", "voltage", 2),
    ("cell4_voltage", "Cell 4 Voltage", "V", "voltage", 2),
    ("current", "Current", "A", "current", 2),
    ("soc", "State of Charge", "%", "battery", 0),
    ("bms_temperature", "BMS Temperature", "°F", "temperature", 1),
    ("battery_temperature", "Battery Temperature", "°F", "temperature", 1),
    # From "&" trace lines (streamed continuously once $traceon is sent) --
    # CAN-charger bus fields. This battery has no CAN charger wired up, so
    # can_charger_voltage/current are observed as a fixed sentinel and
    # can_charger_status/can_status never change; published anyway since
    # they're genuine decoded fields.
    ("remaining_capacity", "Remaining Capacity", "Ah", None, 0),
    ("remaining_time", "Remaining Time", "min", "duration", 0),
    ("can_charger_voltage", "CAN Charger Voltage", "V", "voltage", 1),
    ("can_charger_current", "CAN Charger Current", "A", "current", 1),
    ("can_charger_status", "CAN Charger Status", None, None, None),
    ("can_status", "CAN Status", None, None, None),
]

# From the "$" info line, sent once by the device right after $info and never
# repeated -- published separately (own retained state topic) rather than
# merged into the periodic telemetry above.
INFO_SENSORS = [
    ("total_consumed", "Lifetime Consumed", "Ah", None, 0),
    ("last_fault_code", "Last Fault Code", None, None, None),
    ("highest_recorded_temp", "Highest Recorded Temperature", "°F", "temperature", 0),
    ("lowest_recorded_temp", "Lowest Recorded Temperature", "°F", "temperature", 0),
    ("firmware_version", "Firmware Version", None, None, None),
    ("aging_factor_temp", "Aging Factor (Temp)", None, None, 0),
    ("aging_factor_soc", "Aging Factor (SOC)", None, None, 0),
    ("serial_number", "Serial Number", None, None, None),
]

info_topic = cfg.get("info_topic", f"van/{device_id}/info")

# The "status" field (main Cs telemetry line, and identically-coded but
# cumulative/latched "last_fault_code" from $info) is a 24-bit flag mask.
# Bit meanings pulled from the app's own "advanced" string-array resource
# (com.lithionics.bms base.apk, array/advanced -- dumped with aapt since the
# app's Kotlin StatusCodeTable class references stale/wrong resource IDs and
# can't be trusted). Index 0 = bit 23 (MSB) down to index 23 = bit 0 (LSB);
# blank entries are unused bits. Applied to "status" only -- last_fault_code
# is a lifetime latch (many bits accumulate over time) and isn't meaningfully
# summarized the same way.
STATUS_FLAGS = [
    "", "BMS Temp High", "Overcurrent State", "Charge OFF", "Aux Input State",
    "Cell Temp Low", "Cell Temp High", "AGSR State", "Temp Sensor Fault",
    "CAN Charger Fault", "CAN Charger Present", "AC Power Present",
    "Contactor Flutter", "Pre-Charge Fault", "Contactor Fault",
    "Contactor State", "Power Off State", "Battery Protection", "Low Voltage",
    "Reserve Range", "OptoLoop Open", "NeverDie Reserve", "Charge Detected",
    "High Voltage",
]


def decode_status(hex_code):
    try:
        value = int(hex_code, 16)
    except (ValueError, TypeError):
        return hex_code
    active = [
        label for i, label in enumerate(STATUS_FLAGS)
        if label and (value >> (23 - i)) & 1
    ]
    return ", ".join(active) if active else "OK"


def parse_line(line):
    """Cs-series main telemetry line (no prefix character)."""
    parts = line.split(",")
    try:
        f0 = int(parts[0])
    except ValueError:
        return None
    if not (101 <= f0 <= 9999):
        return None  # not a Cs-series main telemetry line
    try:
        return {
            "voltage": round(f0 * 0.01, 2),
            "cell1": round(int(parts[1]) * 0.01, 2),
            "cell2": round(int(parts[2]) * 0.01, 2),
            "cell3": round(int(parts[3]) * 0.01, 2),
            "cell4": round(int(parts[4]) * 0.01, 2),
            "bms_temp_f": int(parts[5]),
            "batt_temp_f": int(parts[6]),
            "current_a": int(parts[7]),
            "soc_pct": int(parts[8]),
            "status": parts[9] if len(parts) > 9 else "?",
        }
    except (ValueError, IndexError):
        return None


def parse_trace_line(line):
    """'&' trace line: &,batteryId,remaining,remainingTime,canChargerVoltage,
    canChargerCurrent,canChargerStatus,canStatus"""
    parts = line.split(",")
    try:
        return {
            "remaining_capacity": int(parts[2]),
            "remaining_time": int(parts[3]),
            "can_charger_voltage": round(int(parts[4]) * 0.1, 1),
            "can_charger_current": round(int(parts[5]) * 0.1, 1),
            "can_charger_status": parts[6],
            "can_status": parts[7],
        }
    except (ValueError, IndexError):
        return None


def parse_info_line(line):
    """'$' info line (response to $info): $,totalConsumed,lastFaultCode,
    highestRecordedTemp,lowestRecordedTemp,firmwareVersion,agingFactorTemp,
    agingFactorSoc,serialNumber"""
    parts = line.split(",")
    try:
        return {
            "total_consumed": int(parts[1]),
            "last_fault_code": parts[2],
            "highest_recorded_temp": int(parts[3]),
            "lowest_recorded_temp": int(parts[4]),
            "firmware_version": parts[5],
            "aging_factor_temp": int(parts[6]),
            "aging_factor_soc": int(parts[7]),
            "serial_number": parts[8],
        }
    except (ValueError, IndexError):
        return None


def _sensor_config(key, name, unit, device_class, precision, topic_state, device_info, expire_after=None):
    numeric = precision is not None
    payload = {
        "name": name,
        "unique_id": f"{device_id}_{key}",
        "state_topic": topic_state,
        "unit_of_measurement": unit,
        "device_class": device_class,
        "value_template": (
            f"{{{{ value_json.{key} | round({precision}) }}}}"
            if numeric
            else f"{{{{ value_json.{key} }}}}"
        ),
        "device": device_info,
    }
    if expire_after is not None:
        payload["expire_after"] = expire_after
    if numeric:
        payload["suggested_display_precision"] = precision
        payload["state_class"] = "measurement"
    return payload


def publish_discovery(mqc):
    device_info = {
        "identifiers": [device_id],
        "name": device_name,
        "manufacturer": "Lithionics",
        "model": "Li3 BMS",
    }
    for key, name, unit, device_class, precision in SENSORS:
        topic = f"{discovery_prefix}/sensor/{device_id}/{key}/config"
        payload = _sensor_config(
            key, name, unit, device_class, precision, state_topic, device_info,
            expire_after=publish_interval_s * 4,
        )
        mqc.publish(topic, json.dumps(payload), retain=True)

    status_topic = f"{discovery_prefix}/sensor/{device_id}/status/config"
    status_payload = {
        "name": "Status",
        "unique_id": f"{device_id}_status",
        "state_topic": state_topic,
        "value_template": "{{ value_json.status_text }}",
        "device": device_info,
        "expire_after": publish_interval_s * 4,
    }
    mqc.publish(status_topic, json.dumps(status_payload), retain=True)

    for key, name, unit, device_class, precision in INFO_SENSORS:
        topic = f"{discovery_prefix}/sensor/{device_id}/{key}/config"
        payload = _sensor_config(key, name, unit, device_class, precision, info_topic, device_info)
        payload["entity_category"] = "diagnostic"
        mqc.publish(topic, json.dumps(payload), retain=True)


def publish_info(mqc, info):
    """Publish the once-per-connection '$info' fields as their own retained
    message, separate from the periodic telemetry state."""
    mqc.publish(info_topic, json.dumps(info), retain=True)
    print("published info:", info)


def build_state_payload(cs_reading, trace_fields):
    # A non-zero status means the rest of the Cs line's fields are unreliable
    # (observed 2026-08-18: status '69' alongside e.g. current=341, soc=340,
    # cell4_voltage=34013.63) — publish only the status in that case. Trace
    # ("&" line) fields come from a separate message and are published
    # regardless.
    status = cs_reading["status"]
    if status != "000000":
        payload = {"status": status}
    else:
        payload = {
            "voltage": cs_reading["voltage"],
            "cell1_voltage": cs_reading["cell1"],
            "cell2_voltage": cs_reading["cell2"],
            "cell3_voltage": cs_reading["cell3"],
            "cell4_voltage": cs_reading["cell4"],
            "current": cs_reading["current_a"],
            "soc": cs_reading["soc_pct"],
            "bms_temperature": cs_reading["bms_temp_f"],
            "battery_temperature": cs_reading["batt_temp_f"],
            "status": status,
        }
    payload["status_text"] = decode_status(status)
    payload.update(trace_fields)
    return payload


def restart_bluetooth_service():
    print("scan failed; restarting bluetooth.service")
    subprocess.run(["systemctl", "restart", "bluetooth.service"], check=False)
    time.sleep(3)


async def find_device():
    """Scan/connect retry loop. Never gives up; self-heals a wedged
    bluetoothd discovery state along the way. Returns a found device."""
    consecutive_failures = 0
    last_bluetooth_restart = 0.0
    while True:
        print("scanning for device...")
        dev = await BleakScanner.find_device_by_address(ADDR, timeout=20)
        if dev:
            return dev
        consecutive_failures += 1
        print(f"device not found in scan (attempt {consecutive_failures})")
        if consecutive_failures >= STUCK_DISCOVERY_THRESHOLD:
            now = time.time()
            if (now - last_bluetooth_restart) >= BLUETOOTH_RESTART_COOLDOWN_S:
                restart_bluetooth_service()
                last_bluetooth_restart = now
            consecutive_failures = 0
        await asyncio.sleep(BLE_RETRY_DELAY_S)


async def stream_from_device(dev, mqc):
    """Connect to `dev` and publish readings until it disconnects."""
    buf = ""
    last_publish = 0.0
    latest_cs = {}
    latest_trace = {}
    info_published = False

    def notify_handler(_sender, data):
        nonlocal buf, info_published
        buf += data.decode("utf-8", errors="replace")
        while "\r\n" in buf:
            line, buf = buf.split("\r\n", 1)
            line = line.strip()
            if not line:
                continue
            if line.startswith("&"):
                trace = parse_trace_line(line)
                if trace:
                    latest_trace.update(trace)
            elif line.startswith("$"):
                if not info_published:
                    info = parse_info_line(line)
                    if info:
                        publish_info(mqc, info)
                        info_published = True
            else:
                reading = parse_line(line)
                if reading:
                    latest_cs.update(reading)

    async with BleakClient(dev, timeout=15) as client:
        print("connected")
        await client.start_notify(FFE1, notify_handler)
        await client.write_gatt_char(FFE1, b"$traceon\r\n", response=False)
        await asyncio.sleep(2)
        await client.write_gatt_char(FFE1, b"$info\r\n", response=False)

        while client.is_connected:
            await asyncio.sleep(1)
            now = time.time()
            if latest_cs and (now - last_publish) >= publish_interval_s:
                payload = build_state_payload(latest_cs, latest_trace)
                mqc.publish(state_topic, json.dumps(payload))
#                print("published:", payload)
                last_publish = now
    print("disconnected")


async def main():
    mqc = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2, client_id=client_id)
    mqc.username_pw_set(mq_user, mq_pw)
    mqc.connect(broker, mqtt_port, keepalive=60)
    mqc.loop_start()
    publish_discovery(mqc)

    while True:
        dev = await find_device()
        try:
            await stream_from_device(dev, mqc)
        except Exception as e:
            print("connection error:", e)
        await asyncio.sleep(BLE_RETRY_DELAY_S)


if __name__ == "__main__":
    asyncio.run(main())
