#!/usr/bin/env python3
"""Collect Zigbee2MQTT readings in RRD files and render graphs."""

from __future__ import annotations

import argparse
import html
import json
import logging
import math
import re
import shutil
import signal
import subprocess
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import paho.mqtt.client as mqtt
import yaml


LOG = logging.getLogger("zigbee-rrd")
BASE_DIR = Path(__file__).resolve().parent
STOP = threading.Event()


@dataclass(frozen=True)
class Sensor:
    name: str
    topic: str
    slug: str


@dataclass(frozen=True)
class Config:
    mqtt_host: str
    mqtt_port: int
    mqtt_username: str | None
    mqtt_password: str | None
    interval: int
    timeout: int
    data_dir: Path
    graph_dir: Path
    sensors: tuple[Sensor, ...]
    rrd: dict[str, Any]
    graph: dict[str, Any]
    html: dict[str, Any]
    log_level: str
    log_format: str


def safe_slug(name: str) -> str:
    slug = re.sub(r"[^a-zA-Z0-9_-]+", "_", name.strip()).strip("_")
    if not slug:
        raise ValueError(f"Invalid sensor name {name!r}")
    return slug


def load_config(path: Path) -> Config:
    try:
        raw = yaml.safe_load(path.read_text(encoding="utf-8"))
    except FileNotFoundError as error:
        raise SystemExit(
            f"Configuration file {path} was not found"
        ) from error
    except yaml.YAMLError as error:
        raise SystemExit(f"Invalid YAML in {path}: {error}") from error

    if not isinstance(raw, dict):
        raise SystemExit("The YAML root must be a mapping")

    mqtt_config = raw.get("mqtt", {})
    configured_sensors = raw.get("sensors", [])
    if not isinstance(configured_sensors, list) or not configured_sensors:
        raise SystemExit("The configuration requires a non-empty sensors list")

    collection = raw.get("collection", {})
    excluded_prefix = str(collection.get("excluded_name_prefix", "CHANGE_ME"))
    sensors: list[Sensor] = []
    slugs: set[str] = set()
    topics: set[str] = set()
    for item in configured_sensors:
        try:
            name = str(item["name"]).strip()
            topic = str(item["topic"]).strip().strip("/")
        except (KeyError, TypeError) as error:
            raise SystemExit("Every sensor requires a name and topic") from error
        if not name or not topic:
            raise SystemExit("Sensor names and topics cannot be empty")
        if excluded_prefix and name.casefold().startswith(excluded_prefix.casefold()):
            continue
        slug = safe_slug(name)
        if slug in slugs or topic in topics:
            raise SystemExit(f"Duplicate sensor slug or topic: {name!r}, {topic!r}")
        slugs.add(slug)
        topics.add(topic)
        sensors.append(Sensor(name, topic, slug))

    if not sensors:
        raise SystemExit(
            f"No active sensors remain after excluding names beginning with {excluded_prefix!r}"
        )

    storage = raw.get("storage", {})
    logging_config = raw.get("logging", {})
    interval = int(collection.get("interval_seconds", 600))
    timeout = int(collection.get("mqtt_timeout_seconds", 15))
    if interval < 60 or timeout < 1:
        raise SystemExit("interval_seconds must be >= 60 and timeout must be positive")

    def config_path(key: str, default: str) -> Path:
        value = Path(storage.get(key, default)).expanduser()
        return value if value.is_absolute() else path.parent / value

    return Config(
        mqtt_host=str(mqtt_config.get("host", "127.0.0.1")),
        mqtt_port=int(mqtt_config.get("port", 1883)),
        mqtt_username=mqtt_config.get("username"),
        mqtt_password=mqtt_config.get("password"),
        interval=interval,
        timeout=timeout,
        data_dir=config_path("data_directory", "/var/lib/rrd"),
        graph_dir=config_path("graph_directory", "/var/www/html/sensors"),
        sensors=tuple(sensors),
        rrd=raw.get("rrd", {}),
        graph=raw.get("graph", {}),
        html=raw.get("html", {}),
        log_level=str(logging_config.get("level", "INFO")),
        log_format=str(logging_config.get("format", "%(asctime)s %(levelname)s %(message)s")),
    )


def run_rrdtool(*arguments: str) -> None:
    try:
        subprocess.run(["rrdtool", *arguments], check=True, text=True)
    except subprocess.CalledProcessError as error:
        raise RuntimeError(f"rrdtool {' '.join(arguments[:2])} failed") from error


