"""
gaugemqttcontinuous.py — MQTT-based gauge controller for ESP32 / MicroPython

Deploy these files to the ESP32:
  gauge_vid6008.py        — stepper driver
  gaugemqttcontinuous.py  — this file
  umqtt/simple.py         — MicroPython built-in
  umqtt/robust.py         — https://raw.githubusercontent.com/micropython/micropython-lib/master/micropython/umqtt.robust/umqtt/robust.py
  config.json             — configuration (see below)

MQTT topics (all prefixed with mqtt_prefix from config.json):
  .../set                ← HA publishes target value here
  .../state              → ESP32 publishes current value
  .../status             → ESP32 publishes online/offline
  .../zero               ← publish anything to trigger zero
  .../led/red/set        ← ON/OFF
  .../led/green/set      ← ON/OFF
  .../led/backlight/set  ← 0-100

Serial log format: [HH:MM:SS.mmm] LEVEL  message
"""

import network
import utime
import ujson
from umqtt.robust import MQTTClient
from machine import Pin
from neopixel import NeoPixel
from gauge_vid6008 import Gauge

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------


def _ts():
    ms = utime.ticks_ms()
    return f"{(ms // 3600000) % 24:02d}:{(ms // 60000) % 60:02d}:{(ms // 1000) % 60:02d}.{ms % 1000:03d}"


def log(level, msg):
    print(f"[{_ts()}] {level:5s}  {msg}")


_DEBUG = False


def info(msg):
    if _DEBUG:
        log("INFO", msg)


def warn(msg):
    log("WARN", msg)


def log_err(msg):
    log("ERROR", msg)


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


def _load_config():
    try:
        with open("/config.json") as f:
            cfg = ujson.load(f)
        info("Config loaded from /config.json")
        return cfg
    except OSError:
        log_err("config.json not found — cannot continue")
        raise
    except Exception as e:
        log_err(f"config.json parse error: {e} — cannot continue")
        raise


_cfg = _load_config()

DEBUG = _cfg.get("debug", False)
_DEBUG = DEBUG
_WS2812_ORDER = _cfg.get("ws2812_order", "GRB").upper()

WIFI_SSID = _cfg["wifi_ssid"]
WIFI_PASSWORD = _cfg["wifi_password"]
MQTT_BROKER = _cfg["mqtt_broker"]
MQTT_PORT = int(_cfg.get("mqtt_port", 1883))
MQTT_USER = _cfg["mqtt_user"]
MQTT_PASSWORD = _cfg["mqtt_password"]
MQTT_CLIENT_ID = _cfg["mqtt_client_id"]
MQTT_PREFIX = _cfg["mqtt_prefix"]

MICROSTEPS_PER_SECOND = int(_cfg.get("microsteps_per_second", 600))
HEARTBEAT_MS = int(_cfg.get("heartbeat_ms", 10000))
REZERO_INTERVAL_MS = int(_cfg.get("rezero_interval_ms", 3600000))

backlight_cfg = _cfg.get("backlight", {})
BACKLIGHT_PIN = int(backlight_cfg.get("pin", _cfg.get("led_bl_pin", 23)))
BACKLIGHT_LEDS_PER_GAUGE = int(backlight_cfg.get("num_leds_per_gauge", 3))
STATUS_LEDS_PER_GAUGE = int(backlight_cfg.get("num_status_leds_per_gauge", 2))

device_cfg = _cfg.get("device", {})
DEVICE_NAME = device_cfg.get("name", _cfg.get("device_name", "Selsyn Multi"))
DEVICE_MODEL = device_cfg.get(
    "model", _cfg.get("device_model", "Chernobyl Selsyn-inspired gauge")
)
DEVICE_MFR = device_cfg.get(
    "manufacturer", _cfg.get("device_manufacturer", "AdeBaumann")
)
DEVICE_AREA = device_cfg.get("area", _cfg.get("device_area", "Control Panels"))

