#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (C) 2026 Christian Wirtz <doc@snowheaven.de>
# This file is part of the Zigbee2MQTT Checkmk extension.
# License: GPL-2.0-only, see LICENSE in the repository root.
"""Checkmk special agent for Zigbee2MQTT.

Takes one snapshot of a Zigbee2MQTT instance and turns it into agent sections.
Two transports are supported:

  * ``websocket`` (default) -- the WebSocket API of Zigbee2MQTT's own frontend
    (``ws[s]://<host>:<port><base_url>api``). Right after connecting, Zigbee2MQTT
    sends every retained ``bridge/*`` message it holds (``bridge/state``,
    ``bridge/info``, ``bridge/devices``, ``bridge/health``, per-device
    ``<name>/availability``) and then the *cached current state of every device*
    (incl. ``last_seen`` and ``linkquality``). So one short connection yields a
    complete picture; the agent disconnects as soon as every device's state has
    arrived. An optional ``auth_token`` is passed as ``?token=``.
  * ``mqtt`` -- subscribes to ``<base_topic>/#`` on the MQTT broker and listens for
    a few seconds. The ``bridge/*`` topics are retained, but device state topics
    usually are not, so per-device values (last_seen, link quality, battery) are
    only available for devices that happen to report during the listen window.
    Meant as a fallback when the frontend is disabled or unreachable.

``websocket_mqtt_fallback`` tries the WebSocket first and switches to MQTT only
if that fails; the Bridge service then shows that (and why) the fallback was used.

Design notes
------------
* Standard library only: minimal WebSocket (RFC 6455, client side) and MQTT 3.1.1
  (CONNECT/SUBSCRIBE/PUBLISH QoS 0) implementations, TLS via ``ssl``. Nothing is
  vendored, the .mkp stays architecture independent.
* Read-only: the agent never sends a command to Zigbee2MQTT (no ``bridge/request``,
  no ``/set``). Over MQTT it subscribes only.
* If neither transport delivers data, a ``zigbee2mqtt_bridge`` section with
  ``reachable: false`` is still emitted, so the "Zigbee2MQTT Bridge" service goes
  CRIT with the reason instead of the host just showing "no agent output".
* Secrets (auth token, MQTT password) are resolved from the Checkmk password store
  via ``--token-id`` / ``--mqtt-password-id``; the literal ``--token`` /
  ``--mqtt-password`` options exist only for manual CLI testing.
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import re
import socket
import ssl
import struct
import sys
import time
import unicodedata
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import Any
from urllib.parse import quote

Messages = dict[str, Any]


class TransportError(Exception):
    """A transport could not deliver a snapshot (connect, auth, protocol)."""


# -----------------------------------------------------------------------------------
# helpers
# -----------------------------------------------------------------------------------
def _parse_payload(raw: bytes) -> Any:
    text = raw.decode("utf-8", errors="replace")
    try:
        return json.loads(text)
    except ValueError:
        return text


def _open_socket(host: str, port: int, timeout: float, tls: bool, verify: bool) -> socket.socket:
    try:
        sock = socket.create_connection((host, port), timeout=timeout)
    except OSError as exc:
        raise TransportError(f"cannot connect to {host}:{port}: {exc}") from exc
    if not tls:
        return sock
    ctx = ssl.create_default_context()
    if not verify:
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
    try:
        return ctx.wrap_socket(sock, server_hostname=host)
    except (ssl.SSLError, OSError) as exc:
        sock.close()
        raise TransportError(f"TLS handshake with {host}:{port} failed: {exc}") from exc


class _Reader:
    """Buffered exact-length reads from a socket."""

    def __init__(self, sock: socket.socket, initial: bytes = b"") -> None:
        self.sock = sock
        self.buf = initial

    def read(self, n: int) -> bytes:
        while len(self.buf) < n:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise EOFError("connection closed by peer")
            self.buf += chunk
        data, self.buf = self.buf[:n], self.buf[n:]
        return data


# -----------------------------------------------------------------------------------
# WebSocket transport (Zigbee2MQTT frontend API)
# -----------------------------------------------------------------------------------
def ws_path(base_url: str, token: str | None) -> str:
    base = "/" + base_url.strip("/")
    path = base.rstrip("/") + "/api"
    if token:
        path += "?token=" + quote(token, safe="")
    return path


def _ws_send(sock: socket.socket, opcode: int, payload: bytes = b"") -> None:
    # client frames must be masked (RFC 6455 5.3)
    header = bytes([0x80 | opcode])
    n = len(payload)
    if n < 126:
        header += bytes([0x80 | n])
    elif n < 65536:
        header += bytes([0x80 | 126]) + struct.pack(">H", n)
    else:
        header += bytes([0x80 | 127]) + struct.pack(">Q", n)
    mask = os.urandom(4)
    masked = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
    sock.sendall(header + mask + masked)


def ws_snapshot(
    host: str,
    port: int,
    base_url: str,
    token: str | None,
    tls: bool,
    verify: bool,
    timeout: float,
    idle: float = 2.0,
) -> Messages:
    sock = _open_socket(host, port, timeout, tls, verify)
    try:
        key = base64.b64encode(os.urandom(16)).decode()
        request = (
            f"GET {ws_path(base_url, token)} HTTP/1.1\r\n"
            f"Host: {host}:{port}\r\n"
            "Upgrade: websocket\r\n"
            "Connection: Upgrade\r\n"
            f"Sec-WebSocket-Key: {key}\r\n"
            "Sec-WebSocket-Version: 13\r\n"
            "User-Agent: checkmk-agent-zigbee2mqtt\r\n\r\n"
        )
        deadline = time.monotonic() + timeout
        try:
            sock.sendall(request.encode())
            raw = b""
            while b"\r\n\r\n" not in raw:
                chunk = sock.recv(4096)
                if not chunk:
                    raise TransportError("connection closed during WebSocket handshake")
                raw += chunk
                if len(raw) > 65536:
                    raise TransportError("oversized WebSocket handshake response")
        except OSError as exc:
            raise TransportError(f"WebSocket handshake failed: {exc}") from exc
        head, rest = raw.split(b"\r\n\r\n", 1)
        status_line = head.split(b"\r\n", 1)[0].decode(errors="replace")
        if " 101 " not in f"{status_line} ":
            raise TransportError(
                f"WebSocket upgrade refused ({status_line}) -- frontend enabled, port and "
                "base URL correct?"
            )

        reader = _Reader(sock, rest)
        messages: Messages = {}
        expected: set[str] | None = None
        fragments = b""
        while True:
            now = time.monotonic()
            if now >= deadline:
                break
            # once the device list is known, wait only for the outstanding states
            sock.settimeout(min(idle, deadline - now) if messages else deadline - now)
            try:
                b1, b2 = reader.read(2)
                length = b2 & 0x7F
                if length == 126:
                    (length,) = struct.unpack(">H", reader.read(2))
                elif length == 127:
                    (length,) = struct.unpack(">Q", reader.read(8))
                mask = reader.read(4) if b2 & 0x80 else b""
                data = reader.read(length)
            except socket.timeout:
                break
            except (EOFError, OSError) as exc:
                if messages:
                    break
                raise TransportError(f"WebSocket connection lost: {exc}") from exc
            if mask:
                data = bytes(b ^ mask[i % 4] for i, b in enumerate(data))
            opcode = b1 & 0x0F
            if opcode == 0x8:  # close
                code = struct.unpack(">H", data[:2])[0] if len(data) >= 2 else 0
                if code == 4401:
                    raise TransportError(
                        "Zigbee2MQTT frontend rejected the auth token (4401 Unauthorized)"
                    )
                if not messages:
                    raise TransportError(f"WebSocket closed by Zigbee2MQTT (code {code})")
                break
            if opcode == 0x9:  # ping
                _ws_send(sock, 0xA, data)
                continue
            if opcode not in (0x0, 0x1, 0x2):
                continue
            fragments += data
            if not b1 & 0x80:
                continue
            frame, fragments = fragments, b""
            try:
                msg = json.loads(frame)
                topic = msg["topic"]
            except (ValueError, KeyError, TypeError):
                continue
            messages[topic] = msg.get("payload")
            if topic == "bridge/devices" and isinstance(msg.get("payload"), list):
                expected = {
                    d.get("friendly_name")
                    for d in msg["payload"]
                    if isinstance(d, dict) and d.get("type") != "Coordinator"
                }
            if (
                expected is not None
                and expected.issubset(messages)
                and "bridge/state" in messages
                and "bridge/info" in messages
            ):
                break
        try:
            _ws_send(sock, 0x8, struct.pack(">H", 1000))
        except OSError:
            pass
        if not messages:
            raise TransportError("WebSocket delivered no data")
        return messages
    finally:
        sock.close()


# -----------------------------------------------------------------------------------
# MQTT transport (MQTT 3.1.1, subscribe only)
# -----------------------------------------------------------------------------------
def _mqtt_str(value: str) -> bytes:
    data = value.encode("utf-8")
    return struct.pack(">H", len(data)) + data


def _mqtt_remaining_length(n: int) -> bytes:
    out = b""
    while True:
        digit, n = n % 128, n // 128
        out += bytes([digit | (0x80 if n else 0)])
        if not n:
            return out


def _mqtt_packet(first_byte: int, body: bytes) -> bytes:
    return bytes([first_byte]) + _mqtt_remaining_length(len(body)) + body


def _mqtt_read_packet(reader: _Reader) -> tuple[int, bytes]:
    (first,) = reader.read(1)
    multiplier, length = 1, 0
    for _ in range(4):
        (digit,) = reader.read(1)
        length += (digit & 0x7F) * multiplier
        multiplier *= 128
        if not digit & 0x80:
            break
    return first, reader.read(length)


_CONNACK_ERRORS = {
    1: "unacceptable protocol version",
    2: "client identifier rejected",
    3: "server unavailable",
    4: "bad user name or password",
    5: "not authorized",
}


def mqtt_snapshot(
    host: str,
    port: int,
    user: str | None,
    password: str | None,
    base_topic: str,
    tls: bool,
    verify: bool,
    timeout: float,
    listen: float,
) -> Messages:
    sock = _open_socket(host, port, timeout, tls, verify)
    try:
        flags = 0x02  # clean session
        payload = _mqtt_str(f"cmk-agent-zigbee2mqtt-{os.getpid()}")
        if user:
            flags |= 0x80
            payload += _mqtt_str(user)
            if password is not None:
                flags |= 0x40
        if user and password is not None:
            payload += _mqtt_str(password)
        body = _mqtt_str("MQTT") + bytes([4, flags]) + struct.pack(">H", 60) + payload
        reader = _Reader(sock)
        try:
            sock.sendall(_mqtt_packet(0x10, body))
            first, data = _mqtt_read_packet(reader)
            if first >> 4 != 2 or len(data) < 2:
                raise TransportError("MQTT broker sent no CONNACK")
            if data[1] != 0:
                reason = _CONNACK_ERRORS.get(data[1], f"code {data[1]}")
                raise TransportError(f"MQTT broker refused the connection: {reason}")
            topic_filter = f"{base_topic}/#"
            sock.sendall(_mqtt_packet(0x82, struct.pack(">H", 1) + _mqtt_str(topic_filter) + b"\x00"))
        except (OSError, EOFError) as exc:
            raise TransportError(f"MQTT connection failed: {exc}") from exc

        prefix = f"{base_topic}/"
        messages: Messages = {}
        deadline = time.monotonic() + listen
        subscribed = False
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break
            sock.settimeout(remaining)
            try:
                first, data = _mqtt_read_packet(reader)
            except socket.timeout:
                break
            except (EOFError, OSError) as exc:
                if messages:
                    break
                raise TransportError(f"MQTT connection lost: {exc}") from exc
            ptype = first >> 4
            if ptype == 9:  # SUBACK
                if len(data) >= 3 and data[2] == 0x80:
                    raise TransportError(f"MQTT broker rejected the subscription to {topic_filter}")
                subscribed = True
            elif ptype == 3:  # PUBLISH
                (tlen,) = struct.unpack(">H", data[:2])
                topic = data[2 : 2 + tlen].decode("utf-8", errors="replace")
                offset = 2 + tlen + (2 if (first >> 1) & 0x03 else 0)
                if topic.startswith(prefix):
                    messages[topic[len(prefix) :]] = _parse_payload(data[offset:])
        try:
            sock.sendall(b"\xe0\x00")  # DISCONNECT
        except OSError:
            pass
        if not subscribed and not messages:
            raise TransportError("MQTT broker did not confirm the subscription")
        if not messages:
            raise TransportError(
                f"no messages on {topic_filter} -- is the base topic correct and "
                "Zigbee2MQTT connected to this broker?"
            )
        return messages
    finally:
        sock.close()


# -----------------------------------------------------------------------------------
# section building
# -----------------------------------------------------------------------------------
def parse_last_seen(value: Any) -> float | None:
    """Zigbee2MQTT's last_seen: epoch milliseconds or an ISO 8601 string."""
    if isinstance(value, bool):
        return None
    if isinstance(value, (int, float)):
        return value / 1000.0 if value > 1e11 else float(value)
    if isinstance(value, str) and value:
        try:
            dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
        except ValueError:
            return None
        if dt.tzinfo is None:
            dt = dt.replace(tzinfo=timezone.utc)
        return dt.timestamp()
    return None


