TermControlAgent/app/api/v1/endpoints/websocket.py

"""
WebSocket API端点
"""
import asyncio
import json
from datetime import datetime
from typing import Dict, Any, Optional
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, HTTPException, status
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from app.core.websocket.manager import WebSocketManager
from app.core.websocket.client import WebSocketClient
from app.schemas.websocket import SuccessResponse as WebSocketResponse
from app.utils.structured_log import get_structured_logger
from app.utils.api_decorators import handle_api_errors

logger = get_structured_logger(__name__)

router = APIRouter()

# 创建全局websocket_manager实例
websocket_manager = WebSocketManager()

# 定义缺失的请求和响应模型
class SuccessResponse(BaseModel):
    """成功响应模型"""
    success: bool = Field(True, description="操作是否成功")
    message: str = Field(..., description="响应消息")
    data: Optional[Dict[str, Any]] = Field(None, description="响应数据")
    timestamp: str = Field(..., description="时间戳")

class CreateWebSocketClientRequest(BaseModel):
    """创建WebSocket客户端请求模型"""
    name: str = Field(..., description="客户端名称")
    url: str = Field(..., description="WebSocket URL")
    heartbeat_interval: Optional[int] = Field(30, description="心跳间隔(秒)")

def now_iso() -> str:
    return datetime.now().isoformat()

# 创建并连接客户端
@router.post("/websocket/clients", summary="创建并连接WebSocket客户端", response_model=SuccessResponse)
@handle_api_errors
async def create_and_connect_client(request: CreateWebSocketClientRequest):
    """创建客户端并立即连接"""
    try:
        # 创建客户端（自动创建3个Channel和适配器）
        client = await websocket_manager.create_client(request.name, request.url, request.heartbeat_interval)
        
        # 连接客户端
        success = await websocket_manager.connect_client(request.name)
        if not success:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"WebSocket客户端 {request.name} 连接失败"
            )
        
        logger.info(f"WebSocket客户端 {request.name} 创建并连接成功")
        
        return SuccessResponse(
            success=True,
            message=f"WebSocket客户端 {request.name} 创建并连接成功",
            data={
                "name": request.name, 
                "url": request.url, 
                "status": "connected", 
                "heartbeat_interval": request.heartbeat_interval,
                "channels": ["heartbeat", "send", "receive"]
            },
            timestamp=now_iso()
        )
        
    except Exception as e:
        logger.error(f"创建WebSocket客户端失败: {request.name} - {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"创建WebSocket客户端失败: {str(e)}"
        )

# 断开客户端
@router.post("/websocket/clients/{name}/disconnect", summary="断开WebSocket客户端", response_model=SuccessResponse)
@handle_api_errors
async def disconnect_client(name: str):
    """断开已存在客户端"""
    try:
        success = await websocket_manager.disconnect_client(name)
        if not success:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"WebSocket客户端 {name} 不存在或已断开"
            )
        
        logger.info(f"WebSocket客户端 {name} 断开成功")
        
        return SuccessResponse(
            success=True,
            message=f"WebSocket客户端 {name} 断开成功",
            timestamp=now_iso()
        )
        
    except Exception as e:
        logger.error(f"断开WebSocket客户端失败: {name} - {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"断开WebSocket客户端失败: {str(e)}"
        )

# 发送消息
@router.post("/websocket/clients/{name}/send", summary="发送消息到WebSocket客户端", response_model=SuccessResponse)
@handle_api_errors
async def send_message(name: str, message_type: str = "data", data: dict = None, priority: int = 0):
    """发送消息到指定客户端"""
    try:
        if data is None:
            data = {}
        
        success = await websocket_manager.send_message(name, message_type, data, priority)
        if not success:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"发送消息到WebSocket客户端 {name} 失败"
            )
        
        logger.info(f"消息发送成功: {name} -> {message_type}")
        
        return SuccessResponse(
            success=True,
            message=f"消息发送成功",
            data={"client": name, "type": message_type, "priority": priority},
            timestamp=now_iso()
        )
        
    except Exception as e:
        logger.error(f"发送消息失败: {name} -> {message_type} - {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"发送消息失败: {str(e)}"
        )

# 发送心跳
@router.post("/websocket/clients/{name}/heartbeat", summary="发送心跳消息", response_model=SuccessResponse)
@handle_api_errors
async def send_heartbeat(name: str):
    """发送心跳消息到指定客户端"""
    try:
        success = await websocket_manager.send_heartbeat(name)
        if not success:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"发送心跳到WebSocket客户端 {name} 失败"
            )
        
        logger.info(f"心跳发送成功: {name}")
        
        return SuccessResponse(
            success=True,
            message=f"心跳发送成功",
            data={"client": name},
            timestamp=now_iso()
        )
        
    except Exception as e:
        logger.error(f"发送心跳失败: {name} - {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"发送心跳失败: {str(e)}"
        )

# 获取统计信息
@router.get("/websocket/stats", summary="获取WebSocket统计信息")
@handle_api_errors
async def get_stats():
    """获取WebSocket管理器统计信息"""
    try:
        stats = websocket_manager.get_stats()
        return SuccessResponse(
            message="获取统计信息成功",
            data=stats,
            timestamp=now_iso()
        )
        
    except Exception as e:
        logger.error(f"获取统计信息失败: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"获取统计信息失败: {str(e)}"
        )