gauges = []
if "gauges" in _cfg:
    for i, g in enumerate(_cfg["gauges"]):
        led_cfg = g.get("leds", {})
        gauges.append(
            {
                "id": i,
                "name": g.get("name", f"Gauge {i + 1}"),
                "pins": tuple(g.get("pins", [12, 13])),
                "mode": g.get("mode", "stepdir"),
                "min": float(g.get("min", 0)),
                "max": float(g.get("max", 100)),
                "step_us": int(g.get("step_us", 200)),
                "entity_name": g.get("entity_name", f"Gauge {i + 1}"),
                "unit": g.get("unit", ""),
                "red_pin": int(led_cfg.get("red_pin", 33)),
                "green_pin": int(led_cfg.get("green_pin", 32)),
                "ws2812_red": tuple(led_cfg.get("ws2812_red", [255, 0, 0])),
                "ws2812_green": tuple(led_cfg.get("ws2812_green", [0, 255, 0])),
            }
        )
else:
    gauges.append(
        {
            "id": 0,
            "name": "Gauge 1",
            "pins": tuple(_cfg.get("gauge_pins", [12, 13])),
            "mode": _cfg.get("gauge_mode", "stepdir"),
            "min": float(_cfg.get("gauge_min", 0)),
            "max": float(_cfg.get("gauge_max", 7300)),
            "step_us": int(_cfg.get("gauge_step_us", 200)),
            "entity_name": _cfg.get("gauge_entity_name", "Selsyn 1 Power"),
            "unit": _cfg.get("gauge_unit", "W"),
            "red_pin": int(_cfg.get("led_red_pin", 33)),
            "green_pin": int(_cfg.get("led_green_pin", 32)),
            "ws2812_red": tuple(_cfg.get("ws2812_red", [255, 0, 0])),
            "ws2812_green": tuple(_cfg.get("ws2812_green", [0, 255, 0])),
        }
    )

# ---------------------------------------------------------------------------
# Gauge initialization
# ---------------------------------------------------------------------------

gauge_objects = []
for g in gauges:
    gauge_objects.append(
        Gauge(
            pins=g["pins"],
            mode=g["mode"],
            min_val=g["min"],
            max_val=g["max"],
            step_us=g["step_us"],
        )
    )
    info(
        f"Gauge {g['id']}: {g['name']} pins={g['pins']} mode={g['mode']} range=[{g['min']}, {g['max']}]"
    )

gauge_targets = [g["min"] for g in gauges]  # target value per gauge
gauge_last_rezero = [utime.ticks_ms() for _ in gauges]

# ---------------------------------------------------------------------------
# Topics (per-gauge)
# ---------------------------------------------------------------------------


def make_gauge_topics(prefix, gauge_id):
    return {
        "set": f"{prefix}/gauge{gauge_id}/set",
        "state": f"{prefix}/gauge{gauge_id}/state",
        "status": f"{prefix}/gauge{gauge_id}/status",
        "zero": f"{prefix}/gauge{gauge_id}/zero",
        "disc": f"homeassistant/number/{MQTT_CLIENT_ID}_g{gauge_id}/config",
        "led_red": f"{prefix}/gauge{gauge_id}/led/red/set",
        "led_green": f"{prefix}/gauge{gauge_id}/led/green/set",
        "led_bl": f"{prefix}/gauge{gauge_id}/led/backlight/set",
        "led_red_state": f"{prefix}/gauge{gauge_id}/led/red/state",
        "led_green_state": f"{prefix}/gauge{gauge_id}/led/green/state",
        "led_bl_state": f"{prefix}/gauge{gauge_id}/led/backlight/state",
        "led_red_disc": f"homeassistant/switch/{MQTT_CLIENT_ID}_g{gauge_id}_red/config",
        "led_green_disc": f"homeassistant/switch/{MQTT_CLIENT_ID}_g{gauge_id}_green/config",
        "led_bl_disc": f"homeassistant/light/{MQTT_CLIENT_ID}_g{gauge_id}_bl/config",
        "status_red": f"{prefix}/gauge{gauge_id}/status_led/red/set",
        "status_green": f"{prefix}/gauge{gauge_id}/status_led/green/set",
        "status_red_state": f"{prefix}/gauge{gauge_id}/status_led/red/state",
        "status_green_state": f"{prefix}/gauge{gauge_id}/status_led/green/state",
        "status_red_disc": f"homeassistant/switch/{MQTT_CLIENT_ID}_g{gauge_id}_status_red/config",
        "status_green_disc": f"homeassistant/switch/{MQTT_CLIENT_ID}_g{gauge_id}_status_green/config",
    }


