#!/usr/bin/env python3
"""
Example: submit a job and wait for completion via WebSocket /ws instead of polling.

Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com

Usage:
  - Ensure the server has /ws in security.public_paths and a /ws endpoint.
  - Set HOST, PORT, TOKEN to match your server (or use defaults).
  - Run: python -m mcp_proxy_adapter.examples.websocket_examples.client_websocket_job_status

Requires: aiohttp (already a dependency of mcp-proxy-adapter).
"""

from __future__ import annotations

import asyncio
import os
import sys
from pathlib import Path

# Add project root for imports when run as script
_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


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")
    use_https = os.getenv("MCP_HTTPS", "").lower() in ("1", "true", "yes")

    client = JsonRpcClient(
        protocol="https" if use_https else "http",
        host=host,
        port=port,
        token_header="X-API-Key",
        token=token,
    )
    try:
        # 1. Submit a job (example: embed_queue with echo)
        print("Submitting job (embed_queue echo)...")
        r = await client.jsonrpc_call(
            "embed_queue",
            {"command": "echo", "params": {"message": "Hello via WS"}},
        )
        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. Wait for completion via WebSocket
        print("Waiting for completion via WebSocket /ws (timeout 60s)...")
        ev = await client.wait_for_job_via_websocket(job_id, timeout=60.0)
        print("Event:", ev.get("event"), "| Payload:", ev)
        return 0
    except asyncio.TimeoutError as e:
        print("Timeout:", e)
        return 1
    except Exception as e:
        print("Error:", e)
        return 1
    finally:
        await client.close()


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