netx/netx_api/ume_alarm_ws_subscription.py

"""UME alarm WSS subscription state APIs and manual control."""
from __future__ import annotations

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

from sqlalchemy.orm import Session

from . import ume_alarm_ws_state as ws_state
from .config import settings
from .ume_alarm_subscription_store import (
    DEFAULT_SUBSCRIPTION_KEY,
    clear_subscription,
    load_subscription,
    save_subscription,
)
from .ume_alarm_ws_log import (
    append_ws_log,
    is_ume_subscription_missing_error,
    parse_subscription_id_from_already_exists_error,
)
from .ume_client import UMEClient
from .ume_sync_service import (
    apply_alarm_to_current,
    extract_alarm_from_notification,
    normalize_yang_alarm,
    _utc_now_naive,
)

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


def mark_ume_subscription_lost(reason: str) -> None:
    detail = str(reason or "").strip()[:500]
    with ws_state._ume_lost_lock:
        if ws_state._ume_subscription_lost and detail == ws_state._ume_subscription_lost_reason:
            return
        ws_state._ume_subscription_lost = True
        ws_state._ume_subscription_lost_reason = detail
    append_ws_log(
        f"UME 侧订阅已丢失: {detail[:200]}",
        level="warning",
        dedup_cooldown_s=ws_state._LOOP_STATUS_COOLDOWN_S,
    )
    request_ws_reconnect()


def clear_ume_subscription_lost_flag() -> None:
    with ws_state._ume_lost_lock:
        ws_state._ume_subscription_lost = False
        ws_state._ume_subscription_lost_reason = ""


def is_ume_subscription_lost() -> bool:
    with ws_state._ume_lost_lock:
        return bool(ws_state._ume_subscription_lost)


def _server_subscription_lost_fields() -> dict[str, Any]:
    with ws_state._ume_lost_lock:
        lost = bool(ws_state._ume_subscription_lost)
        reason = str(ws_state._ume_subscription_lost_reason or "")
    return {
        "server_subscription_lost": lost,
        "server_subscription_lost_reason": reason,
    }


def _set_ws_connection_state(state: str, *, detail: str = "") -> None:
    with ws_state._ws_connection_lock:
        ws_state._ws_connection_state = str(state or "").strip() or "unknown"
        ws_state._ws_connection_detail = str(detail or "").strip()[:240]


def get_ws_connection_status() -> dict[str, Any]:
    with ws_state._ws_connection_lock:
        state = str(ws_state._ws_connection_state or "init")
        detail = str(ws_state._ws_connection_detail or "")
    return {
        "state": state,
        "label": ws_state._WS_CONNECTION_LABELS.get(state, state),
        "detail": detail,
    }


def _notify_ws_connection(
    state: str,
    *,
    detail: str = "",
    subscription_id: str = "",
    on_status: Callable[[str], None] | None = None,
    log_level: str = "info",
    log_cooldown_s: float = ws_state._WS_NOTIFY_COOLDOWN_S,
) -> None:
    _set_ws_connection_state(state, detail=detail)
    label = ws_state._WS_CONNECTION_LABELS.get(state, state)
    log_msg = str(detail or label).strip()
    notify_key = f"{state}\0{log_msg}\0{subscription_id}"
    now = time.monotonic()
    should_log = True
    with ws_state._WS_NOTIFY_LOCK:
        if notify_key == ws_state._WS_NOTIFY_LAST_KEY and now - ws_state._WS_NOTIFY_LAST_TS < max(1.0, float(log_cooldown_s)):
            should_log = False
        else:
            ws_state._WS_NOTIFY_LAST_KEY = notify_key
            ws_state._WS_NOTIFY_LAST_TS = now
    if should_log:
        append_ws_log(
            log_msg,
            level=log_level,
            subscription_id=subscription_id,
            dedup=False,
        )
    if on_status is not None:
        on_status(f"ws:{state}")


def _parse_ws_message(raw: str) -> dict[str, Any] | None:
    text = str(raw or "").strip()
    if not text:
        return None
    try:
        data = json.loads(text)
    except json.JSONDecodeError:
        return None
    return data if isinstance(data, dict) else None


