monitor.py

# -*- coding: utf-8 -*-
"""TaskMonitor：訂閱 vision_events / skill_events，
驅動 TaskRuntime FSM 跳轉、更新 current_scene、發出 task_progress、觸發 Diagnostics。

對應架構教學文件 §5.3 / §10.4。
"""
from __future__ import annotations

import json
import logging
import threading
import time
from typing import Any, Callable, Dict, List, Optional

from .context import ContextStore
from .diagnostics import TaskDiagnostics
from .events import (
    FaultRecord,
    ProgressEvent,
    SkillEvent,
    VisionFrame,
    VisionObservation,
)
from .runtime import TaskRuntime, TaskState, Transition


logger = logging.getLogger(__name__)


class TaskMonitor:
    """事件驅動的任務監控器。

    - 訂閱 Redis channel `vision_events` / `skill_events`
    - 把事件路由到對應 TaskRuntime
    - 跳轉發生時：更新 context、廣播 task_progress、必要時 raise_fault
    - 同時把 vision 結果聚合成 task:current_scene
    """

    VISION_CHANNEL = "vision_events"
    SKILL_CHANNEL = "skill_events"
    PROGRESS_CHANNEL = "task_progress"
    SCENE_TTL_SEC = 60

    def __init__(
        self,
        collaborator: Any,
        context_store: ContextStore,
        diagnostics: TaskDiagnostics,
        on_progress: Optional[Callable[[ProgressEvent], None]] = None,
        on_done: Optional[Callable[[str, TaskRuntime], None]] = None,
        on_fault: Optional[Callable[[str, TaskRuntime, FaultRecord], None]] = None,
        queue_size_getter: Optional[Callable[[], Dict[str, int]]] = None,
    ) -> None:
        self.collaborator = collaborator
        self.context_store = context_store
        self.diagnostics = diagnostics
        self.on_progress = on_progress
        self.on_done = on_done
        self.on_fault = on_fault
        self.queue_size_getter = queue_size_getter or (lambda: {"active": 0, "pending": 0})

        # 一個 (task_id, package_id) 對應一個 runtime；任一鍵都能查到
        self._runtimes: Dict[str, TaskRuntime] = {}
        # 反向索引：package_id → set of task_ids
        self._pkg_to_tasks: Dict[str, set] = {}
        self._lock = threading.RLock()
        self._listener_threads: List[threading.Thread] = []
        self._stop_event = threading.Event()

    # --- 註冊 / 解除 ------------------------------------------------------
    def attach(
        self,
        task_id: str,
        package_id: str,
        subtask_id: str = "",
        **runtime_kwargs: Any,
    ) -> TaskRuntime:
        with self._lock:
            rt = TaskRuntime(
                task_id=task_id,
                package_id=package_id,
                subtask_id=subtask_id,
                **runtime_kwargs,
            )
            self._runtimes[task_id] = rt
            self._pkg_to_tasks.setdefault(package_id, set()).add(task_id)
            logger.info(
                "[Monitor] attach task_id=%s package_id=%s", task_id, package_id
            )
            return rt

    def detach(self, task_id: str) -> Optional[TaskRuntime]:
        with self._lock:
            rt = self._runtimes.pop(task_id, None)
            if rt:
                pkg_set = self._pkg_to_tasks.get(rt.package_id)
                if pkg_set:
                    pkg_set.discard(task_id)
                    if not pkg_set:
                        self._pkg_to_tasks.pop(rt.package_id, None)
            return rt

    def get(self, task_id: str) -> Optional[TaskRuntime]:
        with self._lock:
            return self._runtimes.get(task_id)

    def snapshot(self) -> Dict[str, Dict[str, Any]]:
        with self._lock:
            return {
                tid: {
                    "package_id": rt.package_id,
                    "state": rt.state.value,
                    "progress": rt.progress,
                    "time_in_state": rt.time_in_state,
                }
                for tid, rt in self._runtimes.items()
            }

    # --- 啟動 / 停止 listener --------------------------------------------
    def start_listeners(self) -> None:
        for ch, handler in (
            (self.VISION_CHANNEL, self._on_vision_raw),
            (self.SKILL_CHANNEL, self._on_skill_raw),
        ):
            t = threading.Thread(
                target=self._listen_loop,
                args=(ch, handler),
                daemon=True,
                name=f"monitor-{ch}",
            )
            t.start()
            self._listener_threads.append(t)
        logger.info("[Monitor] listeners started: %s", [t.name for t in self._listener_threads])

    def stop(self) -> None:
        self._stop_event.set()

    # --- 內部 listener wrapper -------------------------------------------
    def _listen_loop(self, channel: str, handler: Callable[[str], None]) -> None:
        while not self._stop_event.is_set():
            try:
                self.collaborator.listen(channel, handler)
                # listen 通常是阻塞的；正常返回意味著訂閱結束，退一步休眠等重連
                time.sleep(1.0)
            except Exception as e:  # noqa: BLE001
                logger.warning("[Monitor] listen(%s) error: %s, retry in 2s", channel, e)
                time.sleep(2.0)

    # --- Vision handler --------------------------------------------------
    def _on_vision_raw(self, raw: str) -> None:
        try:
            frame = VisionFrame.from_json(raw)
        except Exception as e:  # noqa: BLE001
            logger.warning("[Monitor] bad vision payload: %s", e)
            return
        self._update_current_scene(frame)
        for obs in frame.observations:
            self._route_vision(obs)

    def _update_current_scene(self, frame: VisionFrame) -> None:
        packages = []
        for obs in frame.observations:
            packages.append(
                {
                    "id": obs.package_id,
                    "zone": obs.zone,
                    "barcode_face": obs.barcode_face,
                    "size": obs.size,
                    "score": obs.score,
                }
            )
        snapshot = {
            "ts": frame.ts,
            "frame_id": frame.frame_id,
            "packages": packages,
        }
        self.context_store.update_current_scene(snapshot)

    def _route_vision(self, obs: VisionObservation) -> None:
        with self._lock:
            tids = list(self._pkg_to_tasks.get(obs.package_id, set()))
        for tid in tids:
            self._tick(tid, obs)

    # --- Skill handler ---------------------------------------------------
    def _on_skill_raw(self, raw: str) -> None:
        try:
            evt = SkillEvent.from_json(raw)
        except Exception as e:  # noqa: BLE001
            logger.warning("[Monitor] bad skill payload: %s", e)
            return
        # 沒指定 task_id 就廣播到所有 runtime；指定了就只給對應的
        if evt.task_id:
            self._tick(evt.task_id, evt)
        else:
            with self._lock:
                tids = list(self._runtimes.keys())
            for tid in tids:
                self._tick(tid, evt)

    # --- FSM tick --------------------------------------------------------
    def _tick(self, task_id: str, evt) -> None:
        with self._lock:
            rt = self._runtimes.get(task_id)
            if not rt or rt.is_terminal():
                return
        try:
            trans: Optional[Transition] = rt.evaluate(evt)
        except Exception as e:  # noqa: BLE001
            logger.exception("[Monitor] evaluate error: %s", e)
            self.diagnostics.soft_error(task_id, e)
            return
        if trans is None:
            return
        self._handle_transition(task_id, rt, trans)

    def _handle_transition(
        self, task_id: str, rt: TaskRuntime, trans: Transition
    ) -> None:
        logger.info(
            "[Monitor] task=%s pkg=%s %s -> %s (%s)",
            task_id,
            rt.package_id,
            trans.prev.value,
            trans.curr.value,
            trans.reason,
        )

        qs = self.queue_size_getter()
        progress_evt = ProgressEvent(
            task_id=task_id,
            subtask_id=rt.subtask_id,
            package_id=rt.package_id,
            fsm_state=trans.curr.value,
            fsm_prev_state=trans.prev.value,
            progress=rt.progress,
            queue_active=qs.get("active", 0),
            queue_pending=qs.get("pending", 0),
        )
        self._broadcast_progress(progress_evt)

        # 寫關鍵幀
        ctx = self.context_store.get(task_id)
        if ctx:
            self.context_store.append_keyframe(
                ctx,
                {
                    "ts": trans.ts,
                    "fsm_state": trans.curr.value,
                    "reason": trans.reason,
                    "package_id": rt.package_id,
                },
            )

        # 故障處理
        if trans.curr == TaskState.FAULT:
            code = trans.fault_code or "TM-007"
            fault = self.diagnostics.raise_fault(
                task_id=task_id,
                code=code,
                details={"reason": trans.reason, "prev_state": trans.prev.value},
                package_id=rt.package_id,
                subtask_id=rt.subtask_id,
            )
            if self.on_fault:
                try:
                    self.on_fault(task_id, rt, fault)
                except Exception as e:  # noqa: BLE001
                    logger.exception("[Monitor] on_fault cb error: %s", e)

        # 完成處理
        if trans.curr == TaskState.DONE:
            self.context_store.bump_stat("success_count", 1)
            if self.on_done:
                try:
                    self.on_done(task_id, rt)
                except Exception as e:  # noqa: BLE001
                    logger.exception("[Monitor] on_done cb error: %s", e)

    def _broadcast_progress(self, evt: ProgressEvent) -> None:
        try:
            self.collaborator.send(self.PROGRESS_CHANNEL, json.dumps(evt.to_dict()))
        except Exception as e:  # noqa: BLE001
            logger.warning("[Monitor] broadcast progress failed: %s", e)
        if self.on_progress:
            try:
                self.on_progress(evt)
            except Exception as e:  # noqa: BLE001
                logger.exception("[Monitor] on_progress cb error: %s", e)


__all__ = ["TaskMonitor"]