netx/netx_api/ume_alarm_ws_consumer.py

"""UME alarm WSS consumer session and background loop."""
from __future__ import annotations

import json
import logging
import ssl
import threading
import time
from typing import Any, Callable

from . import ume_alarm_ws_state as ws_state
from .config import settings
from .db import SessionLocal
from .ume_alarm_ws_log import (
    append_ws_log,
    is_startup_alarm_sync_pending,
    is_ume_subscription_missing_error,
)
from .ume_alarm_ws_subscription import (
    _notify_ws_connection,
    _parse_ws_message,
    _wait_for_shared_token,
    clear_ume_subscription_lost_flag,
    get_active_subscription,
    is_ume_subscription_lost,
    mark_ume_subscription_lost,
    process_alarm_notification,
)
from .ume_client import UMEClient
from .ume_sync_service import (
    _alarm_key,
    extract_alarm_from_notification,
    normalize_yang_alarm,
)

_ws_log = logging.getLogger("netx.ume.alarm_ws")


def _run_ws_session(
    client: UMEClient,
    *,
    wss_uri: str,
    subscription_id: str,
    on_status: Callable[[str], None] | None = None,
    stop_event: threading.Event | None = None,
) -> None:
    import websocket

    headers = client.ws_auth_headers()
    header_list = [f"{k}: {v}" for k, v in headers.items() if str(v or "").strip()]

    sslopt: dict[str, Any] | None = None
    if wss_uri.lower().startswith("wss://"):
        if client.verify_tls:
            sslopt = {"cert_reqs": ssl.CERT_REQUIRED}
        else:
            sslopt = {"cert_reqs": ssl.CERT_NONE}

    closed = threading.Event()

    def _on_message(_ws: Any, message: str) -> None:
        payload = _parse_ws_message(message)
        if payload is None:
            append_ws_log(
                "message ignored: invalid JSON",
                level="warning",
                subscription_id=subscription_id,
                dedup_cooldown_s=ws_state._WS_IGNORE_COOLDOWN_S,
            )
            return
        db = SessionLocal()
        try:
            alarm = extract_alarm_from_notification(payload)
            if alarm is None:
                append_ws_log(
                    "收包忽略: 非告警通知",
                    level="warning",
                    subscription_id=subscription_id,
                    dedup_cooldown_s=ws_state._WS_IGNORE_COOLDOWN_S,
                )
                return
            norm = normalize_yang_alarm(alarm) or {}
            alarm_key = _alarm_key(norm) or "?"
            action, changed = apply_alarm_to_current(db, alarm, touch_ts=_utc_now_naive())
            if changed:
                db.commit()
            else:
                db.rollback()
            if not changed:
                return
            label = ws_state._WS_ALARM_ACTION_LABEL.get(action, action)
            status_msg = f"{label} key={alarm_key}"
            append_ws_log(status_msg, subscription_id=subscription_id, dedup=False)
            try:
                from .key_alert_forward import maybe_forward_key_alert

                maybe_forward_key_alert(db, norm=norm, alarm_key=alarm_key, action=action)
            except Exception as fwd_exc:
                append_ws_log(
                    f"key alert forward failed: {str(fwd_exc)[:160]}",
                    level="error",
                    subscription_id=subscription_id,
                )
        except Exception as exc:
            db.rollback()
            append_ws_log(f"alarm apply failed: {str(exc)[:200]}", level="error", subscription_id=subscription_id)
        finally:
            db.close()

    def _on_error(_ws: Any, error: Any) -> None:
        err_text = str(error)
        detail = f"ws error: {err_text[:200]}"
        if is_ume_subscription_missing_error(err_text):
            mark_ume_subscription_lost(err_text)
            _notify_ws_connection(
                "subscription_lost",
                detail=detail,
                subscription_id=subscription_id,
                on_status=on_status,
                log_level="warning",
            )
        else:
            _notify_ws_connection(
                "error",
                detail=detail,
                subscription_id=subscription_id,
                on_status=on_status,
                log_level="error",
            )

    def _on_close(_ws: Any, close_status_code: Any, close_msg: Any) -> None:
        detail = f"ws closed code={close_status_code} msg={str(close_msg or '')[:120]}"
        _notify_ws_connection("disconnected", detail=detail, subscription_id=subscription_id, on_status=on_status)
        closed.set()

    def _on_open(_ws: Any) -> None:
        _notify_ws_connection(
            "connected",
            detail=f"ws connected uri={wss_uri[:120]}",
            subscription_id=subscription_id,
            on_status=on_status,
        )

    ws_app = websocket.WebSocketApp(
        wss_uri,
        header=header_list,
        on_open=_on_open,
        on_message=_on_message,
        on_error=_on_error,
        on_close=_on_close,
    )

    kwargs: dict[str, Any] = {"ping_interval": 30, "ping_timeout": 20}
    if sslopt is not None:
        kwargs["sslopt"] = sslopt

    thread = threading.Thread(
        target=lambda: ws_app.run_forever(**kwargs),
        name=f"ume-alarm-ws-{subscription_id[:8]}",
        daemon=True,
    )
    thread.start()

    while thread.is_alive() and not closed.is_set():
        if stop_event is not None and stop_event.is_set():
            try:
                ws_app.close()
            except Exception:
                pass
            break
        sub_id_now, _ = get_active_subscription()
        if not sub_id_now or sub_id_now != subscription_id:
            try:
                ws_app.close()
            except Exception:
                pass
            break
        if ws_state._ws_wake_event.is_set():
            try:
                ws_app.close()
            except Exception:
                pass
            break
        time.sleep(0.5)

    if thread.is_alive():
        thread.join(timeout=5.0)
    ws_state._ws_wake_event.clear()