gauge_topics = [make_gauge_topics(MQTT_PREFIX, g["id"]) for g in gauges]


T_SET = f"{MQTT_PREFIX}/set"
T_ZERO = f"{MQTT_PREFIX}/zero"


_DEVICE = {
    "identifiers": [MQTT_CLIENT_ID],
    "name": DEVICE_NAME,
    "model": DEVICE_MODEL,
    "manufacturer": DEVICE_MFR,
    "suggested_area": DEVICE_AREA,
}

# ---------------------------------------------------------------------------
# WiFi
# ---------------------------------------------------------------------------

_wifi_check_interval_ms = 30000
_last_wifi_check = 0
_wifi_sta = None

def connect_wifi(ssid, password, timeout_s=15):
    global _wifi_sta
    _wifi_sta = network.WLAN(network.STA_IF)
    _wifi_sta.active(True)
    if _wifi_sta.isconnected():
        ip, mask, gw, dns = _wifi_sta.ifconfig()
        info("WiFi already connected")
        info(f"  IP:{ip}  mask:{mask}  gw:{gw}  dns:{dns}")
        return ip
    info(f"WiFi connecting to '{ssid}' ...")
    _wifi_sta.connect(ssid, password)
    deadline = utime.time() + timeout_s
    while not _wifi_sta.isconnected():
        if utime.time() > deadline:
            log_err(f"WiFi connect timeout after {timeout_s}s")
            raise OSError("WiFi connect timeout")
        utime.sleep_ms(200)
    ip, mask, gw, dns = _wifi_sta.ifconfig()
    mac = ":".join(f"{b:02x}" for b in _wifi_sta.config("mac"))
    info("WiFi connected!")
    info(f"  SSID    : {ssid}")
    info(f"  MAC     : {mac}")
    info(f"  IP      : {ip}  mask:{mask}  gw:{gw}  dns:{dns}")
    return ip


def check_wifi():
    global _wifi_sta, _last_wifi_check
    now = utime.ticks_ms()
    if utime.ticks_diff(now, _last_wifi_check) < _wifi_check_interval_ms:
        return
    _last_wifi_check = now

    if _wifi_sta is None:
        _wifi_sta = network.WLAN(network.STA_IF)

    if _wifi_sta.isconnected():
        return

    log_err("WiFi lost connection — attempting reconnect...")
    try:
        _wifi_sta.active(True)
        _wifi_sta.connect(WIFI_SSID, WIFI_PASSWORD)
        deadline = utime.time() + 15
        while not _wifi_sta.isconnected():
            if utime.time() > deadline:
                log_err("WiFi reconnect timeout")
                return
            utime.sleep_ms(200)
        ip, mask, gw, dns = _wifi_sta.ifconfig()
        info(f"WiFi reconnected! IP:{ip}")
    except Exception as e:
        log_err(f"WiFi reconnect failed: {e}")


num_gauges = len(gauges)
leds_red = []
leds_green = []
for g in gauges:
    leds_red.append(Pin(g["red_pin"], Pin.OUT, value=0))
    leds_green.append(Pin(g["green_pin"], Pin.OUT, value=0))

total_backlight_leds = num_gauges * (BACKLIGHT_LEDS_PER_GAUGE + STATUS_LEDS_PER_GAUGE)
leds_bl = NeoPixel(Pin(BACKLIGHT_PIN), total_backlight_leds)

