#!/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, value_template_field, display_precision)
    ("voltage", "Pack Voltage", "V", "voltage", "voltage", 2),
    ("cell1_voltage", "Cell 1 Voltage", "V", "voltage", "cell1", 2),
    ("cell2_voltage", "Cell 2 Voltage", "V", "voltage", "cell2", 2),
    ("cell3_voltage", "Cell 3 Voltage", "V", "voltage", "cell3", 2),
    ("cell4_voltage", "Cell 4 Voltage", "V", "voltage", "cell4", 2),
    ("current", "Current", "A", "current", "current_a", 2),
    ("soc", "State of Charge", "%", "battery", "soc_pct", 0),
    ("bms_temperature", "BMS Temperature", "°F", "temperature", "bms_temp_f", 1),
    ("battery_temperature", "Battery Temperature", "°F", "temperature", "batt_temp_f", 1),
]


def parse_line(line):
    line = line.strip()
    if not line or line.startswith("&") or line.startswith("$"):
        return None
    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 publish_discovery(mqc):
    device_info = {
        "identifiers": [device_id],
        "name": device_name,
        "manufacturer": "Lithionics",
        "model": "Li3 BMS",
    }
    for key, name, unit, device_class, _field, precision in SENSORS:
        topic = f"{discovery_prefix}/sensor/{device_id}/{key}/config"
        payload = {
            "name": name,
            "unique_id": f"{device_id}_{key}",
            "state_topic": state_topic,
            "unit_of_measurement": unit,
            "device_class": device_class,
            "value_template": f"{{{{ value_json.{key} | round({precision}) }}}}",
            "suggested_display_precision": precision,
            "device": device_info,
            "expire_after": publish_interval_s * 4,
            "state_class": "measurement",
        }
        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)


def build_state_payload(reading):
    return {
        "voltage": reading["voltage"],
        "cell1_voltage": reading["cell1"],
        "cell2_voltage": reading["cell2"],
        "cell3_voltage": reading["cell3"],
        "cell4_voltage": reading["cell4"],
        "current": reading["current_a"],
        "soc": reading["soc_pct"],
        "bms_temperature": reading["bms_temp_f"],
        "battery_temperature": reading["batt_temp_f"],
        "status": reading["status"],
    }


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 = {}

    def notify_handler(_sender, data):
        nonlocal buf
        buf += data.decode("utf-8", errors="replace")
        while "\r\n" in buf:
            line, buf = buf.split("\r\n", 1)
            reading = parse_line(line)
            if reading:
                latest.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 and (now - last_publish) >= publish_interval_s:
                payload = build_state_payload(latest)
                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())