def process_alarm_notification(db: Session, payload: dict[str, Any]) -> tuple[str, bool]:
    alarm = extract_alarm_from_notification(payload)
    if alarm is None:
        return "skipped", False
    return apply_alarm_to_current(db, alarm, touch_ts=_utc_now_naive())


def _delete_subscription_on_ume(client: UMEClient, subscription_id: str) -> None:
    sub_id = str(subscription_id or "").strip()
    if not sub_id:
        return
    client.delete_alarm_subscription(sub_id)


def _wait_for_shared_token(client: UMEClient, *, timeout_s: float = 300.0) -> bool:
    deadline = time.time() + max(5.0, float(timeout_s))
    while time.time() < deadline:
        client._sync_token_from_store()
        if client.has_valid_token():
            return True
        time.sleep(2.0)
    return False


def _set_active_subscription(subscription_id: str, wss_uri: str, *, topic: str = "ALARM") -> None:
    with ws_state._subscription_lock:
        ws_state._subscription_id = str(subscription_id or "").strip()
        ws_state._subscription_uri = str(wss_uri or "").strip()
        ws_state._subscription_topic = str(topic or "ALARM").strip() or "ALARM"


def _clear_active_subscription() -> None:
    with ws_state._subscription_lock:
        ws_state._subscription_id = ""
        ws_state._subscription_uri = ""
        ws_state._subscription_topic = "ALARM"


def get_active_subscription() -> tuple[str, str]:
    with ws_state._subscription_lock:
        return str(ws_state._subscription_id or ""), str(ws_state._subscription_uri or "")


def load_persisted_subscription() -> bool:
    """Load subscription from DB into memory (call on process startup)."""
    db = SessionLocal()
    try:
        loaded = load_subscription(db)
        if loaded is None:
            _clear_active_subscription()
            return False
        sub_id, uri, topic = loaded
        _set_active_subscription(sub_id, uri, topic=topic)
        append_ws_log(f"loaded persisted subscription id={sub_id}", subscription_id=sub_id)
        return True
    finally:
        db.close()


def get_subscription_status() -> dict[str, Any]:
    with ws_state._subscription_lock:
        sub_id = str(ws_state._subscription_id or "")
        uri = str(ws_state._subscription_uri or "")
        topic = str(ws_state._subscription_topic or "ALARM")
    return {
        "active": bool(sub_id and uri),
        "subscription_id": sub_id,
        "wss_uri": uri,
        "topic": topic,
        **_server_subscription_lost_fields(),
    }


def is_wss_active_for_current_alarms() -> bool:
    """True when WSS subscription is active and should own ume_alarms_current updates."""
    if not bool(getattr(settings, "ume_alarm_ws_enabled", True)):
        return False
    if is_ume_subscription_lost():
        return False
    return bool(get_subscription_status().get("active"))


def get_current_alarms_mode() -> str:
    return "wss" if is_wss_active_for_current_alarms() else "rest"


def get_alarms_coordination_status() -> dict[str, Any]:
    wss_active = is_wss_active_for_current_alarms()
    skip_when_ws = bool(getattr(settings, "ume_sync_alarms_current_skip_when_ws", True))
    return {
        "current_alarms_mode": get_current_alarms_mode(),
        "wss_active_for_current_alarms": wss_active,
        "scheduled_sync_skipped": bool(wss_active and skip_when_ws),
    }


def clear_local_alarm_subscription_manual(db: Session) -> dict[str, Any]:
    """Drop persisted/in-memory subscription without calling UME delete."""
    clear_subscription(db, cache_key=DEFAULT_SUBSCRIPTION_KEY)
    db.commit()
    _clear_active_subscription()
    clear_ume_subscription_lost_flag()
    request_ws_reconnect()
    append_ws_log("cleared local subscription record (UME delete skipped)")
    return get_subscription_status()


def request_ws_reconnect() -> None:
    ws_state._ws_wake_event.set()


