context.py

# -*- coding: utf-8 -*-
"""TaskContext：任務級上下文記憶，可持久化到 Redis。

對應 interfaces.md §3 Redis Key 規範。
"""
from __future__ import annotations

import json
import threading
import time
import uuid
from dataclasses import asdict, dataclass, field
from typing import Any, Dict, List, Optional

from .events import FaultRecord


def new_task_id() -> str:
    return uuid.uuid4().hex


@dataclass
class TaskContext:
    task_id: str
    task_text: str = ""
    state: str = "IDLE"  # IDLE / RUNNING / PAUSED / FAULT / RESET / DONE
    created_at: float = field(default_factory=time.time)
    updated_at: float = field(default_factory=time.time)
    queue_active: List[Dict[str, Any]] = field(default_factory=list)
    queue_pending: List[Dict[str, Any]] = field(default_factory=list)
    faults: List[Dict[str, Any]] = field(default_factory=list)
    keyframes: List[Dict[str, Any]] = field(default_factory=list)
    stats: Dict[str, Any] = field(default_factory=dict)

    def touch(self) -> None:
        self.updated_at = time.time()

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    @classmethod
    def create(cls, task_text: str) -> "TaskContext":
        return cls(task_id=new_task_id(), task_text=task_text)


class ContextStore:
    """把 TaskContext 持久化到 Redis 的 wrapper。

    底層用 RoboOS 既有的 Collaborator.record_environment / read_environment，
    這樣不必再多接一個 Redis client。
    """

    META_TTL_SEC = 24 * 3600

    def __init__(self, collaborator: Any) -> None:
        self.collaborator = collaborator
        self._lock = threading.RLock()
        # 進程內快取：避免每次都打 Redis
        self._cache: Dict[str, TaskContext] = {}

    # --- 基本 CRUD --------------------------------------------------------
    def save(self, ctx: TaskContext) -> None:
        with self._lock:
            ctx.touch()
            self._cache[ctx.task_id] = ctx
            try:
                self.collaborator.record_environment(
                    f"task:{ctx.task_id}:meta", json.dumps(ctx.to_dict(), default=str)
                )
            except Exception as e:  # noqa: BLE001
                # Redis 暫時不可用不應阻塞主流程
                print(f"[ContextStore] save warning: {e}")

    def get(self, task_id: str) -> Optional[TaskContext]:
        with self._lock:
            if task_id in self._cache:
                return self._cache[task_id]
            try:
                raw = self.collaborator.read_environment(f"task:{task_id}:meta")
            except Exception as e:  # noqa: BLE001
                print(f"[ContextStore] read warning: {e}")
                return None
            if not raw:
                return None
            data = raw if isinstance(raw, dict) else json.loads(raw)
            ctx = TaskContext(**data)
            self._cache[task_id] = ctx
            return ctx

    def list_active(self) -> List[TaskContext]:
        with self._lock:
            return [c for c in self._cache.values() if c.state == "RUNNING"]

    # --- 場景快照 ----------------------------------------------------------
    def update_current_scene(self, snapshot: Dict[str, Any]) -> None:
        try:
            self.collaborator.record_environment(
                "task:current_scene", json.dumps(snapshot, default=str)
            )
        except Exception as e:  # noqa: BLE001
            print(f"[ContextStore] scene warning: {e}")

    def read_current_scene(self) -> Optional[Dict[str, Any]]:
        try:
            raw = self.collaborator.read_environment("task:current_scene")
        except Exception:  # noqa: BLE001
            return None
        if not raw:
            return None
        return raw if isinstance(raw, dict) else json.loads(raw)

    # --- 故障 / 關鍵幀 -----------------------------------------------------
    def append_fault(self, ctx: TaskContext, fault: FaultRecord) -> None:
        ctx.faults.append(fault.to_dict())
        self.save(ctx)

    def append_keyframe(self, ctx: TaskContext, keyframe: Dict[str, Any]) -> None:
        ctx.keyframes.append(keyframe)
        # 控制長度，避免上下文無限增長
        if len(ctx.keyframes) > 200:
            ctx.keyframes = ctx.keyframes[-200:]
        self.save(ctx)

    # --- 全局統計 ----------------------------------------------------------
    def bump_stat(self, key: str, delta: float = 1) -> None:
        try:
            raw = self.collaborator.read_environment("task:stats")
            stats: Dict[str, Any] = {}
            if raw:
                stats = raw if isinstance(raw, dict) else json.loads(raw)
            stats[key] = float(stats.get(key, 0)) + delta
            self.collaborator.record_environment(
                "task:stats", json.dumps(stats, default=str)
            )
        except Exception as e:  # noqa: BLE001
            print(f"[ContextStore] stat warning: {e}")

    def read_stats(self) -> Dict[str, Any]:
        try:
            raw = self.collaborator.read_environment("task:stats")
        except Exception:  # noqa: BLE001
            return {}
        if not raw:
            return {}
        return raw if isinstance(raw, dict) else json.loads(raw)


__all__ = ["TaskContext", "ContextStore", "new_task_id"]