"""
StreamOutput - 流式输出封装
"""

import json
from datetime import datetime, timezone
from typing import Any, AsyncGenerator, Callable

from dataclasses import dataclass


@dataclass
class StreamEvent:
    """流式事件"""

    type: str  # "output" | "error" | "status" | "complete"
    session_id: str
    data: str
    timestamp: str


class StreamOutput:
    """流式输出处理器"""

    def __init__(
        self,
        session_id: str,
        websocket_sender: Callable[[str, str], Any] | None = None,
    ):
        self.session_id = session_id
        self.websocket_sender = websocket_sender
        self._listeners: list[Callable[[StreamEvent], Any]] = []

    def add_listener(self, listener: Callable[[StreamEvent], Any]):
        """添加事件监听器"""
        self._listeners.append(listener)

    def remove_listener(self, listener: Callable[[StreamEvent], Any]):
        """移除事件监听器"""
        if listener in self._listeners:
            self._listeners.remove(listener)

    async def push(self, event_type: str, data: str):
        """推送事件到 WebSocket 和监听器"""
        event = StreamEvent(
            type=event_type,
            session_id=self.session_id,
            data=data,
            timestamp=datetime.now(timezone.utc).isoformat(),
        )

        # 发送到 WebSocket
        if self.websocket_sender:
            try:
                await self.websocket_sender(self.session_id, json.dumps(event.__dict__))
            except Exception:
                pass

        # 通知监听器
        for listener in self._listeners:
            try:
                result = listener(event)
                if hasattr(result, "__await__"):
                    await result
            except Exception:
                pass

    async def stream_execution(
        self,
        executor,
        prompt: str,
    ) -> AsyncGenerator[str, None]:
        """包装执行器，实现流式输出"""
        async for line in executor.execute(prompt):
            await self.push("output", line)
            yield line

        await self.push("complete", "")

    async def push_status(self, status: str):
        """推送状态消息"""
        await self.push("status", status)

    async def push_error(self, error: str):
        """推送错误消息"""
        await self.push("error", error)