yolo/backend/alerting.py

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import logging
import os
import threading
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from io import BytesIO
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from PIL import Image, ImageDraw, ImageFont

from .storage import DeliveryStatsStore, utc_now_iso

LOGGER = logging.getLogger(__name__)
ALERT_CLASSES = ("fire", "smoke")
CHANNEL_NAMES = {"wechat": "企业微信", "feishu": "飞书"}


@dataclass(slots=True)
class SessionState:
    consecutive: dict[str, int] = field(
        default_factory=lambda: defaultdict(int)
    )
    last_alert_at: dict[str, float] = field(
        default_factory=lambda: defaultdict(float)
    )
    last_seen_at: float = field(default_factory=time.monotonic)


class AlertManager:
    def __init__(
        self,
        webhook_url: str | None = None,
        wechat_webhook_url: str | None = None,
        feishu_webhook_url: str | None = None,
        feishu_secret: str | None = None,
        confirm_frames: int = 3,
        cooldown_seconds: float = 60.0,
        session_ttl_seconds: float = 3600.0,
        stats_store: DeliveryStatsStore | None = None,
    ) -> None:
        self.wechat_webhook_url = (
            wechat_webhook_url
            or webhook_url
            or os.getenv("WECHAT_WEBHOOK_URL")
        )
        self.feishu_webhook_url = (
            feishu_webhook_url or os.getenv("FEISHU_WEBHOOK_URL")
        )
        self.feishu_secret = feishu_secret or os.getenv("FEISHU_SECRET")
        self.confirm_frames = max(1, confirm_frames)
        self.cooldown_seconds = max(0.0, cooldown_seconds)
        self.session_ttl_seconds = max(60.0, session_ttl_seconds)
        self.stats_store = stats_store
        self._states: dict[str, SessionState] = {}
        self._lock = threading.Lock()
        self._delivery_stats = {
            channel: {
                "delivered": 0,
                "failed": 0,
                "tests": 0,
                "last_status": "idle",
                "last_delivery_at": None,
                "last_test_at": None,
                "last_error": None,
            }
            for channel in CHANNEL_NAMES
        }
        self._executor = ThreadPoolExecutor(
            max_workers=4,
            thread_name_prefix="alert-dispatch",
        )

    @property
    def channels(self) -> dict[str, bool]:
        return {
            "wechat": bool(self.wechat_webhook_url),
            "feishu": bool(self.feishu_webhook_url),
        }

    @property
    def enabled(self) -> bool:
        return any(self.channels.values())


    @property
    def channel_status(self) -> list[dict[str, Any]]:
        configured = self.channels
        if self.stats_store is not None:
            saved = self.stats_store.get_all()
            stats = {
                channel: {
                    "delivered": 0,
                    "failed": 0,
                    "tests": 0,
                    "last_status": "idle",
                    "last_delivery_at": None,
                    "last_test_at": None,
                    "last_error": None,
                    **saved.get(channel, {}),
                }
                for channel in CHANNEL_NAMES
            }
        else:
            with self._lock:
                stats = {
                    channel: dict(values)
                    for channel, values in self._delivery_stats.items()
                }
        return [
            {
                "id": channel,
                "name": CHANNEL_NAMES[channel],
                "configured": configured[channel],
                "signed": channel == "feishu" and bool(self.feishu_secret),
                "capabilities": (
                    ["Markdown 告警", "检测截图"]
                    if channel == "wechat"
                    else ["交互卡片", "签名校验"]
                ),
                **stats[channel],
            }
            for channel in CHANNEL_NAMES
        ]

    def test_channel(self, channel: str) -> dict[str, Any]:
        if channel not in CHANNEL_NAMES:
            raise ValueError(f"Unsupported robot channel: {channel}")
        if not self.channels[channel]:
            raise ValueError(f"{CHANNEL_NAMES[channel]}机器人尚未配置")

        try:
            if channel == "wechat":
                self._send_wechat_test()
            else:
                self._send_feishu_test()
        except (HTTPError, URLError, TimeoutError, ValueError) as error:
            self._record_delivery(
                channel,
                success=False,
                error=str(error),
                test=True,
            )
            raise RuntimeError(str(error)) from error
        self._record_delivery(channel, success=True, test=True)
        return next(
            status for status in self.channel_status if status["id"] == channel
        )
    def evaluate(
        self,
        session_id: str,
        detections: list[dict[str, Any]],
        image: Image.Image,
    ) -> dict[str, Any]:
        now = time.monotonic()
        present = {
            detection["class"]
            for detection in detections
            if detection.get("class") in ALERT_CLASSES
        }
        triggered: list[str] = []

        with self._lock:
            self._prune_sessions(now)
            state = self._states.setdefault(session_id, SessionState())
            state.last_seen_at = now
            for class_name in ALERT_CLASSES:
                state.consecutive[class_name] = (
                    state.consecutive[class_name] + 1
                    if class_name in present
                    else 0
                )
                ready = state.consecutive[class_name] >= self.confirm_frames
                cooldown_elapsed = (
                    now - state.last_alert_at[class_name]
                    >= self.cooldown_seconds
                )
                if ready and cooldown_elapsed:
                    state.last_alert_at[class_name] = now
                    triggered.append(class_name)
            consecutive = dict(state.consecutive)

        if triggered and self.enabled:
            self._executor.submit(
                self._send_alerts,
                annotate_image(image, detections),
                triggered,
                detections,
            )
        return {
            "triggered": bool(triggered),
            "classes": triggered,
            "confirmed_frames": self.confirm_frames,
            "consecutive": consecutive,
            "cooldown_seconds": self.cooldown_seconds,
            "notification_enabled": self.enabled,
            "notification_channels": self.channels,
        }

    def reset(self, session_id: str) -> None:
        with self._lock:
            self._states.pop(session_id, None)

    def close(self) -> None:
        """Shut down the alert dispatch executor (call on service shutdown)."""
        self._executor.shutdown(wait=True, cancel_futures=True)

    def _prune_sessions(self, now: float) -> None:
        expired = [
            session_id
            for session_id, state in self._states.items()
            if now - state.last_seen_at > self.session_ttl_seconds
        ]
        for session_id in expired:
            del self._states[session_id]

    def _send_alerts(
        self,
        image: Image.Image,
        triggered: list[str],
        detections: list[dict[str, Any]],
    ) -> None:
        if self.wechat_webhook_url:
            try:
                self._send_wechat_alert(image, triggered, detections)
                self._record_delivery("wechat", success=True)
            except (HTTPError, URLError, TimeoutError, ValueError) as error:
                self._record_delivery("wechat", success=False, error=str(error))
                LOGGER.error("WeChat alert failed: %s", error)
        if self.feishu_webhook_url:
            try:
                self._send_feishu_alert(triggered, detections)
                self._record_delivery("feishu", success=True)
            except (HTTPError, URLError, TimeoutError, ValueError) as error:
                self._record_delivery("feishu", success=False, error=str(error))
                LOGGER.error("Feishu alert failed: %s", error)


    def _record_delivery(
        self,
        channel: str,
        *,
        success: bool,
        error: str | None = None,
        test: bool = False,
    ) -> None:
        if self.stats_store is not None:
            self.stats_store.record(
                channel,
                success=success,
                error=error,
                test=test,
            )
            return
        now = utc_now_iso()
        with self._lock:
            stats = self._delivery_stats[channel]
            stats["last_status"] = "success" if success else "failed"
            stats["last_error"] = None if success else error
            stats["delivered" if success else "failed"] += 1
            if test:
                stats["tests"] += 1
                stats["last_test_at"] = now
            else:
                stats["last_delivery_at"] = now

    def _send_wechat_test(self) -> None:
        if not self.wechat_webhook_url:
            return
        post_json(
            self.wechat_webhook_url,
            {
                "msgtype": "markdown",
                "markdown": {
                    "content": (
                        "燧安机器人连接测试\n"
                        "> 消息通道工作正常\n"
                        f"> 测试时间：{time.strftime('%Y-%m-%d %H:%M:%S')}"
                    )
                },
            },
        )

    def _send_feishu_test(self) -> None:
        if not self.feishu_webhook_url:
            return
        payload: dict[str, Any] = {
            "msg_type": "interactive",
            "card": {
                "header": {
                    "template": "green",
                    "title": {
                        "tag": "plain_text",
                        "content": "燧安机器人连接测试",
                    },
                },
                "elements": [
                    {
                        "tag": "markdown",
                        "content": (
                            "**消息通道工作正常**\n"
                            f"测试时间：{time.strftime('%Y-%m-%d %H:%M:%S')}"
                        ),
                    }
                ],
            },
        }
        if self.feishu_secret:
            timestamp = str(int(time.time()))
            payload["timestamp"] = timestamp
            payload["sign"] = feishu_signature(
                timestamp,
                self.feishu_secret,
            )
        post_json(self.feishu_webhook_url, payload)
    def _send_wechat_alert(
        self,
        image: Image.Image,
        triggered: list[str],
        detections: list[dict[str, Any]],
    ) -> None:
        if not self.wechat_webhook_url:
            return
        target_text, max_confidence = alert_summary(triggered, detections)
        message = (
            f"🔥 烟火检测告警\n"
            f"> 检测目标：{target_text}\n"
            f"> 最高置信度：{max_confidence:.1%}\n"
            f"> 请及时查看现场视频。"
        )
        post_json(
            self.wechat_webhook_url,
            {"msgtype": "markdown", "markdown": {"content": message}},
        )
        image_bytes = encode_jpeg(image)
        post_json(
            self.wechat_webhook_url,
            {
                "msgtype": "image",
                "image": {
                    "base64": base64.b64encode(image_bytes).decode("ascii"),
                    "md5": hashlib.md5(
                        image_bytes,
                        usedforsecurity=False,
                    ).hexdigest(),
                },
            },
        )

    def _send_feishu_alert(
        self,
        triggered: list[str],
        detections: list[dict[str, Any]],
    ) -> None:
        if not self.feishu_webhook_url:
            return
        target_text, max_confidence = alert_summary(triggered, detections)
        payload: dict[str, Any] = {
            "msg_type": "interactive",
            "card": {
                "config": {"wide_screen_mode": True},
                "header": {
                    "template": "red",
                    "title": {
                        "tag": "plain_text",
                        "content": "烟火检测告警",
                    },
                },
                "elements": [
                    {
                        "tag": "markdown",
                        "content": (
                            f"**检测目标：** {target_text}\n"
                            f"**最高置信度：** {max_confidence:.1%}\n"
                            f"**告警时间：** "
                            f"{time.strftime('%Y-%m-%d %H:%M:%S')}\n"
                            f"请及时查看现场视频。"
                        ),
                    }
                ],
            },
        }
        if self.feishu_secret:
            timestamp = str(int(time.time()))
            payload["timestamp"] = timestamp
            payload["sign"] = feishu_signature(
                timestamp,
                self.feishu_secret,
            )
        post_json(self.feishu_webhook_url, payload)