backlight_color = [(0, 0, 0) for _ in range(num_gauges)]
backlight_brightness = [100 for _ in range(num_gauges)]
backlight_on = [False for _ in range(num_gauges)]

status_led_red = [False for _ in range(num_gauges)]
status_led_green = [False for _ in range(num_gauges)]

_bl_dirty_since = None
_BL_SAVE_DELAY_MS = 5000


def _to_pixel(r, g, b):
    """Reorder RGB to match the WS2812 variant's byte order (GRB or RGB)."""
    if _WS2812_ORDER == "GRB":
        return (g, r, b)
    return (r, g, b)


def _backlight_changed(gauge_idx, new_color, new_on, new_brightness):
    return (
        new_color != backlight_color[gauge_idx]
        or new_on != backlight_on[gauge_idx]
        or (new_on and new_brightness != backlight_brightness[gauge_idx])
    )


def _mark_bl_dirty():
    global _bl_dirty_since
    _bl_dirty_since = utime.ticks_ms()


def _apply_backlight(gauge_idx, r, g, b, brightness):
    """Write RGB+brightness to the physical LEDs and mark dirty."""
    scale = brightness / 100
    leds_per_gauge = BACKLIGHT_LEDS_PER_GAUGE + STATUS_LEDS_PER_GAUGE
    base_idx = gauge_idx * leds_per_gauge
    pixel = _to_pixel(int(r * scale), int(g * scale), int(b * scale))
    for j in range(BACKLIGHT_LEDS_PER_GAUGE):
        leds_bl[base_idx + j] = pixel
    _update_status_leds(gauge_idx)
    leds_bl.write()
    _mark_bl_dirty()


def set_backlight_color(gauge_idx, r, g, b, brightness=None):
    global backlight_color, backlight_brightness, backlight_on
    if brightness is None:
        brightness = backlight_brightness[gauge_idx]
    new_on = brightness > 0
    if not _backlight_changed(gauge_idx, (r, g, b), new_on, brightness):
        return
    backlight_color[gauge_idx] = (r, g, b)
    if brightness > 0:
        backlight_brightness[gauge_idx] = brightness
    backlight_on[gauge_idx] = new_on
    _apply_backlight(gauge_idx, r, g, b, brightness)


def set_backlight_brightness(gauge_idx, brightness):
    global backlight_brightness, backlight_on
    clamped = max(0, min(100, brightness))
    new_on = clamped > 0
    if not _backlight_changed(gauge_idx, backlight_color[gauge_idx], new_on, clamped):
        return
    if clamped > 0:
        backlight_brightness[gauge_idx] = clamped
    backlight_on[gauge_idx] = new_on
    r, g, b = backlight_color[gauge_idx]
    _apply_backlight(gauge_idx, r, g, b, clamped)


def _update_status_leds(gauge_idx):
    leds_per_gauge = BACKLIGHT_LEDS_PER_GAUGE + STATUS_LEDS_PER_GAUGE
    base_idx = gauge_idx * leds_per_gauge + BACKLIGHT_LEDS_PER_GAUGE

    g_cfg = gauges[gauge_idx]
    red_color = g_cfg["ws2812_red"]
    green_color = g_cfg["ws2812_green"]

    if status_led_red[gauge_idx]:
        leds_bl[base_idx] = _to_pixel(*red_color)
    else:
        leds_bl[base_idx] = (0, 0, 0)

    if status_led_green[gauge_idx]:
        leds_bl[base_idx + 1] = _to_pixel(*green_color)
    else:
        leds_bl[base_idx + 1] = (0, 0, 0)


def set_status_led(gauge_idx, led_type, state):
    global status_led_red, status_led_green
    if led_type == "red":
        status_led_red[gauge_idx] = state
    elif led_type == "green":
        status_led_green[gauge_idx] = state
    _update_status_leds(gauge_idx)
    leds_bl.write()


