ohmu/dev/zigbee-collector/collector.py

#!/usr/bin/env python3
"""Zigbee andurite koguja.

Kolm voogu:

1. Temperatuur/niiskus (``MQTT_TOPIC``): mõõtmised koondatakse mälus täistunni
   kaupa ja kirjutatakse iga täistunni järel üks agregeeritud rida
   (``CH_TABLE``, vaikimisi ``uus67.temperatuur``).
2. Kontaktandurid (``DOOR_TOPICS``): salvestatakse **ainult olekumuutused**
   (``CH_DOOR_TABLE``, vaikimisi ``uus67.ukseandur``).
3. Häired (``CH_FAULT_TABLE``, vaikimisi ``uus67.vead``): **ainult muutused**
   (``aktiveeritud`` / ``taastatud``). Allikad: aku tühjenemine
   (``battery_low`` või ``battery`` lävend), nõrk signaal (``linkquality``),
   madal pinge (``voltage``), anduri kadumine (viimase teate timeout),
   koguja MQTT-ühenduse kadu ja Zigbee2MQTT bridge'i offline-olek.

Käivitub konteineris, mis on main_ai_network võrgus (ClickHouse nimega
lahenduv). Välised seadistused tulevad keskkonnamuutujatest.
"""

from __future__ import annotations

import json
import logging
import os
import signal
import threading
from datetime import datetime, timedelta
from typing import Any
from zoneinfo import ZoneInfo

import paho.mqtt.client as mqtt
import requests

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)
LOG = logging.getLogger("zigbee-collector")

TZ = ZoneInfo(os.environ.get("TZ_NAME", "Europe/Tallinn"))


def _env_float(name: str, default: float) -> float:
    try:
        return float(os.environ.get(name, default))
    except (TypeError, ValueError):
        return default


MQTT_HOST = os.environ.get("MQTT_HOST", "127.0.0.1")
MQTT_PORT = int(os.environ.get("MQTT_PORT", "1883"))
MQTT_USER = os.environ.get("MQTT_USER", "")
MQTT_PASS = os.environ.get("MQTT_PASS", "")
MQTT_TOPIC = os.environ.get("MQTT_TOPIC", "zigbee2mqtt/Temperatuur")
DOOR_TOPICS = {
    t.strip()
    for t in os.environ.get("DOOR_TOPICS", "").split(",")
    if t.strip()
}
BRIDGE_STATE_TOPIC = os.environ.get("BRIDGE_STATE_TOPIC", "zigbee2mqtt/bridge/state")

CH_URL = os.environ.get("CH_URL", "http://clickhouse:8123").rstrip("/")
CH_USER = os.environ.get("CH_USER", "default")
CH_PASS = os.environ.get("CH_PASS", "")
CH_DB = os.environ.get("CH_DB", "uus67")
CH_TABLE = os.environ.get("CH_TABLE", "temperatuur")
CH_DOOR_TABLE = os.environ.get("CH_DOOR_TABLE", "ukseandur")
CH_FAULT_TABLE = os.environ.get("CH_FAULT_TABLE", "vead")

# Häirete lävendid (hüstereesiga: aktiveerub <= X, taastub >= Y).
BATTERY_ACT = _env_float("BATTERY_LOW_THRESHOLD", 20)
BATTERY_CLR = _env_float("BATTERY_CLEAR_THRESHOLD", 25)
LQ_ACT = _env_float("LINKQUALITY_LOW_THRESHOLD", 20)
LQ_CLR = _env_float("LINKQUALITY_CLEAR_THRESHOLD", 30)
VOLT_ACT = _env_float("VOLTAGE_LOW_THRESHOLD", 2600)
VOLT_CLR = _env_float("VOLTAGE_CLEAR_THRESHOLD", 2700)

# Offline-tuvastuse timeoutid (minutites).
OFFLINE_TEMP_MIN = _env_float("OFFLINE_TIMEOUT_TEMP_MIN", 120)
OFFLINE_CONTACT_MIN = _env_float("OFFLINE_TIMEOUT_CONTACT_MIN", 240)