def alert_summary(
    triggered: list[str],
    detections: list[dict[str, Any]],
) -> tuple[str, float]:
    labels = {"fire": "火焰", "smoke": "烟雾"}
    target_text = "、".join(labels[name] for name in triggered)
    max_confidence = max(
        (
            float(detection.get("confidence", 0.0))
            for detection in detections
            if detection.get("class") in triggered
        ),
        default=0.0,
    )
    return target_text, max_confidence


def feishu_signature(timestamp: str, secret: str) -> str:
    string_to_sign = f"{timestamp}\n{secret}".encode()
    digest = hmac.new(string_to_sign, digestmod=hashlib.sha256).digest()
    return base64.b64encode(digest).decode("ascii")


def post_json(url: str, payload: dict[str, Any]) -> dict[str, Any]:
    request = Request(
        url,
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urlopen(request, timeout=10) as response:
        result = json.loads(response.read().decode("utf-8"))
    error_code = result.get(
        "errcode",
        result.get("code", result.get("StatusCode", 0)),
    )
    if error_code != 0:
        raise ValueError(
            result.get("errmsg")
            or result.get("msg")
            or result.get("StatusMessage")
            or "Unknown robot webhook error"
        )
    return result


def annotate_image(
    image: Image.Image,
    detections: list[dict[str, Any]],
) -> Image.Image:
    annotated = image.copy()
    draw = ImageDraw.Draw(annotated)
    font = ImageFont.load_default()
    colors = {"fire": "#ff3b30", "smoke": "#00a89b"}
    for detection in detections:
        box = detection.get("box")
        if not box or len(box) != 4:
            continue
        class_name = str(detection.get("class", "target"))
        confidence = float(detection.get("confidence", 0.0))
        color = colors.get(class_name, "#ffd166")
        coordinates = tuple(round(value) for value in box)
        draw.rectangle(coordinates, outline=color, width=4)
        draw.text(
            (coordinates[0] + 4, max(0, coordinates[1] - 16)),
            f"{class_name} {confidence:.0%}",
            fill=color,
            font=font,
        )
    return annotated


def encode_jpeg(image: Image.Image, max_bytes: int = 1_900_000) -> bytes:
    working = image.convert("RGB")
    quality = 88
    while True:
        output = BytesIO()
        working.save(output, format="JPEG", quality=quality, optimize=True)
        payload = output.getvalue()
        if len(payload) <= max_bytes:
            return payload
        if quality > 55:
            quality -= 10
            continue
        width, height = working.size
        working = working.resize(
            (max(1, int(width * 0.8)), max(1, int(height * 0.8)))
        )
        quality = 75