def publish_backlight_states(client):
    """Publish current backlight state for all gauges as retained MQTT messages."""
    for i in range(num_gauges):
        gt = gauge_topics[i]
        r, g, b = backlight_color[i]
        brightness = backlight_brightness[i]
        state = {
            "state": "ON" if backlight_on[i] else "OFF",
            "color_mode": "rgb",
            "brightness": int(brightness * 2.55),
            "color": {"r": r, "g": g, "b": b},
        }
        try:
            client.publish(gt["led_bl_state"], ujson.dumps(state), retain=True)
        except Exception as e:
            log_err(f"Backlight state publish failed for gauge {i}: {e}")


def _flush_backlight_state():
    global _bl_dirty_since
    if _bl_dirty_since is None:
        return
    if utime.ticks_diff(utime.ticks_ms(), _bl_dirty_since) < _BL_SAVE_DELAY_MS:
        return
    if client_ref is None:
        return
    publish_backlight_states(client_ref)
    _bl_dirty_since = None
    info("Backlight state flushed to MQTT")


def _publish(topic, payload, retain=False):
    """Safely publish MQTT message, returning True on success."""
    if client_ref is None:
        return False
    try:
        client_ref.publish(topic, payload, retain=retain)
        return True
    except Exception as e:
        log_err(f"MQTT publish failed: {e}")
        return False


# ---------------------------------------------------------------------------
# MQTT callbacks
# ---------------------------------------------------------------------------


def on_message(topic, payload):
    if client_ref is None:
        return
    topic = topic.decode()
    payload = payload.decode().strip()
    info(f"MQTT rx  {topic}  {payload}")

    for i, gt in enumerate(gauge_topics):
        if topic == gt["zero"]:
            info(f"Zero command received for gauge {i}")
            gauge_objects[i].zero()
            gauge_last_rezero[i] = utime.ticks_ms()
            info(f"Zero complete gauge {i}")
            return

        if topic == gt["set"]:
            g = gauges[i]
            try:
                gauge_targets[i] = max(g["min"], min(g["max"], float(payload)))
                info(f"Gauge {i} target → {gauge_targets[i]:.1f}")
            except ValueError:
                warn(f"Invalid set value for gauge {i}: '{payload}'")
            return

        if topic == gt["led_red"]:
            state = payload.upper() == "ON"
            leds_red[i].value(1 if state else 0)
            _publish(gt["led_red_state"], "ON" if state else "OFF", retain=True)
            info(f"Gauge {i} red LED → {'ON' if state else 'OFF'}")
            return

        if topic == gt["led_green"]:
            state = payload.upper() == "ON"
            leds_green[i].value(1 if state else 0)
            _publish(gt["led_green_state"], "ON" if state else "OFF", retain=True)
            info(f"Gauge {i} green LED → {'ON' if state else 'OFF'}")
            return

        if topic == gt["led_bl"]:
            try:
                data = ujson.loads(payload)
                if data.get("state", "ON").upper() == "OFF":
                    set_backlight_brightness(i, 0)
                    _publish(
                        gt["led_bl_state"], ujson.dumps({"state": "OFF"}), retain=True
                    )
                    info(f"Gauge {i} backlight → OFF")
                    return
                color = data.get("color", {})
                r = max(0, min(255, int(color.get("r", backlight_color[i][0]))))
                g = max(0, min(255, int(color.get("g", backlight_color[i][1]))))
                b = max(0, min(255, int(color.get("b", backlight_color[i][2]))))
                raw_br = data.get("brightness", None)
                if raw_br is not None:
                    brightness = max(0, min(100, round(int(raw_br) / 2.55)))
                elif backlight_brightness[i] > 0:
                    brightness = backlight_brightness[i]
                else:
                    brightness = 100
            except Exception as e:
                warn(f"Invalid backlight payload for gauge {i}: '{payload}' ({e})")
                return
            set_backlight_color(i, r, g, b, brightness)
            state = {
                "state": "ON",
                "color_mode": "rgb",
                "brightness": int(brightness * 2.55),
                "color": {"r": r, "g": g, "b": b},
            }
            _publish(gt["led_bl_state"], ujson.dumps(state), retain=True)
            info(f"Gauge {i} backlight → #{r:02x}{g:02x}{b:02x} @ {brightness}%")
            return

        if topic == gt["status_red"]:
            state = payload.upper() == "ON"
            set_status_led(i, "red", state)
            _publish(gt["status_red_state"], "ON" if state else "OFF", retain=True)
            info(f"Gauge {i} status red → {'ON' if state else 'OFF'}")
            return

        if topic == gt["status_green"]:
            state = payload.upper() == "ON"
            set_status_led(i, "green", state)
            _publish(gt["status_green_state"], "ON" if state else "OFF", retain=True)
            info(f"Gauge {i} status green → {'ON' if state else 'OFF'}")
            return

    if topic == T_ZERO:
        for i, g in enumerate(gauge_objects):
            g.zero()
            gauge_last_rezero[i] = utime.ticks_ms()
        info("All gauges zeroed")
        return

    if topic == T_SET:
        try:
            data = ujson.loads(payload)
            if isinstance(data, dict):
                for i, val in enumerate(data.values()):
                    if i < len(gauges):
                        g = gauges[i]
                        gauge_targets[i] = max(g["min"], min(g["max"], float(val)))
                        info(f"Gauge {i} target → {gauge_targets[i]:.1f}")
            else:
                val = float(payload)
                for i in range(len(gauges)):
                    gauge_targets[i] = max(gauges[i]["min"], min(gauges[i]["max"], val))
                info(f"All gauges target → {val:.1f}")
        except Exception:
            try:
                val = float(payload)
                for i in range(len(gauges)):
                    gauge_targets[i] = max(gauges[i]["min"], min(gauges[i]["max"], val))
                info(f"All gauges target → {val:.1f}")
            except:
                warn(f"Invalid set value: '{payload}'")
        return