# Anduri nimi tuletatakse teemast (zigbee2mqtt/Temperatuur -> Temperatuur).
DEVICE = os.environ.get("DEVICE") or MQTT_TOPIC.rsplit("/", 1)[-1]
DOOR_DEVICES = {t.rsplit("/", 1)[-1] for t in DOOR_TOPICS}
MONITORED_DEVICES = {DEVICE, *DOOR_DEVICES}

SYSTEM_DEVICE_BRIDGE = "zigbee2mqtt"
SYSTEM_DEVICE_COLLECTOR = "koguja"

CHECK_INTERVAL_SEC = 20


def current_hour(now: datetime | None = None) -> datetime:
    now = now or datetime.now(TZ)
    return now.replace(minute=0, second=0, microsecond=0)


def _as_float(value: Any) -> float | None:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    return float(value)


def _fault_row(
    ts: datetime,
    device: str,
    tyyp: str,
    olek: str,
    detail: str,
    value: float | None,
) -> dict[str, Any]:
    return {
        "ts": ts.strftime("%Y-%m-%d %H:%M:%S"),
        "device": device,
        "tyyp": tyyp,
        "olek": olek,
        "detail": detail,
        "vaartus": value,
    }


class HourBucket:
    """Ühe tunni jooksul kogutud temperatuuri/niiskuse mõõtmised."""

    def __init__(self, hour: datetime) -> None:
        self.hour = hour
        self.temps: list[float] = []
        self.hums: list[float] = []

    def add(self, payload: dict[str, Any]) -> None:
        if isinstance(payload.get("temperature"), (int, float)):
            self.temps.append(float(payload["temperature"]))
        if isinstance(payload.get("humidity"), (int, float)):
            self.hums.append(float(payload["humidity"]))

    @property
    def empty(self) -> bool:
        return not self.temps and not self.hums

    def row(self) -> dict[str, Any]:
        def avg(values: list[float]) -> float | None:
            return round(sum(values) / len(values), 2) if values else None

        return {
            "hour": self.hour.strftime("%Y-%m-%d %H:%M:%S"),
            "device": DEVICE,
            "temperature": avg(self.temps),
            "humidity": avg(self.hums),
        }


class DoorTracker:
    """Kontaktandurite olekumuutuste jälgija.

    ``last`` hoiab viimast teadaolevat ``contact`` väärtust anduri kohta.
    Südamelöögid (sama väärtus) ei tekita rida. Esimene nähtud väärtus anduri
    kohta, millel puudub varasem olek, salvestatakse baasreana.
    """

    def __init__(self, last_states: dict[str, bool] | None = None) -> None:
        self.last: dict[str, bool] = dict(last_states or {})

    def process(
        self, device: str, payload: dict[str, Any], ts: datetime
    ) -> dict[str, Any] | None:
        contact = payload.get("contact")
        if not isinstance(contact, bool):
            return None
        if device in self.last and self.last[device] == contact:
            return None
        self.last[device] = contact
        return {
            "ts": ts.strftime("%Y-%m-%d %H:%M:%S"),
            "device": device,
            "contact": contact,
            # SNZB-04P: contact=true = suletud, false = avatud (Z2M HA discovery
            # payload_on=false, payload_off=true, device_class=door).
            "olek": "suletud" if contact else "avatud",
        }


