#!/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. After several
# consecutive failures we restart bluetooth.service ourselves to clear that,
# rate-limited 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 = 3       # 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")


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):
    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,
        "expire_after": publish_interval_s * 4,
    }
    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)
        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": "{{ 'OK' if value_json.status == '000000' else value_json.status }}",
        "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.
    if cs_reading["status"] != "000000":
        payload = {"status": cs_reading["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": cs_reading["status"],
        }
    payload.update(trace_fields)
    return payload


def restart_bluetooth_service():
    print("too many consecutive scan failures; 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
            print("RAW", line)
            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())