def run_alarm_ws_consumer_loop(
    client: UMEClient,
    *,
    on_status: Callable[[str], None] | None = None,
    stop_event: threading.Event | None = None,
    is_paused: Callable[[], bool] | None = None,
) -> None:
    """
    Connect WSS only when a persisted/manual subscription exists.
    Never auto-establish subscription; reconnect same uri after disconnect.
    """
    backoff_s = 2.0
    max_backoff_s = 120.0

    def _loop_status(
        state: str,
        *,
        detail: str = "",
        subscription_id: str = "",
        log_level: str = "info",
    ) -> None:
        _notify_ws_connection(
            state,
            detail=detail,
            subscription_id=subscription_id,
            on_status=on_status,
            log_level=log_level,
            log_cooldown_s=ws_state._LOOP_STATUS_COOLDOWN_S,
        )

    while stop_event is None or not stop_event.is_set():
        if is_paused is not None and is_paused():
            _loop_status("paused", log_level="warning")
            time.sleep(1.0)
            continue

        if not ws_state._STARTUP_ALARM_SYNC_GATE.is_set():
            _loop_status("waiting_startup_alarm_sync", log_level="info")
            ws_state._STARTUP_ALARM_SYNC_GATE.wait(timeout=1.0)
            continue

        if is_ume_subscription_lost():
            with ws_state._ume_lost_lock:
                lost_detail = str(ws_state._ume_subscription_lost_reason or "")
            _loop_status(
                "subscription_lost",
                detail=lost_detail or "UME subscription missing on server",
                log_level="warning",
            )
            ws_state._ws_wake_event.wait(timeout=5.0)
            ws_state._ws_wake_event.clear()
            continue

        subscription_id, wss_uri = get_active_subscription()
        if not subscription_id or not wss_uri:
            _loop_status("no_subscription")
            ws_state._ws_wake_event.wait(timeout=2.0)
            ws_state._ws_wake_event.clear()
            continue

        try:
            if not _wait_for_shared_token(client, timeout_s=120.0):
                _loop_status("waiting_token", log_level="warning", subscription_id=subscription_id)
                time.sleep(2.0)
                continue

            _loop_status("connecting", detail=f"connecting {wss_uri[:160]}", subscription_id=subscription_id)
            _run_ws_session(
                client,
                wss_uri=wss_uri,
                subscription_id=subscription_id,
                on_status=on_status,
                stop_event=stop_event,
            )
            backoff_s = 2.0
        except RuntimeError as exc:
            if "ume_ws_no_valid_token" in str(exc):
                _loop_status("waiting_token", log_level="warning", subscription_id=subscription_id)
                time.sleep(2.0)
                continue
            _loop_status("error", detail=f"session error: {str(exc)[:200]}", subscription_id=subscription_id, log_level="error")
        except Exception as exc:
            _loop_status("error", detail=f"session error: {str(exc)[:200]}", subscription_id=subscription_id, log_level="error")

        sub_id_after, uri_after = get_active_subscription()
        if stop_event is not None and stop_event.is_set():
            break
        if not sub_id_after or not uri_after:
            _loop_status("no_subscription")
            continue

        _loop_status("reconnecting", detail=f"reconnect in {int(backoff_s)}s", subscription_id=sub_id_after)
        slept = 0.0
        while slept < backoff_s:
            if stop_event is not None and stop_event.is_set():
                return
            if ws_state._ws_wake_event.is_set():
                ws_state._ws_wake_event.clear()
                break
            time.sleep(0.5)
            slept += 0.5
        backoff_s = min(max_backoff_s, backoff_s * 2.0)


def shutdown_ws_consumer() -> None:
    """Process exit: close WSS only; subscription remains on UME until user cancels."""
    request_ws_reconnect()
    ws_state._ws_wake_event.set()


def start_ume_alarm_ws_consumer(
    client: UMEClient,
    *,
    on_status: Callable[[str], None] | None = None,
    stop_event: threading.Event | None = None,
    is_paused: Callable[[], bool] | None = None,
) -> threading.Thread:
    ev = stop_event if stop_event is not None else threading.Event()

    def _target() -> None:
        if not bool(getattr(settings, "ume_alarm_ws_enabled", True)):
            return
        if not str(client.base_url or "").strip():
            append_ws_log("consumer disabled: no UME base URL", level="warning")
            if on_status is not None:
                on_status("ws:disabled_no_base_url")
            return
        append_ws_log("ws consumer thread started")
        with ws_state._shutdown_lock:
            ws_state._active_client = client
        run_alarm_ws_consumer_loop(
            client,
            on_status=on_status,
            stop_event=ev,
            is_paused=is_paused,
        )

    thread = threading.Thread(target=_target, name="ume-alarm-ws-consumer", daemon=True)
    thread.start()
    return thread