home-automation/app/services/dsmr_ingest.py

"""DSMR MQTT ingest, keyed by durable meter-source identity."""

from __future__ import annotations

import json
import logging
import threading
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING

import sqlalchemy.exc
from sqlalchemy import select
from sqlalchemy.orm import Session

from app.db import get_session_local
from app.models.energy import DsmrReading, Meter
from app.models.meter_source import MeterSource, MeterSourceBinding, MeterSourceChannel

if TYPE_CHECKING:
    from app.config import Settings

logger = logging.getLogger(__name__)


@dataclass(frozen=True, slots=True)
class DsmrSourceSnapshot:
    """The only DSMR runtime configuration a network handler may use."""

    source_id: int
    topic: str
    tariff_topic: str
    sample_interval_s: int
    broker_host: str = ""
    broker_port: int = 1883
    username: str = ""
    password: str = ""
    tls_enabled: bool = False


_subscriptions: dict[int, DsmrSourceSnapshot] = {}
_subscription_lock = threading.RLock()
# A configuration value is not an ownership identity: disable and re-enable
# can produce an equal snapshot.  Each installed handler therefore captures a
# fresh token and verifies object identity before it can write or update tariff.
_subscription_tokens: dict[int, object] = {}
_reconcile_lock = threading.RLock()
_tariffs: dict[int, int] = {}
_tariff_lock = threading.Lock()
# Kept only for legacy direct test callers of set_current_tariff(value). Runtime
# handlers never write this value; production callers resolve a binding first.
_current_tariff: int | None = None
_UNSET = object()


def get_current_tariff(meter_source_id: int | None = None) -> int | None:
    """Return a source tariff, or resolve the active electricity binding.

    The no-argument form is retained for the pre-M8 expose integration.  It
    opens a short session to select the current electricity binding, so a
    source's MQTT callback can never make another source's tariff current.
    ``_current_tariff`` is solely a test-era fallback when no binding database
    is available; runtime MQTT handlers do not update it.
    """
    if meter_source_id is None:
        session_local = get_session_local()
        session = session_local()
        try:
            source_id = _current_electricity_source_id(session, datetime.now(timezone.utc))
            if source_id is not None:
                return get_current_tariff(source_id)
        except Exception:
            logger.debug("DSMR legacy tariff lookup could not resolve a binding", exc_info=True)
        finally:
            session.close()
        return _current_tariff
    with _tariff_lock:
        return _tariffs.get(meter_source_id)


def set_current_tariff(meter_source_id: int, value: int | None | object = _UNSET) -> None:
    """Set or clear an individual source's tariff state."""
    global _current_tariff
    if value is _UNSET:
        # Compatibility with older direct callers. Do not route runtime source
        # updates through this global fallback.
        _current_tariff = meter_source_id if meter_source_id in (1, 2) else None
        return
    with _tariff_lock:
        if value is None:
            _tariffs.pop(meter_source_id, None)
        else:
            _tariffs[meter_source_id] = value


def _current_electricity_source_id(session: Session, at: datetime) -> int | None:
    """Return the DSMR source bound to electricity at ``at``, if any."""
    return session.scalar(
        select(MeterSource.id)
        .join(MeterSourceChannel, MeterSourceChannel.source_id == MeterSource.id)
        .join(MeterSourceBinding, MeterSourceBinding.channel_id == MeterSourceChannel.id)
        .join(Meter, Meter.id == MeterSourceBinding.meter_id)
        .where(
            Meter.commodity == "electricity",
            MeterSourceBinding.started_at <= at,
            (MeterSourceBinding.ended_at.is_(None)) | (MeterSourceBinding.ended_at > at),
        )
        .order_by(MeterSourceBinding.started_at.desc())
        .limit(1)
    )


def get_current_electricity_tariff(session: Session, at: datetime | None = None) -> int | None:
    """Resolve tariff through the current electricity binding, never globally."""
    source_id = _current_electricity_source_id(session, at or datetime.now(timezone.utc))
    if source_id is None:
        return None
    return get_current_tariff(source_id)


def handle_tariff_message(payload_bytes: bytes, meter_source_id: int) -> None:
    """Parse one source's tariff payload without raising in paho's thread."""
    try:
        raw = (
            payload_bytes.decode("utf-8", errors="replace").strip()
            if isinstance(payload_bytes, (bytes, bytearray))
            else str(payload_bytes).strip()
        )
        value = int(raw)
        if value in (1, 2):
            set_current_tariff(meter_source_id, value)
    except Exception:
        logger.debug("DSMR tariff payload ignored for source_id=%s", meter_source_id)


def _snapshot(source: MeterSource) -> DsmrSourceSnapshot:
    config = source.config
    return DsmrSourceSnapshot(
        source_id=source.id,
        topic=str(config.get("topic", "dsmr/json")),
        tariff_topic=str(config.get("tariff_topic", "")),
        sample_interval_s=int(config.get("sample_interval_s", 10)),
        broker_host=str(config.get("broker_host", "")),
        broker_port=int(config.get("broker_port", 1883)),
        username=str(config.get("username", "")),
        password=str(config.get("password", "")),
        tls_enabled=bool(config.get("tls_enabled", False)),
    )


def _enabled_snapshots() -> list[DsmrSourceSnapshot]:
    session_local = get_session_local()
    session = session_local()
    try:
        sources = session.scalars(
            select(MeterSource).where(MeterSource.kind == "dsmr_mqtt", MeterSource.enabled.is_(True))
        ).all()
        return [_snapshot(source) for source in sources]
    finally:
        session.close()