# ---------------------------------------------------------------------------
# MQTT connect + discovery
# ---------------------------------------------------------------------------


def _subscribe_all(c):
    c.subscribe(f"{MQTT_PREFIX}/set")
    c.subscribe(f"{MQTT_PREFIX}/zero")
    for i in range(num_gauges):
        prefix = f"{MQTT_PREFIX}/gauge{i}"
        c.subscribe(f"{prefix}/set")
        c.subscribe(f"{prefix}/zero")
        c.subscribe(f"{prefix}/led/red/set")
        c.subscribe(f"{prefix}/led/green/set")
        c.subscribe(f"{prefix}/led/backlight/set")
        c.subscribe(f"{prefix}/status_led/red/set")
        c.subscribe(f"{prefix}/status_led/green/set")


def connect_mqtt():
    global client_ref, _mqtt_connected
    info(f"Connecting to MQTT broker {MQTT_BROKER}:{MQTT_PORT} ...")
    client = MQTTClient(
        client_id=MQTT_CLIENT_ID,
        server=MQTT_BROKER,
        port=MQTT_PORT,
        user=MQTT_USER,
        password=MQTT_PASSWORD,
        keepalive=30,
    )
    # Don't set last will - it might be causing issues
    # client.set_last_will(T_STATUS, b"offline", retain=True, qos=0)
    client.set_callback(on_message)
    client.connect()
    client_ref = client
    _mqtt_connected = True
    info(f"MQTT connected  client_id={MQTT_CLIENT_ID}")


_mqtt_check_interval_ms = 30000
_last_mqtt_check = 0
client_ref = None
_mqtt_connected = False



