scheduler.py

# -*- coding: utf-8 -*-
"""TaskScheduler：包住既有 GlobalAgent，提供狀態機 + 隊列調度 + 模式切換 + 重規劃。

對應架構教學文件 §5.2 / §10.6、interfaces.md §8.3。
"""
from __future__ import annotations

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

from .context import ContextStore, TaskContext, new_task_id
from .diagnostics import DiagnoseDecision, TaskDiagnostics
from .monitor import TaskMonitor
from .queue import TaskQueueManager
from .runtime import TaskRuntime, TaskState


logger = logging.getLogger(__name__)


class SchedState(str, Enum):
    IDLE = "IDLE"
    RUNNING = "RUNNING"
    PAUSED = "PAUSED"
    FAULT = "FAULT"
    RESET = "RESET"
    DONE = "DONE"


class SchedMode(str, Enum):
    TELEOP = "teleop"
    AUTO = "auto"
    SEMI = "semi"


# 包裹 ID 正則（interfaces.md §附錄 A：P_yyyyMMddNNNNN）
PACKAGE_ID_RE = re.compile(r"P[_-]?\w{5,}")


def extract_package_id(subtask: Dict[str, Any]) -> Optional[str]:
    """從 subtask dict 取出 package_id，優先用顯式欄位，否則從字串正則。"""
    pid = subtask.get("package_id")
    if pid:
        return pid
    text = subtask.get("subtask") or ""
    m = PACKAGE_ID_RE.search(text)
    return m.group(0) if m else None


