#!/usr/bin/env python3
"""
NAME
    ws_example_runner – run a single /ws WebSocket test for MCP Proxy Adapter

SYNOPSIS
    asyncio.run(run_ws_example(protocol, host, port, token=..., cert=..., key=..., ca=...))
    exit_code = run_ws_example_sync(...)

DESCRIPTION
    Connects to the server's /ws endpoint (bidirectional channel), sends
    subscribe and unsubscribe messages for a test job_id, and closes.
    Used by mode-specific scripts (http_basic, https_token, mtls_roles, etc.)
    to verify that /ws works regardless of protocol (http/https/mtls) and
    security (no auth, token, token+roles).

    The /ws endpoint is protocol- and security-agnostic: it is registered
    when job_push.enabled or transport.ws_enabled is true in config, and
    is listed in public paths so the WebSocket upgrade is not blocked by
    token or mTLS middleware. Clients may optionally send the same token
    or use the same TLS client cert for identity (e.g. for subscribe auth).

PROTOCOL
    protocol must be one of: "http", "https", "mtls".
    - http:  ws://host:port/ws, no TLS.
    - https: wss://host:port/ws, server TLS; no client cert unless cert/key given.
    - mtls:  wss://host:port/ws, server TLS and client certificate.

SECURITY
    - No auth: pass token=None. /ws is public; upgrade succeeds without header.
    - Token: pass token="<api-key>" and token_header="X-API-Key" (or use default).
      Header is sent on the WebSocket upgrade request; server may use it for
      identity in subscribe checks.
    - Roles: same as token; roles are applied by server if configured.
    - mTLS: pass cert= and key= (and optionally ca=). Client certificate is
      sent during TLS handshake; server may use it for identity.

OPTIONS (arguments)
    protocol   "http" | "https" | "mtls"
    host       Server host (e.g. "localhost").
    port       Server port (e.g. 8080 for http, 8443 for https/mtls).
    token      Optional API key for X-API-Key (or Authorization: Bearer).
    token_header  Header name for token (default "X-API-Key").
    cert       Path to client certificate (for mtls or optional https).
    key        Path to client private key (for mtls or optional https).
    ca         Path to CA certificate (optional; if omitted, TLS verify disabled for examples).
    check_hostname  Whether to verify server hostname in TLS (default False for examples).

RETURN VALUE
    run_ws_example_sync returns 0 on success (connect, subscribe, unsubscribe, close).
    Returns 1 on connection failure, invalid response, or exception.

EXAMPLES
    # HTTP basic (no auth)
    run_ws_example_sync("http", "localhost", 8080)

    # HTTPS with token
    run_ws_example_sync("https", "localhost", 8443, token="admin-secret-key-https")

    # mTLS with client cert
    run_ws_example_sync("mtls", "localhost", 8443, token="admin-secret-key-mtls",
                        cert="mtls_certificates/client/test-client.crt",
                        key="mtls_certificates/client/test-client.key",
                        ca="mtls_certificates/ca/ca.crt")

SEE ALSO
    http_basic.py, http_token.py, http_token_roles.py, https_basic.py,
    https_token.py, https_token_roles.py, mtls_basic.py, mtls_roles.py,
    mcp_proxy_adapter.client.jsonrpc_client.ws_job_status,
    docs/.../bidirectional_job_push/

Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com
"""

from __future__ import annotations

import asyncio
import json
import sys
from pathlib import Path
from typing import Optional

# Ensure package root on path when run as script
_SCRIPT_DIR = Path(__file__).resolve().parent
_PKG_ROOT = _SCRIPT_DIR.parent.parent.parent
if str(_PKG_ROOT) not in sys.path:
    sys.path.insert(0, str(_PKG_ROOT))

from mcp_proxy_adapter.client.jsonrpc_client.transport import JsonRpcTransport
from mcp_proxy_adapter.client.jsonrpc_client.ws_job_status import (
    open_bidirectional_ws_channel,
)


async def run_ws_example(
    protocol: str,
    host: str = "localhost",
    port: int = 8080,
    token: Optional[str] = None,
    token_header: str = "X-API-Key",
    cert: Optional[str] = None,
    key: Optional[str] = None,
    ca: Optional[str] = None,
    check_hostname: bool = False,
) -> int:
    """
    Connect to /ws, send subscribe then unsubscribe for a test job_id, then close.

    Returns 0 on success, 1 on failure.
    """
    print(
        f"WS example: {protocol} {host}:{port} "
        f"(token={'yes' if token else 'no'}, cert={'yes' if cert else 'no'})"
    )
    transport = JsonRpcTransport(
        protocol=protocol,
        host=host,
        port=port,
        token_header=token_header if token else None,
        token=token,
        cert=cert,
        key=key,
        ca=ca,
        check_hostname=check_hostname,
    )
    try:
        channel = open_bidirectional_ws_channel(transport)
        async with channel:
            test_job_id = "ws-example-test-job"
            await channel.send_json({"action": "subscribe", "job_id": test_job_id})
            # Optionally wait for one message (server may not push if job doesn't exist)
            recv = channel.receive_iter()
            try:
                msg = await asyncio.wait_for(recv.__anext__(), 2.0)
                if isinstance(msg, dict):
                    print(f"  received: {json.dumps(msg)[:200]}")
            except asyncio.TimeoutError:
                pass
            await channel.send_json({"action": "unsubscribe", "job_id": test_job_id})
        print("  /ws connect, subscribe, unsubscribe, close: OK")
        return 0
    except Exception as e:
        print(f"  /ws failed: {e}")
        return 1


def run_ws_example_sync(
    protocol: str,
    host: str = "localhost",
    port: int = 8080,
    token: Optional[str] = None,
    token_header: str = "X-API-Key",
    cert: Optional[str] = None,
    key: Optional[str] = None,
    ca: Optional[str] = None,
    check_hostname: bool = False,
) -> int:
    """Synchronous wrapper: run run_ws_example via asyncio.run(). Returns 0 or 1."""
    return asyncio.run(
        run_ws_example(
            protocol=protocol,
            host=host,
            port=port,
            token=token,
            token_header=token_header,
            cert=cert,
            key=key,
            ca=ca,
            check_hostname=check_hostname,
        )
    )


if __name__ == "__main__":
    # Default: HTTP basic on 8080
    import argparse

    p = argparse.ArgumentParser(
        description="Run /ws example (default: http localhost 8080)"
    )
    p.add_argument("--protocol", default="http", choices=("http", "https", "mtls"))
    p.add_argument("--host", default="localhost")
    p.add_argument("--port", type=int, default=8080)
    p.add_argument("--token", default=None)
    p.add_argument("--token-header", default="X-API-Key")
    p.add_argument("--cert", default=None)
    p.add_argument("--key", default=None)
    p.add_argument("--ca", default=None)
    args = p.parse_args()
    code = run_ws_example_sync(
        protocol=args.protocol,
        host=args.host,
        port=args.port,
        token=args.token,
        token_header=args.token_header,
        cert=args.cert,
        key=args.key,
        ca=args.ca,
    )
    sys.exit(code)