def check_mqtt():
    global client_ref, _mqtt_connected, _last_mqtt_check
    now = utime.ticks_ms()
    if utime.ticks_diff(now, _last_mqtt_check) < _mqtt_check_interval_ms:
        return _mqtt_connected
    _last_mqtt_check = now

    if client_ref is None:
        return False

    try:
        client_ref.ping()
        _mqtt_connected = True
        return True
    except Exception as e:
        log_err(f"MQTT connection lost: {e}")
        _mqtt_connected = False

    # Try to reconnect
    log_err("Attempting MQTT reconnection...")
    for attempt in range(3):
        try:
            client_ref = MQTTClient(
                client_id=MQTT_CLIENT_ID,
                server=MQTT_BROKER,
                port=MQTT_PORT,
                user=MQTT_USER,
                password=MQTT_PASSWORD,
                keepalive=30,
            )
            client_ref.set_callback(on_message)
            client_ref.connect()
            _mqtt_connected = True
            info("MQTT reconnected!")
            publish_discovery(client_ref)
            _subscribe_all(client_ref)
            publish_state(client_ref)
            publish_backlight_states(client_ref)
            return True
        except Exception as e2:
            log_err(f"MQTT reconnect attempt {attempt + 1} failed: {e2}")
            utime.sleep_ms(2000)

    log_err("MQTT reconnection failed after 3 attempts")
    return False



def publish_discovery(client):
    """Publish all HA MQTT discovery payloads for gauges and LEDs."""
    _dev_ref = _DEVICE

    for i, g in enumerate(gauges):
        gt = gauge_topics[i]

        client.publish(
            gt["disc"],
            ujson.dumps(
                {
                    "name": g["entity_name"],
                    "unique_id": f"{MQTT_CLIENT_ID}_g{i}",
                    "cmd_t": gt["set"],
                    "stat_t": gt["state"],
                    "avty_t": gt["status"],
                    "min": g["min"],
                    "max": g["max"],
                    "step": 1,
                    "unit_of_meas": g["unit"],
                    "icon": "mdi:gauge",
                    "dev": _dev_ref,
                }
            ),
            retain=True,
        )
        info(f"Discovery: gauge {i} ({g['name']})")
        
        # Process MQTT messages between gauges
        for _ in range(5):
            client.check_msg()
            utime.sleep_ms(10)
        
        client.publish(
            gt["led_red_disc"],
            ujson.dumps(
                {
                    "name": f"{g['name']} Red LED",
                    "uniq_id": f"{MQTT_CLIENT_ID}_g{i}_red",
                    "cmd_t": gt["led_red"],
                    "stat_t": gt["led_red_state"],
                    "pl_on": "ON",
                    "pl_off": "OFF",
                    "icon": "mdi:led-on",
                    "dev": _dev_ref,
                    "ret": True,
                }
            ),
            retain=True,
        )
        info(f"Discovery: gauge {i} red LED")
        
        client.publish(
            gt["led_green_disc"],
            ujson.dumps(
                {
                    "name": f"{g['name']} Green LED",
                    "uniq_id": f"{MQTT_CLIENT_ID}_g{i}_green",
                    "cmd_t": gt["led_green"],
                    "stat_t": gt["led_green_state"],
                    "pl_on": "ON",
                    "pl_off": "OFF",
                    "icon": "mdi:led-on",
                    "dev": _dev_ref,
                    "ret": True,
                }
            ),
            retain=True,
        )
        info(f"Discovery: gauge {i} green LED")
        
        # Process MQTT messages
        for _ in range(5):
            client.check_msg()
            utime.sleep_ms(10)

        client.publish(
            gt["led_bl_disc"],
            ujson.dumps(
                {
                    "name": f"{g['name']} Backlight",
                    "uniq_id": f"{MQTT_CLIENT_ID}_g{i}_bl",
                    "cmd_t": gt["led_bl"],
                    "stat_t": gt["led_bl_state"],
                    "schema": "json",
                    "supported_color_modes": ["rgb"],
                    "icon": "mdi:led-strip",
                    "dev": _dev_ref,
                    "ret": True,
                }
            ),
            retain=True,
        )
        info(f"Discovery: gauge {i} backlight")

        client.publish(
            gt["status_red_disc"],
            ujson.dumps(
                {
                    "name": f"{g['name']} Status Red",
                    "uniq_id": f"{MQTT_CLIENT_ID}_g{i}_status_red",
                    "cmd_t": gt["status_red"],
                    "stat_t": gt["status_red_state"],
                    "pl_on": "ON",
                    "pl_off": "OFF",
                    "icon": "mdi:led-on",
                    "dev": _dev_ref,
                    "ret": True,
                }
            ),
            retain=True,
        )
        info(f"Discovery: gauge {i} status red")

        client.publish(
            gt["status_green_disc"],
            ujson.dumps(
                {
                    "name": f"{g['name']} Status Green",
                    "uniq_id": f"{MQTT_CLIENT_ID}_g{i}_status_green",
                    "cmd_t": gt["status_green"],
                    "stat_t": gt["status_green_state"],
                    "pl_on": "ON",
                    "pl_off": "OFF",
                    "icon": "mdi:led-on",
                    "dev": _dev_ref,
                    "ret": True,
                }
            ),
            retain=True,
        )
        info(f"Discovery: gauge {i} status green")

        # Process between gauges to avoid MQTT blocking
        for _ in range(5):
            client.check_msg()
            utime.sleep_ms(10)