def _state_of(payload: Any) -> str | None:
    if isinstance(payload, dict):
        state = payload.get("state")
        return str(state) if state is not None else None
    if isinstance(payload, str):
        return payload
    return None


_INFO_KEYS = (
    "version",
    "commit",
    "coordinator",
    "network",
    "log_level",
    "permit_join",
    "permit_join_end",
    "restart_required",
    "os",
    "mqtt",
)


def build_bridge(messages: Messages, transport: str, fallback_reason: str | None) -> dict[str, Any]:
    info = messages.get("bridge/info") if isinstance(messages.get("bridge/info"), dict) else {}
    config = info.get("config") if isinstance(info.get("config"), dict) else {}
    advanced = config.get("advanced") if isinstance(config.get("advanced"), dict) else {}
    availability = config.get("availability")
    health_cfg = config.get("health") if isinstance(config.get("health"), dict) else {}
    serial = config.get("serial") if isinstance(config.get("serial"), dict) else {}
    bridge: dict[str, Any] = {
        "serial_port": serial.get("port"),
        "adapter": serial.get("adapter"),
        "reachable": True,
        "transport": transport,
        "state": _state_of(messages.get("bridge/state")),
        "info": {k: info[k] for k in _INFO_KEYS if k in info},
        "last_seen_mode": advanced.get("last_seen"),
        "availability_enabled": (
            bool(availability.get("enabled")) if isinstance(availability, dict) else bool(availability)
        ),
        "health_interval_min": health_cfg.get("interval"),
        "health_reset_on_check": health_cfg.get("reset_on_check"),
    }
    if fallback_reason:
        bridge["fallback_reason"] = fallback_reason
    health = messages.get("bridge/health")
    if isinstance(health, dict):
        bridge["health"] = {k: v for k, v in health.items() if k != "devices"}
    return bridge