def apply_dsmr_subscription(_settings: "Settings | None" = None) -> None:
    """Reconcile enabled DSMR source subscriptions from the database.

    The optional, ignored settings parameter preserves the old config-route
    call shape while ensuring flat DSMR settings are no longer read.
    """
    from app.integrations.mqtt import mqtt_manager

    try:
        desired = {snapshot.source_id: snapshot for snapshot in _enabled_snapshots()}
    except Exception:
        logger.exception("DSMR subscription reconcile failed while reading sources")
        return

    # Do not retain _subscription_lock while stopping MQTT clients: a callback
    # may currently hold it through its complete dispatch, and MqttManager
    # waits for that callback before teardown returns.
    with _reconcile_lock:
        # Every source owns a distinct MQTT client, so equal topics from
        # different sources/brokers are dispatchable.  A telegram and tariff
        # topic on the *same* client would overwrite one handler, however.
        for source_id, snapshot in list(desired.items()):
            if snapshot.tariff_topic and snapshot.topic == snapshot.tariff_topic:
                logger.error("DSMR source_id=%s rejected: telegram/tariff topic collision", source_id)
                desired.pop(source_id)

        with _subscription_lock:
            stale_source_ids = [
                source_id
                for source_id, current in _subscriptions.items()
                if desired.get(source_id) != current
            ]
            for source_id in stale_source_ids:
                _subscriptions.pop(source_id, None)
                _subscription_tokens.pop(source_id, None)
                set_current_tariff(source_id, None)

        for source_id in stale_source_ids:
            mqtt_manager.remove_source(source_id)

        for source_id, snapshot in desired.items():
            with _subscription_lock:
                current = _subscriptions.get(source_id)
            if current == snapshot and mqtt_manager.source_is_active(source_id):
                continue
            if current is not None:
                # The client went inactive outside reconcile. Invalidate its
                # old token before rebuilding the same snapshot.
                with _subscription_lock:
                    if _subscriptions.get(source_id) == current:
                        _subscriptions.pop(source_id, None)
                        _subscription_tokens.pop(source_id, None)
                mqtt_manager.remove_source(source_id)
            token = object()
            handlers = {
                snapshot.topic: lambda payload, captured=snapshot, captured_token=token: (
                    handle_captured_message(payload, captured, captured_token)
                )
            }
            if snapshot.tariff_topic:
                handlers[snapshot.tariff_topic] = (
                    lambda payload, captured=snapshot, captured_token=token: (
                        handle_captured_tariff_message(payload, captured, captured_token)
                    )
                )
            with _subscription_lock:
                _subscriptions[source_id] = snapshot
                _subscription_tokens[source_id] = token
            applied = mqtt_manager.replace_source(
                source_id,
                host=snapshot.broker_host,
                port=snapshot.broker_port,
                username=snapshot.username,
                password=snapshot.password,
                tls_enabled=snapshot.tls_enabled,
                subscriptions=handlers,
            )
            if not applied:
                with _subscription_lock:
                    if _subscription_tokens.get(source_id) is token:
                        _subscriptions.pop(source_id, None)
                        _subscription_tokens.pop(source_id, None)


def handle_message(payload_bytes: bytes, snapshot: DsmrSourceSnapshot) -> None:
    """Persist one down-sampled frame under its captured source identity."""
    try:
        _handle_message_inner(payload_bytes, snapshot)
    except Exception:
        logger.exception("DSMR ingest handler failed for source_id=%s (swallowed)", snapshot.source_id)


def handle_captured_message(
    payload_bytes: bytes, snapshot: DsmrSourceSnapshot, token: object | None = None
) -> None:
    """Run a broker callback only while its exact source generation is active."""
    with _subscription_lock:
        if token is not None:
            if _subscription_tokens.get(snapshot.source_id) is not token:
                return
        elif _subscriptions.get(snapshot.source_id) != snapshot:
            return
        handle_message(payload_bytes, snapshot)


def handle_captured_tariff_message(
    payload_bytes: bytes, snapshot: DsmrSourceSnapshot, token: object | None = None
) -> None:
    """Ignore tariff callbacks retained from a removed/replaced source."""
    with _subscription_lock:
        if token is not None:
            if _subscription_tokens.get(snapshot.source_id) is not token:
                return
        elif _subscriptions.get(snapshot.source_id) != snapshot:
            return
        handle_tariff_message(payload_bytes, snapshot.source_id)


def _handle_message_inner(payload_bytes: bytes, snapshot: DsmrSourceSnapshot) -> None:
    try:
        data = json.loads(payload_bytes)
    except (json.JSONDecodeError, ValueError):
        return
    if not isinstance(data, dict):
        return
    try:
        raw_ts = data["timestamp"]
        ts_utc = datetime.fromisoformat(raw_ts.replace("Z", "+00:00"))
        if ts_utc.tzinfo is None:
            ts_utc = ts_utc.replace(tzinfo=timezone.utc)
    except (KeyError, ValueError, TypeError, AttributeError):
        return
    if snapshot.sample_interval_s > 0 and ts_utc.second % snapshot.sample_interval_s:
        return
    telegram_id = data.get("id")
    if telegram_id is not None and not isinstance(telegram_id, int):
        telegram_id = None

    session_local = get_session_local()
    session = session_local()
    try:
        exists = session.scalar(
            select(DsmrReading.id).where(
                DsmrReading.meter_source_id == snapshot.source_id,
                DsmrReading.recorded_at == ts_utc,
            )
        )
        if exists is None:
            session.add(
                DsmrReading(
                    meter_source_id=snapshot.source_id,
                    recorded_at=ts_utc,
                    telegram_id=telegram_id,
                    payload=data,
                )
            )
            session.commit()
    except sqlalchemy.exc.IntegrityError:
        session.rollback()
    except Exception:
        session.rollback()
        logger.exception("DSMR database write failed for source_id=%s", snapshot.source_id)
    finally:
        session.close()