ohmu/dev/zigbee-collector/collector.py

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

Kaks 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``). Perioodilised
   südamelöögid (sama ``contact``) jäetakse vahele. Käivitusel seemendatakse
   viimane teadaolev olek andmebaasist; kui rida pole, salvestatakse esimene
   nähtud olek baasreana.

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
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"))

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()
}

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")

# Anduri nimi tuletatakse teemast (zigbee2mqtt/Temperatuur -> Temperatuur).
DEVICE = os.environ.get("DEVICE") or MQTT_TOPIC.rsplit("/", 1)[-1]

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_int(value: Any) -> int | None:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    return int(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] = []
        self.battery: int | None = None
        self.linkquality: int | None = None

    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"]))
        if isinstance(payload.get("battery"), (int, float)):
            self.battery = int(payload["battery"])
        if isinstance(payload.get("linkquality"), (int, float)):
            self.linkquality = int(payload["linkquality"])

    @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),
            "temp_min": min(self.temps) if self.temps else None,
            "temp_max": max(self.temps) if self.temps else None,
            "humidity_min": min(self.hums) if self.hums else None,
            "humidity_max": max(self.hums) if self.hums else None,
            "samples": max(len(self.temps), len(self.hums)),
            "battery": self.battery,
            "linkquality": self.linkquality,
        }


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
        battery_low = payload.get("battery_low")
        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",
            "battery": _as_int(payload.get("battery")),
            "voltage": _as_int(payload.get("voltage")),
            "linkquality": _as_int(payload.get("linkquality")),
            "battery_low": battery_low if isinstance(battery_low, bool) else None,
        }


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 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 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:
        resp = requests.post(
            f"{CH_URL}/",
            params={"database": CH_DB, "query": query},
            auth=(CH_USER, CH_PASS),
            timeout=15,
        )
        resp.raise_for_status()
        states: dict[str, bool] = {}
        for line in resp.text.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 main() -> int:
    lock = threading.Lock()
    stop = threading.Event()
    bucket = HourBucket(current_hour())
    temp_writer = ClickHouseWriter(CH_TABLE)
    door_writer = ClickHouseWriter(CH_DOOR_TABLE)
    door_tracker = DoorTracker(load_last_states())

    def on_connect(client, userdata, flags, reason_code, properties=None):
        if reason_code == 0:
            topics = [MQTT_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)
        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)

    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)
            elif msg.topic in DOOR_TOPICS:
                device = msg.topic.rsplit("/", 1)[-1]
                row = door_tracker.process(device, payload, now)
                if row is not None:
                    LOG.info("Olekumuutus: %s -> %s", device, row["olek"])
                    door_writer.enqueue(row)

    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)",
        CH_DB,
        CH_TABLE,
        CH_DB,
        CH_DOOR_TABLE,
        len(DOOR_TOPICS),
    )

    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())
        temp_writer.flush()
        door_writer.flush()

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


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