def publish_state(client):
    for i, g in enumerate(gauge_objects):
        gt = gauge_topics[i]
        val = g.get()
        client.publish(gt["state"], str(val))


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main():
    info("=" * 48)
    info("Gauge MQTT controller starting")
    info("=" * 48)

    connect_wifi(WIFI_SSID, WIFI_PASSWORD)

    # Connect MQTT (no subscriptions yet — keeps broker silent during discovery)
    connect_mqtt()

    # Publish discovery — broker has nothing to send back yet
    info("Publishing discovery...")
    publish_discovery(client_ref)

    # Subscribe now — retained messages will start arriving from here
    _subscribe_all(client_ref)
    info("Draining initial retained messages...")
    for _ in range(50):
        client_ref.check_msg()
        utime.sleep_ms(20)

    # Now initialize gauges
    info("Zeroing gauges on startup ...")
    for i, g in enumerate(gauge_objects):
        g.zero()
        info(f"Zeroed gauge {i}")
    info("Zero complete")

    info("Publishing state...")
    publish_state(client_ref)
    utime.sleep_ms(50)
    for _ in range(5):
        client_ref.check_msg()
        utime.sleep_ms(20)

    info("Entering main loop")
    info("-" * 48)

    try:
        import ota
        ota.mark_ok()
    except:
        pass

    # Initialize variables for main loop
    last_heartbeat = utime.ticks_ms()
    now = 0
    was_moving = False

    while True:
        try:
            now = utime.ticks_ms()

            check_wifi()

            if not check_mqtt():
                utime.sleep_ms(1000)
                continue

            client_ref.check_msg()
            _flush_backlight_state()

            pending = [g.steps_toward(gauge_targets[i],limit=50) for i, g in enumerate(gauge_objects)]

            moved_any = any(s != 0 for s in pending)
            if moved_any:
                was_moving = True
                delay_us = 1_000_000 // MICROSTEPS_PER_SECOND
                for tick in range(max(abs(s) for s in pending)):
                    for i, g in enumerate(gauge_objects):
                        if tick < abs(pending[i]):
                            g.step(1 if pending[i] > 0 else -1)
                    utime.sleep_us(delay_us)
            else:
                if was_moving:
                    publish_state(client_ref)
                    was_moving = False
                utime.sleep_ms(10)

            if utime.ticks_diff(now, last_heartbeat) >= HEARTBEAT_MS:
                info(f"Heartbeat: {gauge_targets}")
                publish_state(client_ref)
                last_heartbeat = now

        except Exception as e:
            import sys
            sys.print_exception(e)
            log_err(f"Main loop error: {e} — continuing")
            utime.sleep_ms(100)


main()