def build_devices(messages: Messages, exclude: Sequence[re.Pattern[str]]) -> list[dict[str, Any]]:
    devices = messages.get("bridge/devices")
    if not isinstance(devices, list):
        return []
    health = messages.get("bridge/health")
    health_devices = health.get("devices", {}) if isinstance(health, dict) else {}
    out = []
    for dev in devices:
        if not isinstance(dev, dict) or dev.get("type") == "Coordinator":
            continue
        name = dev.get("friendly_name") or dev.get("ieee_address")
        if not name or any(p.search(name) for p in exclude):
            continue
        definition = dev.get("definition") if isinstance(dev.get("definition"), dict) else {}
        interview_state = dev.get("interview_state")
        if interview_state is None and "interview_completed" in dev:
            # Zigbee2MQTT < 2.1 only has the boolean flags
            if dev.get("interviewing"):
                interview_state = "IN_PROGRESS"
            else:
                interview_state = "SUCCESSFUL" if dev.get("interview_completed") else "FAILED"
        entry: dict[str, Any] = {
            "name": name,
            "ieee_address": dev.get("ieee_address"),
            "type": dev.get("type"),
            "power_source": dev.get("power_source"),
            "vendor": definition.get("vendor"),
            "model": definition.get("model"),
            "description": definition.get("description") or dev.get("description"),
            "model_id": dev.get("model_id"),
            "software_build_id": dev.get("software_build_id"),
            "supported": dev.get("supported"),
            "disabled": bool(dev.get("disabled")),
            "interview_state": interview_state,
            "availability": _state_of(messages.get(f"{name}/availability")),
            "has_state": False,
        }
        state = messages.get(name)
        if isinstance(state, dict):
            entry["has_state"] = True
            entry["last_seen"] = parse_last_seen(state.get("last_seen"))
            for key in ("linkquality", "battery", "battery_low"):
                if key in state and not isinstance(state[key], dict):
                    entry[key] = state[key]
            update = state.get("update")
            if isinstance(update, dict):
                entry["update"] = {
                    k: update.get(k) for k in ("state", "installed_version", "latest_version")
                }
        stats = health_devices.get(dev.get("ieee_address"))
        if isinstance(stats, dict):
            entry["health"] = stats
        out.append(entry)
    return out