def establish_alarm_subscription_manual(
    client: UMEClient,
    db: Session,
    *,
    force_reestablish: bool = False,
) -> dict[str, Any]:
    """Manually establish ALARM subscription (UME API + persist). Does not open WSS."""
    if force_reestablish or is_ume_subscription_lost():
        clear_local_alarm_subscription_manual(db)

    existing = load_subscription(db)
    if existing is not None:
        sub_id, uri, topic = existing
        _set_active_subscription(sub_id, uri, topic=topic)
        request_ws_reconnect()
        append_ws_log(f"establish skipped: already in DB id={sub_id}", subscription_id=sub_id)
        status_out = get_subscription_status()
        return {**status_out, "already_exists": True}

    mem_id, mem_uri = get_active_subscription()
    if mem_id and mem_uri:
        request_ws_reconnect()
        append_ws_log(f"establish skipped: in-memory id={mem_id}", subscription_id=mem_id)
        status_out = get_subscription_status()
        return {**status_out, "already_exists": True}

    if not _wait_for_shared_token(client, timeout_s=120.0):
        raise RuntimeError("ume_ws_no_valid_token:wait_timeout")
    client._sync_token_from_store()
    if not client.has_valid_token():
        raise RuntimeError("ume_ws_no_valid_token")

    topic = str(getattr(settings, "ume_notification_topic", "ALARM") or "ALARM").strip() or "ALARM"
    try:
        sub_id, uri = client.establish_alarm_subscription(topic=topic)
    except RuntimeError as exc:
        orphan_id = parse_subscription_id_from_already_exists_error(str(exc))
        if not orphan_id:
            raise
        append_ws_log(
            f"establish: orphan on UME id={orphan_id}, deleting then retry",
            level="warning",
            subscription_id=orphan_id,
        )
        try:
            _delete_subscription_on_ume(client, orphan_id)
        except RuntimeError as del_exc:
            raise RuntimeError(
                f"ume_orphan_subscription_delete_failed:id={orphan_id}:{str(del_exc)[:180]}"
            ) from del_exc
        sub_id, uri = client.establish_alarm_subscription(topic=topic)
    save_subscription(db, subscription_id=sub_id, wss_uri=uri, topic=topic)
    db.commit()

    with ws_state._shutdown_lock:
        ws_state._active_client = client
    _set_active_subscription(sub_id, uri, topic=topic)
    clear_ume_subscription_lost_flag()
    request_ws_reconnect()
    append_ws_log(f"manual establish subscription id={sub_id}", subscription_id=sub_id)
    return {**get_subscription_status(), "already_exists": False}


def cancel_alarm_subscription_manual(
    client: UMEClient,
    db: Session,
    *,
    force_clear_local: bool = False,
) -> dict[str, Any]:
    """Manually cancel subscription (UME delete + clear DB + drop WSS)."""
    loaded = load_subscription(db)
    sub_id, _uri = get_active_subscription()
    if not sub_id and loaded is not None:
        sub_id = str(loaded[0] or "")
    ume_already_missing = False
    if sub_id:
        client.refresh_if_needed()
        try:
            _delete_subscription_on_ume(client, sub_id)
        except RuntimeError as exc:
            if is_ume_subscription_missing_error(str(exc)):
                ume_already_missing = True
                mark_ume_subscription_lost(str(exc))
                if not force_clear_local:
                    return {
                        **get_subscription_status(),
                        "ok": False,
                        "ume_already_missing": True,
                        "needs_local_cleanup": True,
                        "message": (
                            "UME 侧告警订阅已不存在或已过期（服务器订阅已丢失）。"
                            "请确认是否清除本地订阅记录，然后重新建立订阅。"
                        ),
                    }
            else:
                raise

    clear_subscription(db, cache_key=DEFAULT_SUBSCRIPTION_KEY)
    db.commit()
    _clear_active_subscription()
    clear_ume_subscription_lost_flag()
    request_ws_reconnect()
    append_ws_log(
        f"manual cancel subscription id={sub_id or '(none)'}",
        subscription_id=sub_id,
    )
    status_out = get_subscription_status()
    return {
        **status_out,
        "ok": True,
        "cleared_local": True,
        "ume_already_missing": ume_already_missing,
    }