class FaultTracker:
    """Häirete olekumuutuste jälgija (ainult muutused).

    ``current_keys`` hoiab iga seadme jaoks parajasti aktiivsete häirete
    tüüpe. Iga hindamise järel emiteeritakse read ainult üleminekute kohta
    (``aktiveeritud`` / ``taastatud``). Käivitusel seemendatakse aktiivsed
    häired andmebaasist, et restart ei dubleeriks kirjeid.
    """

    def __init__(self, active: dict[str, set[str]] | None = None) -> None:
        self.payload_faults: dict[str, dict[str, tuple[str, float | None]]] = {}
        self.offline_faults: dict[str, tuple[str, float | None]] = {}
        self.current_keys: dict[str, set[str]] = {
            device: set(tyybid) for device, tyybid in (active or {}).items()
        }

    def _low(self, device: str, tyyp: str, value: float, activate_at: float, clear_at: float) -> bool:
        if tyyp in self.current_keys.get(device, set()):
            return value < clear_at
        return value <= activate_at

    def _derive(self, device: str, payload: dict[str, Any]) -> dict[str, tuple[str, float | None]]:
        faults: dict[str, tuple[str, float | None]] = {}
        battery_low = payload.get("battery_low")
        battery = _as_float(payload.get("battery"))
        if battery_low is True:
            faults["low_battery"] = ("battery_low=true", battery)
        elif battery is not None and self._low(device, "low_battery", battery, BATTERY_ACT, BATTERY_CLR):
            faults["low_battery"] = (f"battery {battery:.0f}%", battery)
        lq = _as_float(payload.get("linkquality"))
        if lq is not None and self._low(device, "weak_signal", lq, LQ_ACT, LQ_CLR):
            faults["weak_signal"] = (f"linkquality {lq:.0f}", lq)
        voltage = _as_float(payload.get("voltage"))
        if voltage is not None and self._low(device, "low_voltage", voltage, VOLT_ACT, VOLT_CLR):
            faults["low_voltage"] = (f"voltage {voltage:.0f} mV", voltage)
        return faults

    def _reconcile(self, device: str, ts: datetime) -> list[dict[str, Any]]:
        current: dict[str, tuple[str, float | None]] = {}
        current.update(self.payload_faults.get(device, {}))
        if device in self.offline_faults:
            current["offline"] = self.offline_faults[device]
        prev = self.current_keys.get(device, set())
        rows: list[dict[str, Any]] = []
        for tyyp, (detail, value) in current.items():
            if tyyp not in prev:
                rows.append(_fault_row(ts, device, tyyp, "aktiveeritud", detail, value))
        for tyyp in prev - set(current):
            rows.append(_fault_row(ts, device, tyyp, "taastatud", "", None))
        self.current_keys[device] = set(current)
        return rows

    def on_device_message(
        self, device: str, payload: dict[str, Any], ts: datetime
    ) -> list[dict[str, Any]]:
        self.payload_faults[device] = self._derive(device, payload)
        self.offline_faults.pop(device, None)
        return self._reconcile(device, ts)

    def set_offline(
        self, device: str, offline: bool, detail: str, ts: datetime
    ) -> list[dict[str, Any]]:
        if offline:
            self.offline_faults[device] = (detail, None)
        else:
            self.offline_faults.pop(device, None)
        return self._reconcile(device, ts)

    def set_system_fault(
        self,
        device: str,
        tyyp: str,
        active: bool,
        detail: str,
        ts: datetime,
    ) -> list[dict[str, Any]]:
        faults = dict(self.payload_faults.get(device, {}))
        if active:
            faults[tyyp] = (detail, None)
        else:
            faults.pop(tyyp, None)
        self.payload_faults[device] = faults
        return self._reconcile(device, ts)


class ClickHouseWriter:
    def __init__(self, table: str) -> None:
        self.table = table
        self.pending: list[dict[str, Any]] = []

    def enqueue(self, row: dict[str, Any]) -> None:
        self.pending.append(row)

    def enqueue_many(self, rows: list[dict[str, Any]]) -> None:
        self.pending.extend(rows)

    def flush(self) -> None:
        if not self.pending:
            return
        body = "\n".join(json.dumps(r, ensure_ascii=False) for r in self.pending)
        try:
            resp = requests.post(
                f"{CH_URL}/",
                params={
                    "database": CH_DB,
                    "query": f"INSERT INTO {CH_DB}.{self.table} FORMAT JSONEachRow",
                },
                data=body.encode("utf-8"),
                auth=(CH_USER, CH_PASS),
                timeout=15,
            )
            resp.raise_for_status()
            LOG.info("ClickHouse'i (%s) kirjutatud %d rida", self.table, len(self.pending))
            self.pending.clear()
        except Exception:
            LOG.exception(
                "ClickHouse'i (%s) kirjutamine ebaõnnestus, jätan %d rida järjekorda",
                self.table,
                len(self.pending),
            )


def _ch_query(query: str) -> str:
    resp = requests.post(
        f"{CH_URL}/",
        params={"database": CH_DB, "query": query},
        auth=(CH_USER, CH_PASS),
        timeout=15,
    )
    resp.raise_for_status()
    return resp.text