_UMLAUTS = str.maketrans({"ä": "ae", "ö": "oe", "ü": "ue", "Ä": "Ae", "Ö": "Oe", "Ü": "Ue", "ß": "ss"})


def piggyback_name(device: Mapping[str, Any], naming: str, prefix: str) -> str:
    raw = device.get("ieee_address") if naming == "ieee_address" else device.get("name")
    text = unicodedata.normalize("NFKD", str(raw or "").translate(_UMLAUTS))
    text = "".join(c for c in text if not unicodedata.combining(c))
    text = re.sub(r"[^A-Za-z0-9_.-]+", "_", text).strip("_.-")
    return f"{prefix}{text or 'unnamed'}"


def render_output(
    bridge: Mapping[str, Any],
    devices: Sequence[Mapping[str, Any]] | None,
    piggyback: bool = False,
    naming: str = "friendly_name",
    prefix: str = "",
) -> str:
    lines = ["<<<zigbee2mqtt_bridge:sep(0)>>>", json.dumps(bridge, sort_keys=True)]
    if devices is None:
        return "\n".join(lines) + "\n"
    lines += [
        "<<<zigbee2mqtt_devices:sep(0)>>>",
        json.dumps({"scope": "bridge", "piggyback": piggyback, "devices": list(devices)}, sort_keys=True),
    ]
    if piggyback:
        used: set[str] = set()
        for dev in devices:
            host = piggyback_name(dev, naming, prefix)
            if host in used:
                host = f"{host}_{str(dev.get('ieee_address') or '').removeprefix('0x')}"
            used.add(host)
            lines += [
                f"<<<<{host}>>>>",
                "<<<zigbee2mqtt_devices:sep(0)>>>",
                json.dumps({"scope": "device", "piggyback": False, "devices": [dev]}, sort_keys=True),
                "<<<<>>>>",
            ]
    return "\n".join(lines) + "\n"


