xml-pipeline/xml_pipeline/server/app.py

"""
aiohttp-based HTTP/WebSocket server.

Provides:
- REST API for authentication
- WebSocket for console/GUI message sending
- Integration with SystemPipeline for message injection
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import TYPE_CHECKING, Optional, Callable

try:
    from aiohttp import web, WSMsgType
    AIOHTTP_AVAILABLE = True
except ImportError:
    AIOHTTP_AVAILABLE = False
    web = None
    WSMsgType = None

from ..auth.users import get_user_store, UserStore
from ..auth.sessions import get_session_manager, SessionManager, Session

if TYPE_CHECKING:
    from ..message_bus.stream_pump import StreamPump
    from ..message_bus.system_pipeline import SystemPipeline

logger = logging.getLogger(__name__)


def auth_middleware():
    @web.middleware
    async def middleware(request, handler):
        if request.path in ("/auth/login", "/health"):
            return await handler(request)

        auth_header = request.headers.get("Authorization", "")
        if not auth_header.startswith("Bearer "):
            return web.json_response({"error": "Missing Authorization"}, status=401)

        token = auth_header[7:]
        session = request.app["session_manager"].validate(token)

        if not session:
            return web.json_response({"error": "Invalid token"}, status=401)

        request["session"] = session
        return await handler(request)

    return middleware


async def handle_login(request):
    try:
        data = await request.json()
    except:
        return web.json_response({"error": "Invalid JSON"}, status=400)

    username = data.get("username", "")
    password = data.get("password", "")

    if not username or not password:
        return web.json_response({"error": "Credentials required"}, status=400)

    user = request.app["user_store"].authenticate(username, password)
    if not user:
        return web.json_response({"error": "Invalid credentials"}, status=401)

    session = request.app["session_manager"].create(user.username, user.role)
    return web.json_response(session.to_dict())


async def handle_logout(request):
    session = request["session"]
    request.app["session_manager"].revoke(session.token)
    return web.json_response({"message": "Logged out"})


async def handle_me(request):
    session = request["session"]
    return web.json_response({
        "username": session.username,
        "role": session.role,
        "expires_at": session.expires_at.isoformat(),
    })


async def handle_health(request):
    return web.json_response({"status": "ok"})


async def handle_websocket(request):
    session = request["session"]
    pump = request.app.get("pump")
    system_pipeline = request.app.get("system_pipeline")

    ws = web.WebSocketResponse()
    await ws.prepare(request)

    # Track this WebSocket for response delivery
    ws_id = id(ws)
    request.app["websockets"][ws_id] = {
        "ws": ws,
        "user": session.username,
        "threads": set(),  # Thread IDs this client is subscribed to
    }

    await ws.send_json({"type": "connected", "username": session.username})

    try:
        async for msg in ws:
            if msg.type == WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                    resp = await handle_ws_msg(
                        data, session, pump, system_pipeline,
                        request.app["websockets"][ws_id]
                    )
                    await ws.send_json(resp)
                except Exception as e:
                    logger.exception(f"WebSocket error: {e}")
                    await ws.send_json({"type": "error", "error": str(e)})
    finally:
        # Cleanup on disconnect
        del request.app["websockets"][ws_id]

    return ws


async def handle_ws_msg(data, session, pump, system_pipeline, ws_state):
    """
    Handle WebSocket message.

    Message types:
        ping        - Keepalive
        status      - Get server status
        listeners   - List available listeners
        targets     - Alias for listeners
        send        - Send message to pipeline (@target or explicit)
    """
    t = data.get("type", "")

    if t == "ping":
        return {"type": "pong"}

    elif t == "status":
        from ..memory import get_context_buffer
        stats = get_context_buffer().get_stats()
        return {"type": "status", "threads": stats["thread_count"]}

    elif t == "listeners" or t == "targets":
        if not pump:
            return {"type": "listeners", "listeners": []}
        return {"type": "listeners", "listeners": list(pump.listeners.keys())}

    elif t == "send":
        # Send message to pipeline
        if not system_pipeline:
            return {"type": "error", "error": "Pipeline not available"}

        # Support two formats:
        # 1. {"type": "send", "raw": "@greeter Dan"}
        # 2. {"type": "send", "target": "greeter", "content": "Dan"}
        raw = data.get("raw")
        if raw:
            # Parse @target message format
            try:
                thread_id = await system_pipeline.inject_console(
                    raw=raw,
                    user=session.username,
                )
            except ValueError as e:
                return {"type": "error", "error": str(e)}
        else:
            target = data.get("target")
            content = data.get("content", data.get("text", data.get("message", "")))

            if not target:
                return {"type": "error", "error": "Missing target"}
            if not content:
                return {"type": "error", "error": "Missing content"}

            try:
                thread_id = await system_pipeline.inject_raw(
                    target=target,
                    content=content,
                    source="websocket",
                    user=session.username,
                )
            except ValueError as e:
                return {"type": "error", "error": str(e)}

        # Track thread for response delivery
        ws_state["threads"].add(thread_id)

        return {
            "type": "sent",
            "thread_id": thread_id,
            "target": data.get("target") or raw.split()[0].lstrip("@") if raw else None,
        }

    return {"type": "error", "error": f"Unknown message type: {t}"}


def create_app(pump=None, system_pipeline=None):
    """
    Create the aiohttp application.

    Args:
        pump: StreamPump instance (optional)
        system_pipeline: SystemPipeline instance (optional, created from pump if not provided)
    """
    if not AIOHTTP_AVAILABLE:
        raise RuntimeError("aiohttp not installed")

    app = web.Application(middlewares=[auth_middleware()])
    app["user_store"] = get_user_store()
    app["session_manager"] = get_session_manager()
    app["pump"] = pump
    app["websockets"] = {}  # Track connected WebSocket clients

    # Create SystemPipeline if pump provided but system_pipeline not
    if pump and not system_pipeline:
        from ..message_bus.system_pipeline import SystemPipeline
        system_pipeline = SystemPipeline(pump)

    app["system_pipeline"] = system_pipeline

    app.router.add_post("/auth/login", handle_login)
    app.router.add_post("/auth/logout", handle_logout)
    app.router.add_get("/auth/me", handle_me)
    app.router.add_get("/health", handle_health)
    app.router.add_get("/ws", handle_websocket)

    return app


async def run_server(pump=None, host="127.0.0.1", port=8765):
    """
    Run the server.

    Args:
        pump: StreamPump instance for message handling
        host: Bind address
        port: Port number
    """
    app = create_app(pump)
    runner = web.AppRunner(app)
    await runner.setup()

    site = web.TCPSite(runner, host, port)
    await site.start()

    print(f"Server on http://{host}:{port}")

    try:
        while True:
            await asyncio.sleep(3600)
    except asyncio.CancelledError:
        pass
    finally:
        await runner.cleanup()