def load_last_states() -> dict[str, bool]:
    """Loeb andmebaasist iga anduri viimase teadaoleva ``contact`` väärtuse."""
    query = (
        f"SELECT device, argMax(contact, ts) AS contact "
        f"FROM {CH_DB}.{CH_DOOR_TABLE} GROUP BY device FORMAT JSONEachRow"
    )
    try:
        states: dict[str, bool] = {}
        for line in _ch_query(query).splitlines():
            if not line.strip():
                continue
            obj = json.loads(line)
            value = obj["contact"]
            states[obj["device"]] = value if isinstance(value, bool) else bool(int(value))
        LOG.info("Seemendatud %d anduri viimane olek", len(states))
        return states
    except Exception:
        LOG.exception("Viimase oleku laadimine ebaõnnestus, alustan tühjalt")
        return {}


def load_active_faults() -> dict[str, set[str]]:
    """Loeb andmebaasist iga seadme parajasti aktiivsed häirete tüübid."""
    query = (
        f"SELECT device, tyyp, argMax(olek, ts) AS olek "
        f"FROM {CH_DB}.{CH_FAULT_TABLE} GROUP BY device, tyyp FORMAT JSONEachRow"
    )
    try:
        active: dict[str, set[str]] = {}
        for line in _ch_query(query).splitlines():
            if not line.strip():
                continue
            obj = json.loads(line)
            if obj.get("olek") == "aktiveeritud":
                active.setdefault(obj["device"], set()).add(obj["tyyp"])
        LOG.info("Seemendatud %d seadme aktiivsed häired", len(active))
        return active
    except Exception:
        LOG.exception("Aktiivsete häirete laadimine ebaõnnestus, alustan tühjalt")
        return {}


def load_last_seen() -> dict[str, datetime]:
    """Loeb andmebaasist iga anduri viimase kirje aja (offline-baas)."""
    query = (
        "SELECT device, max(ts) AS ts FROM ("
        f"SELECT device, ts FROM {CH_DB}.{CH_DOOR_TABLE} "
        "UNION ALL "
        f"SELECT device, hour AS ts FROM {CH_DB}.{CH_TABLE}"
        ") GROUP BY device FORMAT JSONEachRow"
    )
    try:
        seen: dict[str, datetime] = {}
        for line in _ch_query(query).splitlines():
            if not line.strip():
                continue
            obj = json.loads(line)
            seen[obj["device"]] = datetime.strptime(obj["ts"], "%Y-%m-%d %H:%M:%S").replace(
                tzinfo=TZ
            )
        LOG.info("Seemendatud %d anduri viimane teate aeg", len(seen))
        return seen
    except Exception:
        LOG.exception("Viimaste teate aegade laadimine ebaõnnestus, alustan tühjalt")
        return {}