# -----------------------------------------------------------------------------------
# main
# -----------------------------------------------------------------------------------
def parse_arguments(argv: Sequence[str]) -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--host", required=True, help="Zigbee2MQTT frontend host")
    p.add_argument("--port", type=int, default=8080, help="frontend port (default 8080)")
    p.add_argument("--protocol", choices=("http", "https"), default="http")
    p.add_argument("--no-tls-verify", action="store_true", help="frontend: do not verify the TLS certificate")
    p.add_argument("--base-url", default="/", help="frontend base_url (default /)")
    tok = p.add_mutually_exclusive_group()
    tok.add_argument("--token-id", help="password store reference of the frontend auth_token")
    tok.add_argument("--token", help="frontend auth_token (manual testing only)")
    p.add_argument(
        "--transport",
        choices=("websocket", "websocket_mqtt_fallback", "mqtt"),
        default="websocket",
    )
    p.add_argument("--mqtt-host", help="MQTT broker (default: --host)")
    p.add_argument("--mqtt-port", type=int, default=None, help="default 1883, or 8883 with --mqtt-tls")
    p.add_argument("--mqtt-tls", action="store_true")
    p.add_argument("--mqtt-no-tls-verify", action="store_true")
    p.add_argument("--mqtt-user")
    mpw = p.add_mutually_exclusive_group()
    mpw.add_argument("--mqtt-password-id", help="password store reference of the MQTT password")
    mpw.add_argument("--mqtt-password", help="MQTT password (manual testing only)")
    p.add_argument("--mqtt-base-topic", default="zigbee2mqtt")
    p.add_argument("--mqtt-listen", type=float, default=5.0, help="seconds to collect MQTT messages")
    p.add_argument("--timeout", type=float, default=10.0, help="connect/read timeout in seconds")
    p.add_argument("--piggyback", action="store_true", help="also emit every device as a piggyback host")
    p.add_argument("--piggyback-naming", choices=("friendly_name", "ieee_address"), default="friendly_name")
    p.add_argument("--piggyback-prefix", default="")
    p.add_argument("--exclude", action="append", default=[], help="regex on the friendly name (repeatable)")
    p.add_argument("--replay", help=argparse.SUPPRESS)  # JSON {topic: payload} instead of a connection
    p.add_argument("--debug", action="store_true", help="raise exceptions instead of reporting them")
    return p.parse_args(argv)


