"""VidaAgent — 24/7 持久化 Agent Daemon operator。

该 operator 位于 L4（sage-middleware），统一编排：

- L3 算法执行核心：``AsyncReActLoop``
- L4 记忆接入：``VidaMemoryBridge``
- L4 周期反思：``VidaReflectionEngine``

核心能力：

1. 事件循环 + 异步消息队列（后台 daemon 消费）
2. 状态持久化（pending/inflight）与崩溃恢复
3. 优雅关闭（停止接收、持久化、可选等待队列）
4. 与三层记忆读写联动
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
import uuid
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


@dataclass(slots=True)
class VidaMessage:
    """VidaAgent 队列消息。"""

    message_id: str
    user_id: str
    content: str
    metadata: dict[str, Any] = field(default_factory=dict)
    created_at: float = field(default_factory=time.time)


@dataclass(slots=True)
class VidaResult:
    """VidaAgent 处理结果。"""

    message_id: str
    answer: str = ""
    error: str = ""
    created_at: float = field(default_factory=time.time)

    @property
    def ok(self) -> bool:
        """结果是否成功。"""
        return not self.error


class VidaAgent:
    """24/7 持久化 Agent Daemon。

    Args:
        react_loop: ``AsyncReActLoop`` 实例。
        memory_bridge: ``VidaMemoryBridge`` 实例。
        reflection_engine: 可选 ``VidaReflectionEngine`` 实例。
        config: 配置项，支持：

            - ``state_file`` (str): 状态文件路径，默认 ``.vida_agent_state.json``。
            - ``working_top_k`` (int): 工作记忆检索条数，默认 10。
            - ``episodic_top_k`` (int): 情节记忆检索条数，默认 5。
            - ``semantic_top_k`` (int): 语义记忆检索条数，默认 5。
            - ``working_memory_capacity`` (int): 工作记忆容量阈值，默认 20。
            - ``drain_timeout_seconds`` (float): shutdown 等待队列清空超时，默认 5.0。
            - ``persist_on_each_step`` (bool): 每次状态变化后落盘，默认 True。
            - ``close_memory_on_shutdown`` (bool): shutdown 时关闭 memory bridge，默认 True。
    """

    _STATE_VERSION = 1

    def __init__(
        self,
        react_loop: Any,
        memory_bridge: Any,
        reflection_engine: Any | None = None,
        config: dict[str, Any] | None = None,
    ) -> None:
        self._react_loop = react_loop
        self._memory_bridge = memory_bridge
        self._reflection_engine = reflection_engine
        cfg = config or {}

        self._state_file = Path(cfg.get("state_file", ".vida_agent_state.json"))
        self._working_top_k = int(cfg.get("working_top_k", 10))
        self._episodic_top_k = int(cfg.get("episodic_top_k", 5))
        self._semantic_top_k = int(cfg.get("semantic_top_k", 5))
        self._working_memory_capacity = int(cfg.get("working_memory_capacity", 20))
        self._drain_timeout_seconds = float(cfg.get("drain_timeout_seconds", 5.0))
        self._persist_on_each_step = bool(cfg.get("persist_on_each_step", True))
        self._close_memory_on_shutdown = bool(cfg.get("close_memory_on_shutdown", True))

        self._queue: asyncio.Queue[str] = asyncio.Queue()
        self._pending: dict[str, VidaMessage] = {}
        self._inflight: VidaMessage | None = None
        self._results: dict[str, VidaResult] = {}
        self._waiters: dict[str, asyncio.Future[VidaResult]] = {}

        self._consumer_task: asyncio.Task[None] | None = None
        self._running = False
        self._accepting = False

        self._processed_count = 0
        self._failed_count = 0
        self._working_writes_since_reflect = 0

    async def start(self) -> None:
        """启动 daemon：恢复状态并开始消费队列。"""
        if self._running:
            return

        self._running = True
        self._accepting = True

        await self._load_state()
        self._consumer_task = asyncio.create_task(self._consumer_loop(), name="vida_agent_consumer")

        if self._reflection_engine is not None:
            await self._reflection_engine.start()

        logger.info("VidaAgent 已启动: pending=%d", self.pending_count)

    async def shutdown(self, *, drain: bool = True) -> None:
        """优雅关闭 daemon。

        Args:
            drain: 是否等待队列在超时内尽量清空。
        """
        if not self._running:
            return

        self._accepting = False

        if drain:
            await self._drain_queue(timeout=self._drain_timeout_seconds)

        self._running = False

        if self._consumer_task is not None and not self._consumer_task.done():
            self._consumer_task.cancel()
            try:
                await asyncio.gather(self._consumer_task, return_exceptions=True)
            finally:
                self._consumer_task = None

        if self._reflection_engine is not None:
            await self._reflection_engine.stop()

        await self._persist_state()
        await self._memory_bridge.persist_all()

        if self._close_memory_on_shutdown:
            await self._memory_bridge.close()

        logger.info(
            "VidaAgent 已关闭: processed=%d failed=%d", self._processed_count, self._failed_count
        )

    async def submit(
        self,
        user_id: str,
        content: str,
        metadata: dict[str, Any] | None = None,
    ) -> str:
        """提交消息到队列。

        Returns:
            message_id
        """
        if not self._accepting:
            raise RuntimeError("VidaAgent is not accepting new messages")

        message = VidaMessage(
            message_id=uuid.uuid4().hex,
            user_id=user_id,
            content=content,
            metadata=metadata or {},
        )

        self._pending[message.message_id] = message
        await self._queue.put(message.message_id)

        if self._persist_on_each_step:
            await self._persist_state()

        return message.message_id

    async def wait_result(self, message_id: str, timeout: float | None = None) -> VidaResult:
        """等待消息处理结果。"""
        if message_id in self._results:
            return self._results[message_id]

        fut = self._waiters.get(message_id)
        if fut is None:
            fut = asyncio.get_running_loop().create_future()
            self._waiters[message_id] = fut

        if timeout is None:
            return await fut
        return await asyncio.wait_for(fut, timeout=timeout)

    async def ask(
        self,
        user_id: str,
        content: str,
        metadata: dict[str, Any] | None = None,
        timeout: float | None = None,
    ) -> VidaResult:
        """提交并等待结果。"""
        message_id = await self.submit(user_id=user_id, content=content, metadata=metadata)
        return await self.wait_result(message_id=message_id, timeout=timeout)

    async def handle_event(self, event: Any) -> VidaResult:
        """统一事件处理入口（供 TriggerManager 调用）。

        支持入参：
        - VidaEvent（具备 ``event_type/source/payload`` 属性）
        - dict（包含同名键）

        事件 payload 中可选字段：
        - ``user_id``
        - ``content`` / ``message``
        - ``timeout``
        """
        if isinstance(event, dict):
            event_type = str(event.get("event_type", "unknown"))
            source = str(event.get("source", "unknown"))
            payload_raw = event.get("payload", {})
        else:
            event_type = str(getattr(event, "event_type", "unknown"))
            source = str(getattr(event, "source", "unknown"))
            payload_raw = getattr(event, "payload", {})

        payload = payload_raw if isinstance(payload_raw, dict) else {"value": payload_raw}
        user_id = str(payload.get("user_id", "system"))

        content: str
        if "content" in payload:
            content = str(payload["content"])
        elif "message" in payload:
            content = str(payload["message"])
        elif "body" in payload:
            body = payload["body"]
            content = body if isinstance(body, str) else json.dumps(body, ensure_ascii=False)
        else:
            content = json.dumps(payload, ensure_ascii=False)

        timeout_val = payload.get("timeout")
        timeout = float(timeout_val) if timeout_val is not None else None

        metadata = {
            "source": source,
            "event_type": event_type,
            "event_payload": payload,
        }
        return await self.ask(user_id=user_id, content=content, metadata=metadata, timeout=timeout)

    async def _consumer_loop(self) -> None:
        """后台消费循环。"""
        while self._running:
            try:
                message_id = await self._queue.get()
            except asyncio.CancelledError:
                break

            message = self._pending.get(message_id)
            if message is None:
                self._queue.task_done()
                continue

            self._inflight = message
            if self._persist_on_each_step:
                await self._persist_state()

            try:
                result = await self._handle_message(message)
                self._processed_count += 1
            except Exception as exc:  # noqa: BLE001
                self._failed_count += 1
                result = VidaResult(message_id=message.message_id, error=str(exc))
                logger.exception("VidaAgent 消息处理失败: message_id=%s", message.message_id)

            self._results[message.message_id] = result
            self._pending.pop(message.message_id, None)
            self._inflight = None

            waiter = self._waiters.pop(message.message_id, None)
            if waiter is not None and not waiter.done():
                waiter.set_result(result)

            self._queue.task_done()

            if self._persist_on_each_step:
                await self._persist_state()

    async def _handle_message(self, message: VidaMessage) -> VidaResult:
        """处理单条消息：读记忆 → ReAct 推理 → 写记忆。"""
        await self._memory_bridge.store_working(
            message.content,
            metadata={"role": "user", "user_id": message.user_id, **message.metadata},
        )
        self._working_writes_since_reflect += 1

        working = await self._memory_bridge.recall_working("", top_k=self._working_top_k)
        episodic = await self._memory_bridge.recall_episodic(
            message.content,
            top_k=self._episodic_top_k,
        )
        semantic = await self._memory_bridge.recall_semantic(
            message.content,
            top_k=self._semantic_top_k,
        )

        memory_context = {
            "working": working,
            "episodic": episodic,
            "semantic": semantic,
            "user_id": message.user_id,
        }

        answer = await self._react_loop.run_with_tools(
            query=message.content,
            tools=[],
            memory_context=memory_context,
        )

        await self._memory_bridge.store_working(
            answer,
            metadata={"role": "assistant", "user_id": message.user_id},
        )
        self._working_writes_since_reflect += 1

        await self._memory_bridge.store_episodic(
            f"Q: {message.content}\nA: {answer}",
            metadata={"user_id": message.user_id, "source": "vida_agent"},
        )

        if (
            self._reflection_engine is not None
            and self._working_writes_since_reflect >= self._working_memory_capacity
        ):
            self._working_writes_since_reflect = 0
            try:
                await self._reflection_engine.reflect()
            except Exception:  # noqa: BLE001
                logger.exception("VidaAgent 触发反思失败")

        return VidaResult(message_id=message.message_id, answer=answer)

    async def _drain_queue(self, timeout: float) -> None:
        """等待队列清空到超时。"""
        if self._queue.empty():
            return

        try:
            await asyncio.wait_for(self._queue.join(), timeout=timeout)
        except TimeoutError:
            logger.warning("VidaAgent 等待队列清空超时: timeout=%.2fs", timeout)

    async def _load_state(self) -> None:
        """从状态文件恢复 pending/inflight 消息。"""
        if not self._state_file.exists():
            return

        try:
            with self._state_file.open("r", encoding="utf-8") as f:
                raw = json.load(f)
        except Exception as exc:  # noqa: BLE001
            raise RuntimeError(f"Failed to load VidaAgent state: {exc}") from exc

        if int(raw.get("version", -1)) != self._STATE_VERSION:
            raise RuntimeError("Unsupported VidaAgent state version")

        self._processed_count = int(raw.get("processed_count", 0))
        self._failed_count = int(raw.get("failed_count", 0))
        self._working_writes_since_reflect = int(raw.get("working_writes_since_reflect", 0))

        pending = raw.get("pending", [])
        inflight = raw.get("inflight")

        recovered: list[VidaMessage] = []
        for item in pending:
            recovered.append(VidaMessage(**item))
        if inflight:
            recovered.append(VidaMessage(**inflight))

        for message in recovered:
            self._pending[message.message_id] = message
            await self._queue.put(message.message_id)

        logger.info("VidaAgent 状态恢复完成: recovered=%d", len(recovered))

    async def _persist_state(self) -> None:
        """持久化当前 daemon 状态到磁盘（原子写）。"""
        self._state_file.parent.mkdir(parents=True, exist_ok=True)

        payload = {
            "version": self._STATE_VERSION,
            "pending": [asdict(msg) for msg in self._pending.values()],
            "inflight": asdict(self._inflight) if self._inflight is not None else None,
            "processed_count": self._processed_count,
            "failed_count": self._failed_count,
            "working_writes_since_reflect": self._working_writes_since_reflect,
            "saved_at": time.time(),
        }

        tmp_file = self._state_file.with_suffix(self._state_file.suffix + ".tmp")
        with tmp_file.open("w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)
        tmp_file.replace(self._state_file)

    @property
    def is_running(self) -> bool:
        """daemon 是否运行中。"""
        return self._consumer_task is not None and not self._consumer_task.done()

    @property
    def is_accepting(self) -> bool:
        """是否接受新消息。"""
        return self._accepting

    @property
    def pending_count(self) -> int:
        """当前 pending 消息数。"""
        return len(self._pending)

    @property
    def processed_count(self) -> int:
        """成功处理消息数。"""
        return self._processed_count

    @property
    def failed_count(self) -> int:
        """处理失败消息数。"""
        return self._failed_count