def main() -> int:
    lock = threading.Lock()
    stop = threading.Event()
    bucket = HourBucket(current_hour())
    temp_writer = ClickHouseWriter(CH_TABLE)
    door_writer = ClickHouseWriter(CH_DOOR_TABLE)
    fault_writer = ClickHouseWriter(CH_FAULT_TABLE)
    door_tracker = DoorTracker(load_last_states())
    fault_tracker = FaultTracker(load_active_faults())

    start = datetime.now(TZ)
    last_seen: dict[str, datetime] = {}
    for device, seen in load_last_seen().items():
        last_seen[device] = seen
    for device in MONITORED_DEVICES:
        last_seen.setdefault(device, start)

    def offline_timeout(device: str) -> float:
        if device == DEVICE:
            return OFFLINE_TEMP_MIN
        return OFFLINE_CONTACT_MIN

    def on_connect(client, userdata, flags, reason_code, properties=None):
        if reason_code == 0:
            topics = [MQTT_TOPIC, BRIDGE_STATE_TOPIC, *sorted(DOOR_TOPICS)]
            LOG.info("MQTT ühendatud (%s:%s), tellin %s", MQTT_HOST, MQTT_PORT, ", ".join(topics))
            for topic in topics:
                client.subscribe(topic, qos=1)
            if not stop.is_set():
                with lock:
                    fault_writer.enqueue_many(
                        fault_tracker.set_system_fault(
                            SYSTEM_DEVICE_COLLECTOR,
                            "mqtt_disconnected",
                            False,
                            "",
                            datetime.now(TZ),
                        )
                    )
        else:
            LOG.error("MQTT ühendus ebaõnnestus: %s", reason_code)

    def on_disconnect(client, userdata, flags, reason_code, properties=None):
        LOG.warning("MQTT ühendus katkes: %s", reason_code)
        if not stop.is_set():
            with lock:
                fault_writer.enqueue_many(
                    fault_tracker.set_system_fault(
                        SYSTEM_DEVICE_COLLECTOR,
                        "mqtt_disconnected",
                        True,
                        f"MQTT ühendus katkes: {reason_code}",
                        datetime.now(TZ),
                    )
                )

    def on_message(client, userdata, msg):
        try:
            payload = json.loads(msg.payload.decode("utf-8"))
        except Exception:
            LOG.exception("Vigane MQTT teade teemal %s", msg.topic)
            return
        if not isinstance(payload, dict):
            return
        now = datetime.now(TZ)
        with lock:
            if msg.topic == MQTT_TOPIC:
                bucket.add(payload)
                last_seen[DEVICE] = now
                fault_writer.enqueue_many(fault_tracker.on_device_message(DEVICE, payload, now))
            elif msg.topic in DOOR_TOPICS:
                device = msg.topic.rsplit("/", 1)[-1]
                last_seen[device] = now
                row = door_tracker.process(device, payload, now)
                if row is not None:
                    LOG.info("Olekumuutus: %s -> %s", device, row["olek"])
                    door_writer.enqueue(row)
                fault_writer.enqueue_many(fault_tracker.on_device_message(device, payload, now))
            elif msg.topic == BRIDGE_STATE_TOPIC:
                offline = payload.get("state") == "offline"
                fault_writer.enqueue_many(
                    fault_tracker.set_system_fault(
                        SYSTEM_DEVICE_BRIDGE,
                        "bridge_offline",
                        offline,
                        "zigbee2mqtt bridge offline",
                        now,
                    )
                )

    client = mqtt.Client(
        mqtt.CallbackAPIVersion.VERSION2,
        client_id="zigbee-collector",
    )
    if MQTT_USER:
        client.username_pw_set(MQTT_USER, MQTT_PASS)
    client.on_connect = on_connect
    client.on_disconnect = on_disconnect
    client.on_message = on_message
    client.reconnect_delay_set(min_delay=1, max_delay=60)
    client.connect_async(MQTT_HOST, MQTT_PORT, keepalive=60)
    client.loop_start()

    def shutdown(signum, frame):
        LOG.info("Lõpetamissignaal (%s), sulgen", signum)
        stop.set()

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)

    LOG.info(
        "Koguja käivitatud: temp=%s.%s, uksed=%s.%s (%d teemat), häired=%s.%s",
        CH_DB,
        CH_TABLE,
        CH_DB,
        CH_DOOR_TABLE,
        len(DOOR_TOPICS),
        CH_DB,
        CH_FAULT_TABLE,
    )

    while not stop.is_set():
        stop.wait(CHECK_INTERVAL_SEC)
        now = datetime.now(TZ)
        with lock:
            if current_hour(now) != bucket.hour:
                finished = bucket
                bucket = HourBucket(current_hour(now))
            else:
                finished = None
            if finished is not None and not finished.empty:
                temp_writer.enqueue(finished.row())
            for device in sorted(MONITORED_DEVICES):
                seen = last_seen.get(device, start)
                timeout = timedelta(minutes=offline_timeout(device))
                offline = (now - seen) > timeout
                detail = (
                    f"viimane teade {seen.strftime('%Y-%m-%d %H:%M:%S')}"
                    if offline
                    else ""
                )
                fault_writer.enqueue_many(
                    fault_tracker.set_offline(device, offline, detail, now)
                )
        temp_writer.flush()
        door_writer.flush()
        fault_writer.flush()

    temp_writer.flush()
    door_writer.flush()
    fault_writer.flush()
    client.loop_stop()
    client.disconnect()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())