def create_rrd(path: Path, config: Config) -> None:
    if path.exists():
        return
    rrd = config.rrd
    data_sources = rrd.get("data_sources", {})
    archives = rrd.get("archives", [])
    if not data_sources or not archives:
        raise SystemExit("rrd.data_sources and rrd.archives must be configured")

    def data_source(name: str) -> str:
        settings = data_sources[name]
        minimum = settings.get("minimum", "U")
        maximum = settings.get("maximum", "U")
        return (
            f"DS:{name}:{settings['type']}:{settings['heartbeat_seconds']}:"
            f"{minimum}:{maximum}"
        )

    archive_arguments = [
        "RRA:{consolidation_function}:{xfiles_factor}:{steps}:{rows}".format(**archive)
        for archive in archives
    ]
    run_rrdtool(
        "create", str(path), "--start", str(rrd["start"]),
        "--step", str(rrd.get("step_seconds", config.interval)),
        data_source("temperature"), data_source("humidity"), data_source("battery"),
        *archive_arguments,
    )
    LOG.info("Created %s", path)


def number_or_unknown(reading: dict[str, Any], key: str) -> str:
    value = reading.get(key)
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return "U"
    number = float(value)
    return str(number) if math.isfinite(number) else "U"


def update_rrd(path: Path, reading: dict[str, Any]) -> None:
    values = ":".join(
        number_or_unknown(reading, key)
        for key in ("temperature", "humidity", "battery")
    )
    run_rrdtool(
        "update", str(path), "--template", "temperature:humidity:battery", f"N:{values}"
    )


def render_graphs(sensor: Sensor, rrd_path: Path, config: Config) -> None:
    graph = config.graph
    temperature = graph["temperature"]
    humidity = graph["humidity"]
    battery = graph["battery"]
    temperature_axis = graph["axes"]["temperature"]
    percentage_axis = graph["axes"]["percentage"]
    temperature_minimum = float(temperature_axis["minimum"])
    temperature_maximum = float(temperature_axis["maximum"])
    percentage_minimum = float(percentage_axis["minimum"])
    percentage_maximum = float(percentage_axis["maximum"])
    temperature_span = temperature_maximum - temperature_minimum
    percentage_span = percentage_maximum - percentage_minimum
    if temperature_span <= 0 or percentage_span <= 0:
        raise SystemExit("Graph axis maximums must be greater than their minimums")

    percentage_to_temperature = temperature_span / percentage_span
    right_axis_scale = percentage_span / temperature_span
    right_axis_shift = percentage_minimum - temperature_minimum * right_axis_scale
    percentage_transform = (
        f"{percentage_minimum},-,{percentage_to_temperature},*,{temperature_minimum},+"
    )
    reference_lines = [
        f"HRULE:{float(line['value'])}#{line['color']}"
        for line in temperature_axis.get("reference_lines", [])
    ]
    optional_flags = []
    if graph.get("slope_mode", False):
        optional_flags.append("--slope-mode")

    for suffix, period in graph["periods"].items():
        output = config.graph_dir / f"{sensor.slug}_{suffix}.png"
        run_rrdtool(
            "graph", str(output), "--start", str(period["start"]),
            "--end", str(graph["end"]),
            "--title", f"{sensor.name} - {period['label']}",
            "--vertical-label", str(temperature_axis["label"]),
            "--right-axis", f"{right_axis_scale}:{right_axis_shift}",
            "--right-axis-label", str(percentage_axis["label"]),
            "--lower-limit", str(temperature_minimum),
            "--upper-limit", str(temperature_maximum),
            "--rigid",
            "--width", str(graph["width"]), "--height", str(graph["height"]),
            *optional_flags,
            f"DEF:temp={rrd_path}:temperature:AVERAGE",
            f"DEF:humidity={rrd_path}:humidity:AVERAGE",
            f"DEF:battery={rrd_path}:battery:AVERAGE",
            f"CDEF:humidity_scaled=humidity,{percentage_transform}",
            f"CDEF:battery_scaled=battery,{percentage_transform}",
            f"AREA:temp#{temperature['color']}:{temperature['legend']}",
            *reference_lines,
            f"LINE{humidity['line_width']}:humidity_scaled#{humidity['color']}:{humidity['legend']}",
            f"LINE{battery['line_width']}:battery_scaled#{battery['color']}:{battery['legend']}:dashes={battery['dashes']}",
            f"GPRINT:temp:LAST:{temperature['last_value_format']}",
            f"GPRINT:humidity:LAST:{humidity['last_value_format']}",
            f"GPRINT:battery:LAST:{battery['last_value_format']}",
        )


