leocrm/app/plugins/builtins/ai_ui_control/routes.py

"""REST and WebSocket routes for AI UI Control.

Task 4.2: WebSocket-Endpoint für KI-UI-Steuerung
Backend WebSocket /ws/ai-ui-control. Authenticated via session cookie.
AI agent sends commands via REST, frontend receives via WebSocket.
"""

from __future__ import annotations

import json
import logging
import uuid

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, HTTPException, Request
from fastapi.responses import JSONResponse

from app.plugins.builtins.ai_ui_control.schemas import (
    UICommand,
    UICommandCreate,
    UICommandFeedback,
    UICommandResponse,
    UICommandStatus,
    UICommandStatusResponse,
    UICommandType,
)

logger = logging.getLogger(__name__)

router = APIRouter(tags=["ai-ui-control"])


# ─── REST endpoints (for AI agents) ───

@router.post("/command", response_model=UICommandResponse)
async def send_ui_command(
    request: Request,
    body: UICommandCreate,
):
    """Send a UI command to the frontend.

    Called by AI agents to control the UI. The command is forwarded to
    the frontend via WebSocket. If the user is not online, returns pending status.

    Authentication: requires valid session (same-user commands only).
    """
    from app.config import get_settings
    from app.core.auth import get_session_data, get_redis
    from app.core.service_container import get_container

    settings = get_settings()
    session_id = request.cookies.get(settings.session_cookie_name)
    if not session_id:
        raise HTTPException(status_code=401, detail="Not authenticated")

    redis = get_redis()
    session_data = await get_session_data(redis, session_id)
    if session_data is None:
        raise HTTPException(status_code=401, detail="Session expired")

    user_id = session_data["user_id"]

    container = get_container()
    if not container.has("ai_ui_control_ws"):
        raise HTTPException(status_code=503, detail="AI UI Control not available")

    ws_manager = container.get("ai_ui_control_ws")

    command_id = str(uuid.uuid4())
    command_dict = {
        "command_id": command_id,
        "action": body.action.value,
        "path": body.path,
        "entity": body.entity,
        "filter": body.filter,
        "contact_id": body.contact_id,
        "modal": body.modal,
        "tab": body.tab,
        "section": body.section,
        "key": body.key,
        "value": body.value,
        "description": body.description,
    }

    delivered_id = await ws_manager.send_command(user_id, command_dict)

    if delivered_id:
        return UICommandResponse(
            command_id=delivered_id,
            status=UICommandStatus.delivered,
            action=body.action,
            message="Command delivered to frontend",
        )
    else:
        # User not online — store as pending
        ws_manager._feedback[command_id] = {
            "command_id": command_id,
            "status": "pending",
            "action": body.action.value,
            "message": "User not online, command pending",
        }
        return UICommandResponse(
            command_id=command_id,
            status=UICommandStatus.pending,
            action=body.action,
            message="User not online, command pending",
        )


@router.get("/command/{command_id}/status", response_model=UICommandStatusResponse)
async def get_command_status(
    request: Request,
    command_id: str,
):
    """Poll the status of a previously sent command.

    AI agents call this to check if the frontend has executed the command.
    """
    from app.config import get_settings
    from app.core.auth import get_session_data, get_redis
    from app.core.service_container import get_container

    settings = get_settings()
    session_id = request.cookies.get(settings.session_cookie_name)
    if not session_id:
        raise HTTPException(status_code=401, detail="Not authenticated")

    redis = get_redis()
    session_data = await get_session_data(redis, session_id)
    if session_data is None:
        raise HTTPException(status_code=401, detail="Session expired")

    container = get_container()
    if not container.has("ai_ui_control_ws"):
        raise HTTPException(status_code=503, detail="AI UI Control not available")

    ws_manager = container.get("ai_ui_control_ws")
    feedback = ws_manager.get_feedback(command_id)

    if feedback is None:
        raise HTTPException(status_code=404, detail="Command not found")

    status = UICommandStatus(feedback.get("status", "pending"))
    action = None
    if feedback.get("action"):
        try:
            action = UICommandType(feedback["action"])
        except ValueError:
            pass

    fb_model = None
    if status in (UICommandStatus.success, UICommandStatus.failed):
        fb_model = UICommandFeedback(
            command_id=feedback.get("command_id", command_id),
            status=status,
            action=action,
            current_path=feedback.get("current_path"),
            current_tab=feedback.get("current_tab"),
            message=feedback.get("message"),
            error=feedback.get("error"),
            data=feedback.get("data"),
        )

    return UICommandStatusResponse(
        command_id=command_id,
        status=status,
        action=action,
        feedback=fb_model,
    )


@router.get("/online-users")
async def get_online_users(request: Request):
    """Check which users are currently online (have active frontend WS connections)."""
    from app.config import get_settings
    from app.core.auth import get_session_data, get_redis
    from app.core.service_container import get_container

    settings = get_settings()
    session_id = request.cookies.get(settings.session_cookie_name)
    if not session_id:
        raise HTTPException(status_code=401, detail="Not authenticated")

    redis = get_redis()
    session_data = await get_session_data(redis, session_id)
    if session_data is None:
        raise HTTPException(status_code=401, detail="Session expired")

    container = get_container()
    if not container.has("ai_ui_control_ws"):
        raise HTTPException(status_code=503, detail="AI UI Control not available")

    ws_manager = container.get("ai_ui_control_ws")
    return {"online_users": ws_manager.get_online_users()}


# ─── WebSocket endpoint (for frontend) ───

@router.websocket("/ws")
async def ai_ui_control_ws(websocket: WebSocket):
    """WebSocket endpoint for AI UI control.

    Frontend connects here to receive UI commands from AI agents.
    Frontend sends feedback back through this WebSocket after executing commands.

    Authentication: via session cookie (same pattern as kommunikation plugin).
    """
    from app.config import get_settings
    from app.core.auth import get_session_data, get_redis
    from app.core.service_container import get_container

    settings = get_settings()
    session_id = websocket.cookies.get(settings.session_cookie_name)
    if not session_id:
        await websocket.close(code=4001, reason="Not authenticated")
        return

    redis = get_redis()
    session_data = await get_session_data(redis, session_id)
    if session_data is None:
        await websocket.close(code=4001, reason="Session expired")
        return

    user_id = session_data["user_id"]
    tenant_id = session_data["tenant_id"]

    container = get_container()
    if not container.has("ai_ui_control_ws"):
        await websocket.close(code=4003, reason="AI UI Control not available")
        return

    ws_manager = container.get("ai_ui_control_ws")
    await ws_manager.connect(websocket, user_id)

    try:
        while True:
            data = await websocket.receive_text()
            msg = json.loads(data)
            msg_type = msg.get("type")

            if msg_type == "ping":
                await websocket.send_text(json.dumps({"type": "pong"}))
            elif msg_type == "feedback":
                # Frontend sends feedback after executing a command
                ws_manager.store_feedback(msg)
            elif msg_type == "status":
                # Frontend requests current state info
                await websocket.send_text(json.dumps({
                    "type": "status",
                    "online": True,
                    "user_id": user_id,
                }))
    except WebSocketDisconnect:
        await ws_manager.disconnect(websocket, user_id)
    except Exception:
        logger.exception("AI UI Control WebSocket error")
        await ws_manager.disconnect(websocket, user_id)