#!/usr/bin/env python3
"""
Example: bidirectional WebSocket channel — send (subscribe/unsubscribe) and receive (events).

Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com

Shows both directions on the same /ws connection:
- Client -> server: subscribe, optional unsubscribe
- Server -> client: iterate over events (job_completed, job_failed, progress, etc.)

Usage:
  - Ensure the server has /ws in security.public_paths and a /ws endpoint.
  - Set MCP_HOST, MCP_PORT, MCP_TOKEN to match your server.
  - Run: python -m mcp_proxy_adapter.examples.websocket_examples.client_bidirectional_websocket
"""

from __future__ import annotations

import asyncio
import os
import sys
from pathlib import Path

_root = Path(__file__).resolve().parent
while _root != _root.parent and not (_root / "pyproject.toml").exists():
    _root = _root.parent
if (_root / "pyproject.toml").exists():
    sys.path.insert(0, str(_root))

from mcp_proxy_adapter.client.jsonrpc_client import (
    JsonRpcClient,
    open_bidirectional_ws_channel,
)

# Terminal events: stop the receive loop
TERMINAL_EVENTS = frozenset({"job_completed", "job_failed", "job_stopped"})


async def main() -> int:
    host = os.getenv("MCP_HOST", "127.0.0.1")
    port = int(os.getenv("MCP_PORT", "8080"))
    token = os.getenv("MCP_TOKEN", "admin-secret-key")

    client = JsonRpcClient(
        protocol="http",
        host=host,
        port=port,
        token_header="X-API-Key",
        token=token,
    )
    try:
        # 1. Submit a job to get a job_id
        print("Submitting job (embed_queue echo)...")
        r = await client.jsonrpc_call(
            "embed_queue",
            {"command": "echo", "params": {"message": "Bidirectional WS demo"}},
        )
        out = client._extract_result(r)
        job_id = out.get("data", {}).get("job_id")
        if not job_id:
            print("No job_id in response:", out)
            return 1
        print("Job submitted:", job_id)

        # 2. Bidirectional channel: send subscribe, then receive until terminal event
        print("Opening bidirectional WebSocket channel to /ws ...")
        channel = open_bidirectional_ws_channel(client, receive_timeout=60.0)
        async with channel:
            # Direction 1: client -> server
            await channel.send_json({"action": "subscribe", "job_id": job_id})
            print("Sent: subscribe", job_id)

            # Direction 2: server -> client
            async for msg in channel.receive_iter():
                ev = msg.get("event") or msg.get("type")
                print("Received:", ev, "|", msg)
                if ev in TERMINAL_EVENTS:
                    break

            # Direction 1 again: optional unsubscribe
            await channel.send_json({"action": "unsubscribe", "job_id": job_id})
            print("Sent: unsubscribe", job_id)

        print("Channel closed.")
        return 0
    except Exception as e:
        print("Error:", e)
        return 1
    finally:
        await client.close()


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))