def render_dashboard(config: Config) -> None:
    settings = config.html
    filename = str(settings.get("filename", "index.html"))
    if Path(filename).name != filename:
        raise SystemExit("html.filename must be a filename, not a path")

    title = str(settings.get("title", "Zigbee sensor history"))
    periods = config.graph["periods"]
    default_period = str(settings.get("default_period", next(iter(periods))))
    if default_period not in periods:
        raise SystemExit(f"html.default_period {default_period!r} is not a graph period")
    refresh_seconds = int(settings.get("refresh_seconds", config.interval))
    if refresh_seconds < 0:
        raise SystemExit("html.refresh_seconds cannot be negative")

    buttons = "\n".join(
        f'<button type="button" data-period="{html.escape(key)}">'
        f'{html.escape(str(period["label"]))}</button>'
        for key, period in periods.items()
    )
    panels: list[str] = []
    for period_key, period in periods.items():
        cards = "\n".join(
            "<article class=\"sensor-card\">"
            f"<h2>{html.escape(sensor.name)}</h2>"
            f'<a href="{html.escape(sensor.slug)}_{html.escape(period_key)}.png">'
            f'<img src="{html.escape(sensor.slug)}_{html.escape(period_key)}.png" '
            f'data-source="{html.escape(sensor.slug)}_{html.escape(period_key)}.png" '
            f'alt="{html.escape(sensor.name)} — {html.escape(str(period["label"]))}">'
            "</a></article>"
            for sensor in config.sensors
        )
        panels.append(
            f'<section class="period-panel" data-period="{html.escape(period_key)}" hidden>'
            f'<div class="sensor-grid">{cards}</div></section>'
        )

    document = f"""<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>{html.escape(title)}</title>
  <style>
    :root {{ color-scheme: light dark; font-family: system-ui, sans-serif; }}
    body {{ max-width: 1800px; margin: 0 auto; padding: 1.25rem; background: #f3f5f7; color: #17202a; }}
    header {{ display: flex; align-items: center; justify-content: space-between; gap: 1rem; flex-wrap: wrap; }}
    h1 {{ margin: 0; font-size: clamp(1.5rem, 3vw, 2.25rem); }}
    .period-picker {{ display: flex; gap: .5rem; flex-wrap: wrap; }}
    button {{ border: 1px solid #aab2bd; border-radius: .45rem; padding: .6rem 1rem; background: white; color: #17202a; cursor: pointer; font: inherit; }}
    button[aria-selected="true"] {{ background: #2457a6; border-color: #2457a6; color: white; }}
    .sensor-grid {{ display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 1rem; margin-top: 1.25rem; }}
    .sensor-card {{ overflow: hidden; border: 1px solid #d5d8dc; border-radius: .65rem; background: white; box-shadow: 0 2px 8px #0001; }}
    .sensor-card h2 {{ margin: 0; padding: .75rem 1rem .25rem; font-size: 1.1rem; }}
    .sensor-card img {{ display: block; width: 100%; height: auto; }}
    @media (max-width: 900px) {{ .sensor-grid {{ grid-template-columns: 1fr; }} }}
    @media (prefers-color-scheme: dark) {{
      body {{ background: #111820; color: #eef2f5; }}
      .sensor-card {{ background: #1b2530; border-color: #344454; }}
      button {{ background: #1b2530; color: #eef2f5; border-color: #526577; }}
    }}
  </style>
</head>
<body>
  <header>
    <h1>{html.escape(title)}</h1>
    <nav class="period-picker" aria-label="Graph period">{buttons}</nav>
  </header>
  <main>{''.join(panels)}</main>
  <script>
    const periods = {json.dumps(list(periods))};
    const configuredDefault = {json.dumps(default_period)};
    const stored = localStorage.getItem('sensor-period');
    const initial = periods.includes(stored) ? stored : configuredDefault;

    function selectPeriod(period) {{
      if (!periods.includes(period)) return;
      document.querySelectorAll('.period-panel').forEach((panel) => {{
        panel.hidden = panel.dataset.period !== period;
      }});
      document.querySelectorAll('[data-period]').forEach((element) => {{
        if (element.tagName === 'BUTTON') {{
          element.setAttribute('aria-selected', String(element.dataset.period === period));
        }}
      }});
      localStorage.setItem('sensor-period', period);
    }}

    document.querySelectorAll('.period-picker button').forEach((button) => {{
      button.addEventListener('click', () => selectPeriod(button.dataset.period));
    }});
    selectPeriod(initial);

    const refreshSeconds = {refresh_seconds};
    if (refreshSeconds > 0) {{
      setInterval(() => {{
        const stamp = Date.now();
        document.querySelectorAll('img[data-source]').forEach((image) => {{
          image.src = `${{image.dataset.source}}?t=${{stamp}}`;
        }});
      }}, refreshSeconds * 1000);
    }}
  </script>
</body>
</html>
"""
    output = config.graph_dir / filename
    temporary = output.with_suffix(output.suffix + ".tmp")
    temporary.write_text(document, encoding="utf-8")
    temporary.replace(output)
    LOG.info("Generated dashboard %s", output)


