#!/usr/bin/env python3
"""
Minimal test server for WebSocket client examples: embed_queue + /ws.

Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com

Serves:
- GET /health
- POST /api/jsonrpc: method "embed_queue" returns job_id; others return error.
- GET /ws: WebSocket; on subscribe sends job_completed after short delay.

Run: python -m mcp_proxy_adapter.examples.websocket_examples.run_ws_test_server [--port 8080]
Then run client_websocket_job_status and client_bidirectional_websocket.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import uuid
from aiohttp import web

JOBS: dict = {}  # job_id -> asyncio.Event to trigger completion


async def health(_request: web.Request) -> web.Response:
    return web.json_response({"status": "ok"})


async def jsonrpc(request: web.Request) -> web.Response:
    try:
        body = await request.json()
    except Exception:
        return web.json_response(
            {"jsonrpc": "2.0", "error": {"code": -32700, "message": "Parse error"}, "id": None},
            status=200,
        )
    method = body.get("method")
    params = body.get("params") or {}
    req_id = body.get("id", 1)
    if method == "embed_queue":
        job_id = str(uuid.uuid4())
        JOBS[job_id] = asyncio.Event()
        return web.json_response({
            "jsonrpc": "2.0",
            "result": {
                "success": True,
                "data": {"job_id": job_id, "command": params.get("command", "echo")},
            },
            "id": req_id,
        })
    return web.json_response({
        "jsonrpc": "2.0",
        "error": {"code": -32601, "message": f"Method not found: {method}"},
        "id": req_id,
    })


async def websocket_handler(request: web.Request) -> web.StreamResponse:
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    try:
        async for msg in ws:
            if msg.type == web.WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                except json.JSONDecodeError:
                    continue
                action = data.get("action")
                job_id = data.get("job_id")
                if action == "subscribe" and job_id:
                    ev = JOBS.get(job_id)
                    if ev:
                        await asyncio.sleep(0.3)
                        await ws.send_str(json.dumps({
                            "event": "job_completed",
                            "job_id": job_id,
                            "result": {"message": "done"},
                        }))
                        ev.set()
                    else:
                        await ws.send_str(json.dumps({
                            "event": "job_failed",
                            "job_id": job_id,
                            "error": "job not found",
                        }))
    finally:
        await ws.close()
    return ws


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8080)
    parser.add_argument("--host", default="127.0.0.1")
    args = parser.parse_args()
    app = web.Application()
    app.router.add_get("/health", health)
    app.router.add_post("/api/jsonrpc", jsonrpc)
    app.router.add_get("/ws", websocket_handler)
    print(f"WS test server: http://{args.host}:{args.port} (embed_queue + /ws)")
    web.run_app(app, host=args.host, port=args.port)


if __name__ == "__main__":
    main()