class TaskScheduler:
    """任務調度狀態機。

    生命週期：
      IDLE → start_task() → RUNNING → (pause/resume) → DONE / FAULT / RESET

    與外部模塊的關係：
      - 透過 `global_agent` 觸發規劃 + 派發（沿用 RoboOS 既有路徑）
      - 透過 `monitor` 註冊 / 解除 TaskRuntime
      - 透過 `queue` 管理子任務
      - 接收 `diagnostics.scheduler_action_callback` 的決策來決定下一步
    """

    def __init__(
        self,
        global_agent: Any,
        context_store: ContextStore,
        queue: TaskQueueManager,
        monitor: TaskMonitor,
        diagnostics: TaskDiagnostics,
        max_replan_attempts: int = 3,
    ) -> None:
        self.global_agent = global_agent
        self.context_store = context_store
        self.queue = queue
        self.monitor = monitor
        self.diagnostics = diagnostics
        self.max_replan_attempts = max_replan_attempts

        self.state = SchedState.IDLE
        self.mode = SchedMode.AUTO
        self.current_task_id: Optional[str] = None
        self._replan_attempts: Dict[str, int] = {}
        self._lock = threading.RLock()

        # 把 monitor 的回呼接好（任務完成 / 故障時 scheduler 介入）
        self.monitor.on_done = self._on_runtime_done  # type: ignore[assignment]
        self.monitor.on_fault = self._on_runtime_fault  # type: ignore[assignment]
        # 讓 monitor 在發 progress 時知道隊列大小
        self.monitor.queue_size_getter = lambda: {
            "active": self.queue.active_count,
            "pending": self.queue.pending_count,
        }

    # --- 對外：控制 API ---------------------------------------------------
    def start_task(self, task_text: str, refresh: bool = True) -> Dict[str, Any]:
        """接收自然語言任務，呼叫 Planner 拆解、入隊、創建 runtimes。"""
        with self._lock:
            if self.mode == SchedMode.TELEOP:
                self.diagnostics.raise_fault(
                    task_id=self.current_task_id or "n/a",
                    code="TM-008",
                    details={"task_text": task_text},
                )
                return {"status": "blocked", "reason": "teleop_mode"}

            task_id = new_task_id()
            ctx = TaskContext(task_id=task_id, task_text=task_text, state="RUNNING")
            self.context_store.save(ctx)
            self.current_task_id = task_id

            try:
                reasoning = self.global_agent.publish_global_task(
                    task_text, refresh=refresh, task_id=task_id
                )
            except Exception as e:  # noqa: BLE001
                logger.exception("[Scheduler] planner error: %s", e)
                self.diagnostics.raise_fault(
                    task_id=task_id,
                    code="TM-001",
                    details={"exception": repr(e)},
                )
                self.state = SchedState.FAULT
                return {"status": "error", "reason": "plan_fail", "task_id": task_id}

            subtasks: List[Dict[str, Any]] = list(
                reasoning.get("subtask_list", []) if isinstance(reasoning, dict) else []
            )
            if not subtasks:
                self.diagnostics.raise_fault(task_id=task_id, code="TM-001")
                self.state = SchedState.FAULT
                return {
                    "status": "error",
                    "reason": "no_subtasks",
                    "task_id": task_id,
                    "reasoning": reasoning,
                }

            self._ingest_subtasks(task_id, subtasks, ctx)
            self.state = SchedState.RUNNING
            return {
                "status": "ok",
                "task_id": task_id,
                "subtask_count": len(subtasks),
                "reasoning": reasoning,
            }

    def pause(self) -> None:
        with self._lock:
            if self.state == SchedState.RUNNING:
                self.state = SchedState.PAUSED
                logger.info("[Scheduler] paused")

    def resume(self) -> None:
        with self._lock:
            if self.state in (SchedState.PAUSED, SchedState.FAULT):
                self.state = SchedState.RUNNING
                logger.info("[Scheduler] resumed")

    def abort(self, reason: str = "") -> None:
        with self._lock:
            logger.warning("[Scheduler] abort: %s", reason)
            self.state = SchedState.FAULT
            self.queue.clear()
            if self.current_task_id:
                ctx = self.context_store.get(self.current_task_id)
                if ctx:
                    ctx.state = "FAULT"
                    self.context_store.save(ctx)

    def reset(self) -> None:
        with self._lock:
            logger.info("[Scheduler] reset")
            self.queue.clear()
            self.state = SchedState.RESET
            self._replan_attempts.clear()
            # detach 所有 runtime
            for tid in list(self.monitor.snapshot().keys()):
                self.monitor.detach(tid)
            self.state = SchedState.IDLE

    def switch_mode(self, mode: str) -> Dict[str, Any]:
        with self._lock:
            try:
                self.mode = SchedMode(mode)
            except ValueError:
                return {"status": "error", "reason": f"unknown_mode:{mode}"}
            logger.info("[Scheduler] mode -> %s", self.mode.value)
            return {"status": "ok", "mode": self.mode.value}

    # --- 對外：狀態查詢 ---------------------------------------------------
    def status(self) -> Dict[str, Any]:
        with self._lock:
            return {
                "scheduler_state": self.state.value,
                "mode": self.mode.value,
                "current_task_id": self.current_task_id,
                "queue": self.queue.snapshot(),
                "runtimes": self.monitor.snapshot(),
                "stats": self.context_store.read_stats(),
            }

    # --- 重規劃 / 重試 ----------------------------------------------------
    def request_replan(self, task_id: str, reason: str) -> None:
        with self._lock:
            attempts = self._replan_attempts.get(task_id, 0)
            if attempts >= self.max_replan_attempts:
                logger.warning(
                    "[Scheduler] replan exhausted for %s (%d), aborting",
                    task_id,
                    attempts,
                )
                self.abort(reason=f"replan_exhausted:{reason}")
                return
            self._replan_attempts[task_id] = attempts + 1
            logger.info(
                "[Scheduler] replan task=%s reason=%s attempt=%d",
                task_id,
                reason,
                attempts + 1,
            )
            ctx = self.context_store.get(task_id)
            text = (ctx.task_text if ctx else "Continue task") + f" (replan: {reason})"
            self.queue.clear()
            try:
                reasoning = self.global_agent.publish_global_task(
                    text, refresh=False, task_id=task_id
                )
                subtasks = (
                    reasoning.get("subtask_list", [])
                    if isinstance(reasoning, dict)
                    else []
                )
                if ctx:
                    self._ingest_subtasks(task_id, subtasks, ctx)
            except Exception as e:  # noqa: BLE001
                logger.exception("[Scheduler] replan failed: %s", e)
                self.diagnostics.raise_fault(
                    task_id=task_id, code="TM-001", details={"reason": reason}
                )

    def retry_current(self, task_id: str) -> None:
        rt = self.monitor.get(task_id)
        if not rt:
            return
        # 把當前 state 計數 reset，等下一次事件
        rt.force(rt.state, reason="retry")
        logger.info("[Scheduler] retry task=%s state=%s", task_id, rt.state.value)

    # --- Diagnostics → Scheduler 的回呼 -----------------------------------
    def diagnostics_callback(self, decision: DiagnoseDecision) -> None:
        action = decision.action
        tid = self.current_task_id or ""
        logger.info("[Scheduler] diagnostics decision: %s action=%s", decision.code, action)
        if action == "retry":
            self.retry_current(tid)
        elif action == "replan":
            if tid:
                self.request_replan(tid, decision.reason)
        elif action == "pause":
            self.pause()
        elif action == "abort":
            self.abort(decision.reason)
        # "none" 則不做事

    # --- Runtime 事件 hook ------------------------------------------------
    def _on_runtime_done(self, task_id: str, rt: TaskRuntime) -> None:
        with self._lock:
            logger.info("[Scheduler] runtime DONE: task=%s pkg=%s", task_id, rt.package_id)
            self.monitor.detach(task_id)
            # 移除這個 package 對應的所有 subtasks（PICK / REORIENT / PUSH）
            self._remove_subtasks_for_package(rt.package_id)
            if self.queue.active_count == 0 and self.queue.pending_count == 0:
                self.state = SchedState.DONE
                ctx = self.context_store.get(task_id)
                if ctx:
                    ctx.state = "DONE"
                    self.context_store.save(ctx)

    def _remove_subtasks_for_package(self, package_id: str) -> int:
        """從 queue 中移除所有屬於同一個 package 的 subtask。"""
        snap = self.queue.snapshot()
        removed = 0
        for sub in snap["active"] + snap["pending"]:
            if extract_package_id(sub) == package_id:
                if self.queue.remove(sub["subtask_id"]):
                    removed += 1
        return removed

    def _on_runtime_fault(self, task_id: str, rt: TaskRuntime, fault) -> None:
        logger.warning(
            "[Scheduler] runtime FAULT: task=%s pkg=%s code=%s",
            task_id,
            rt.package_id,
            getattr(fault, "code", "?"),
        )
        # 實際處理由 diagnostics_callback 走（已經在 raise_fault 時觸發了）

    # --- 內部：把規劃結果灌進隊列 + 創建 runtime --------------------------
    def _ingest_subtasks(
        self, task_id: str, subtasks: List[Dict[str, Any]], ctx: TaskContext
    ) -> None:
        seen_packages: set = set()
        for i, sub in enumerate(subtasks):
            if "subtask_id" not in sub:
                sub["subtask_id"] = f"t{i + 1}"
            self.queue.push(sub)
            pkg = extract_package_id(sub)
            if pkg and pkg not in seen_packages:
                self.monitor.attach(
                    task_id=task_id,
                    package_id=pkg,
                    subtask_id=sub["subtask_id"],
                )
                seen_packages.add(pkg)
        ctx.queue_active = self.queue.snapshot()["active"]
        ctx.queue_pending = self.queue.snapshot()["pending"]
        self.context_store.save(ctx)


__all__ = ["TaskScheduler", "SchedState", "SchedMode", "extract_package_id"]