def read_sensors(config: Config) -> dict[str, dict[str, Any]]:
    readings: dict[str, dict[str, Any]] = {}
    complete = threading.Event()
    topics = {sensor.topic for sensor in config.sensors}
    lock = threading.Lock()

    def on_connect(client, userdata, flags, reason_code, properties):
        if reason_code != 0:
            LOG.error("MQTT connection failed: %s", reason_code)
            complete.set()
            return
        for topic in topics:
            client.subscribe(topic)

    def on_message(client, userdata, message):
        if message.topic not in topics:
            return
        try:
            payload = json.loads(message.payload.decode("utf-8"))
            if not isinstance(payload, dict):
                raise ValueError("payload is not an object")
        except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as error:
            LOG.warning("Ignoring invalid message on %s: %s", message.topic, error)
            return
        with lock:
            readings[message.topic] = payload
            if readings.keys() >= topics:
                complete.set()

    client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
    if config.mqtt_username is not None:
        client.username_pw_set(config.mqtt_username, config.mqtt_password)
    client.on_connect = on_connect
    client.on_message = on_message
    try:
        client.connect(config.mqtt_host, config.mqtt_port)
        client.loop_start()
        complete.wait(config.timeout)
    except OSError as error:
        LOG.error("Could not connect to %s:%s: %s", config.mqtt_host, config.mqtt_port, error)
    finally:
        client.disconnect()
        client.loop_stop()

    missing = topics - readings.keys()
    if missing:
        LOG.warning("No reading received for: %s", ", ".join(sorted(missing)))
    return readings


def collect_once(config: Config) -> None:
    config.data_dir.mkdir(parents=True, exist_ok=True)
    config.graph_dir.mkdir(parents=True, exist_ok=True)
    readings = read_sensors(config)
    for sensor in config.sensors:
        rrd_path = config.data_dir / f"{sensor.slug}.rrd"
        create_rrd(rrd_path, config)
        reading = readings.get(sensor.topic)
        if reading is not None:
            update_rrd(rrd_path, reading)
            LOG.info(
                "%s: temperature=%s C humidity=%s %% battery=%s %%",
                sensor.name, reading.get("temperature", "unknown"),
                reading.get("humidity", "unknown"), reading.get("battery", "unknown"),
            )
        render_graphs(sensor, rrd_path, config)
    render_dashboard(config)


def render_only(config: Config) -> None:
    config.graph_dir.mkdir(parents=True, exist_ok=True)
    for sensor in config.sensors:
        rrd_path = config.data_dir / f"{sensor.slug}.rrd"
        if rrd_path.exists():
            render_graphs(sensor, rrd_path, config)
        else:
            LOG.warning("Skipping %s; %s does not exist", sensor.name, rrd_path)
    render_dashboard(config)


def handle_signal(signum, frame) -> None:
    LOG.info("Stopping")
    STOP.set()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=BASE_DIR / "config.yaml")
    modes = parser.add_mutually_exclusive_group()
    modes.add_argument("--once", action="store_true", help="collect and graph once")
    modes.add_argument("--graph-only", action="store_true", help="only regenerate graphs")
    parser.add_argument("--verbose", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if shutil.which("rrdtool") is None:
        raise SystemExit("rrdtool is not installed; run: sudo apt install rrdtool")
    config = load_config(args.config.resolve())
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else config.log_level.upper(),
        format=config.log_format,
    )
    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)

    if args.graph_only:
        render_only(config)
        return 0
    while not STOP.is_set():
        started = time.monotonic()
        try:
            collect_once(config)
        except RuntimeError:
            LOG.exception("Collection cycle failed")
            if args.once:
                return 1
        if args.once:
            break
        STOP.wait(max(0.0, config.interval - (time.monotonic() - started)))
    return 0


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