hd-crane-gs/backend/app/api/v1/module_business/crane/controller.py

# -*- coding: utf-8 -*-

import json
import asyncio
from starlette.websockets import WebSocketDisconnect
from fastapi import APIRouter, Depends, UploadFile, Body, Path, WebSocket
from fastapi.responses import StreamingResponse, JSONResponse
from app.common.response import SuccessResponse, StreamResponse
from app.core.dependencies import AuthPermission
from app.api.v1.module_system.auth.schema import AuthSchema
from app.core.base_params import PaginationQueryParam
from app.utils.common_util import bytes2file_response
from app.core.logger import log
from app.core.base_schema import BatchSetAvailable

from .service import BizCraneService
from .schema import BizCraneCreateSchema, BizCraneUpdateSchema, BizCraneQueryParam

BizCraneRouter = APIRouter(prefix='/crane', tags=["行车信息模块"]) 

@BizCraneRouter.get("/detail/{id}", summary="获取行车信息详情", description="获取行车信息详情")
async def get_crane_detail_controller(
    id: int = Path(..., description="ID"),
    auth: AuthSchema = Depends(AuthPermission(["module_business:crane:query"]))
) -> JSONResponse:
    """获取行车信息详情接口"""
    result_dict = await BizCraneService.detail_crane_service(auth=auth, id=id)
    log.info(f"获取行车信息详情成功 {id}")
    return SuccessResponse(data=result_dict, msg="获取行车信息详情成功")

@BizCraneRouter.get("/list", summary="查询行车信息列表", description="查询行车信息列表")
async def get_crane_list_controller(
    page: PaginationQueryParam = Depends(),
    search: BizCraneQueryParam = Depends(),
    auth: AuthSchema = Depends(AuthPermission(["module_business:crane:query"]))
) -> JSONResponse:
    """查询行车信息列表接口（数据库分页）"""
    result_dict = await BizCraneService.page_crane_service(
        auth=auth,
        page_no=page.page_no if page.page_no is not None else 1,
        page_size=page.page_size if page.page_size is not None else 10,
        search=search,
        order_by=page.order_by
    )
    log.info("查询行车信息列表成功")
    return SuccessResponse(data=result_dict, msg="查询行车信息列表成功")

@BizCraneRouter.post("/create", summary="创建行车信息", description="创建行车信息")
async def create_crane_controller(
    data: BizCraneCreateSchema,
    auth: AuthSchema = Depends(AuthPermission(["module_business:crane:create"]))
) -> JSONResponse:
    """创建行车信息接口"""
    result_dict = await BizCraneService.create_crane_service(auth=auth, data=data)
    log.info("创建行车信息成功")
    return SuccessResponse(data=result_dict, msg="创建行车信息成功")

@BizCraneRouter.put("/update/{id}", summary="修改行车信息", description="修改行车信息")
async def update_crane_controller(
    data: BizCraneUpdateSchema,
    id: int = Path(..., description="ID"),
    auth: AuthSchema = Depends(AuthPermission(["module_business:crane:update"]))
) -> JSONResponse:
    """修改行车信息接口"""
    result_dict = await BizCraneService.update_crane_service(auth=auth, id=id, data=data)
    log.info("修改行车信息成功")
    return SuccessResponse(data=result_dict, msg="修改行车信息成功")

@BizCraneRouter.delete("/delete", summary="删除行车信息", description="删除行车信息")
async def delete_crane_controller(
    ids: list[int] = Body(..., description="ID列表"),
    auth: AuthSchema = Depends(AuthPermission(["module_business:crane:delete"]))
) -> JSONResponse:
    """删除行车信息接口"""
    await BizCraneService.delete_crane_service(auth=auth, ids=ids)
    log.info(f"删除行车信息成功: {ids}")
    return SuccessResponse(msg="删除行车信息成功")

@BizCraneRouter.patch("/available/setting", summary="批量修改行车信息状态", description="批量修改行车信息状态")
async def batch_set_available_crane_controller(
    data: BatchSetAvailable,
    auth: AuthSchema = Depends(AuthPermission(["module_business:crane:patch"]))
) -> JSONResponse:
    """批量修改行车信息状态接口"""
    await BizCraneService.set_available_crane_service(auth=auth, data=data)
    log.info(f"批量修改行车信息状态成功: {data.ids}")
    return SuccessResponse(msg="批量修改行车信息状态成功")

@BizCraneRouter.post('/export', summary="导出行车信息", description="导出行车信息")
async def export_crane_list_controller(
    search: BizCraneQueryParam = Depends(),
    auth: AuthSchema = Depends(AuthPermission(["module_business:crane:export"]))
) -> StreamingResponse:
    """导出行车信息接口"""
    result_dict_list = await BizCraneService.list_crane_service(search=search, auth=auth)
    export_result = await BizCraneService.batch_export_crane_service(obj_list=result_dict_list)
    log.info('导出行车信息成功')

    return StreamResponse(
        data=bytes2file_response(export_result),
        media_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
        headers={
            'Content-Disposition': 'attachment; filename=biz_crane.xlsx'
        }
    )

@BizCraneRouter.post('/import', summary="导入行车信息", description="导入行车信息")
async def import_crane_list_controller(
    file: UploadFile,
    auth: AuthSchema = Depends(AuthPermission(["module_business:crane:import"]))
) -> JSONResponse:
    """导入行车信息接口"""
    batch_import_result = await BizCraneService.batch_import_crane_service(file=file, auth=auth, update_support=True)
    log.info("导入行车信息成功")
    
    return SuccessResponse(data=batch_import_result, msg="导入行车信息成功")

@BizCraneRouter.post('/download/template', summary="获取行车信息导入模板", description="获取行车信息导入模板", dependencies=[Depends(AuthPermission(["module_business:crane:download"]))])
async def export_crane_template_controller() -> StreamingResponse:
    """获取行车信息导入模板接口"""
    import_template_result = await BizCraneService.import_template_download_crane_service()
    log.info('获取行车信息导入模板成功')

    return StreamResponse(
        data=bytes2file_response(import_template_result),
        media_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
        headers={'Content-Disposition': 'attachment; filename=biz_crane_template.xlsx'}
    )

@BizCraneRouter.websocket("/ws/online", name="天车在线状态检测", dependencies=[])
async def websocket_chat_controller(
        websocket: WebSocket,
):
    await websocket.accept()
    is_closed = False
    try:
        log.info("前端WebSocket连接成功，开始推送测试数据")
        # 构造测试数据（完全匹配前端预期的格式）
        test_messages = [
            # 测试报警消息（gc/alert）
            {
                "topic": "gc/crane_status",
                "data": [
                    {"crane_no": "11111111", "is_online": True},
                    {"crane_no": "123456", "is_online": False}
                ]
            }
        ]
        # 循环推送测试数据（每隔3秒推一次，模拟实时更新）
        while True:
            if is_closed:
                break
            for msg in test_messages:
                # 将字典转为JSON字符串（前端能解析）
                await websocket.send_text(json.dumps(msg))
                log.info(f"已推送测试数据: {msg['topic']}")
                await asyncio.sleep(5)  # 每3秒推一条，可调整间隔
    except WebSocketDisconnect:
        is_closed = True
        log.info("前端主动断开WebSocket连接")
    except Exception as e:
        is_closed = True
        log.error(f"WebSocket异常: {str(e)}", exc_info=True)
    finally:
        if not is_closed:
            try:
                await websocket.close()
                log.info("关闭WebSocket连接")
            except RuntimeError:
                log.debug("WebSocket已关闭，无需重复操作")