def _secret(ref: str | None, literal: str | None) -> str | None:
    if ref:
        from cmk.password_store.v1_unstable import dereference_secret

        return dereference_secret(ref).reveal()
    return literal


def collect(args: argparse.Namespace) -> tuple[Messages, str, str | None]:
    if args.replay:
        with open(args.replay, encoding="utf-8") as f:
            return json.load(f), "replay", None

    fallback_reason = None
    if args.transport in ("websocket", "websocket_mqtt_fallback"):
        try:
            return (
                ws_snapshot(
                    args.host,
                    args.port,
                    args.base_url,
                    _secret(args.token_id, args.token),
                    args.protocol == "https",
                    not args.no_tls_verify,
                    args.timeout,
                ),
                "websocket",
                None,
            )
        except TransportError as exc:
            if args.transport == "websocket":
                raise
            fallback_reason = f"WebSocket: {exc}"

    port = args.mqtt_port or (8883 if args.mqtt_tls else 1883)
    try:
        messages = mqtt_snapshot(
            args.mqtt_host or args.host,
            port,
            args.mqtt_user,
            _secret(args.mqtt_password_id, args.mqtt_password),
            args.mqtt_base_topic.strip("/"),
            args.mqtt_tls,
            not args.mqtt_no_tls_verify,
            args.timeout,
            args.mqtt_listen,
        )
    except TransportError as exc:
        if fallback_reason:
            raise TransportError(f"{fallback_reason}; MQTT fallback: {exc}") from exc
        raise
    return messages, "mqtt", fallback_reason


def main(argv: Sequence[str] | None = None) -> int:
    args = parse_arguments(sys.argv[1:] if argv is None else argv)
    exclude = [re.compile(p) for p in args.exclude]
    try:
        messages, transport, fallback_reason = collect(args)
    except TransportError as exc:
        if args.debug:
            raise
        sys.stdout.write(
            render_output({"reachable": False, "transport": args.transport, "error": str(exc)}, None)
        )
        return 0
    bridge = build_bridge(messages, transport, fallback_reason)
    devices = build_devices(messages, exclude) if "bridge/devices" in messages else None
    sys.stdout.write(
        render_output(bridge, devices, args.piggyback, args.piggyback_naming, args.piggyback_prefix)
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())
