"""
Terminal WebSocket Router - 终端 WebSocket 端点
"""

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from app.agents.tools.terminal_engine import pty_manager

router = APIRouter(prefix="/ws/terminal", tags=["terminal"])


class ConnectionManager:
    """WebSocket 连接管理器"""

    def __init__(self):
        self.active_connections: dict[str, WebSocket] = {}

    async def connect(self, session_id: str, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[session_id] = websocket

    def disconnect(self, session_id: str):
        if session_id in self.active_connections:
            del self.active_connections[session_id]

    async def send(self, session_id: str, data: str):
        if session_id in self.active_connections:
            await self.active_connections[session_id].send_text(data)

    def is_connected(self, session_id: str) -> bool:
        return session_id in self.active_connections


manager = ConnectionManager()


@router.websocket("/{session_id}")
async def terminal_websocket(websocket: WebSocket, session_id: str):
    """终端 WebSocket 端点"""
    await manager.connect(session_id, websocket)

    try:
        # 获取该 session 的输出队列
        queue = pty_manager._output_queues.get(session_id)
        if queue:
            # 异步任务：转发 PTY 输出到 WebSocket
            async def forward_output():
                while True:
                    try:
                        data = await asyncio.wait_for(queue.get(), timeout=0.1)
                        if data is None:
                            await manager.send(session_id, "[SESSION_END]")
                            break
                        await manager.send(session_id, data)
                    except asyncio.TimeoutError:
                        # 检查连接是否还活跃
                        if not manager.is_connected(session_id):
                            break
                    except Exception:
                        break

            import asyncio

            forward_task = asyncio.create_task(forward_output())

            # 主循环：接收用户输入
            try:
                while True:
                    data = await websocket.receive_text()
                    # 接收用户输入，写入 PTY
                    await pty_manager.write(session_id, data + "\n")
            except WebSocketDisconnect:
                pass
            finally:
                forward_task.cancel()
    except Exception:
        pass
    finally:
        manager.disconnect(session_id)