"""
WebSocket endpoint /ws for job push: subscribe/unsubscribe by job_id.

Uses app.state.job_push_registry (SubscriberRegistry). Each connection is wrapped
in a Subscriber; on disconnect unsubscribe_all_for is called.
See core/job_push/events.py for message format; core/job_push/subscriptions.py for registry.

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

from __future__ import annotations

import json

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

from mcp_proxy_adapter.core.logging import get_global_logger
from mcp_proxy_adapter.core.job_push.events import (
    ACTION_SUBSCRIBE,
    ACTION_UNSUBSCRIBE,
    ACTION_SUBSCRIBE_ERROR,
)
from mcp_proxy_adapter.core.job_push.subscriptions import Subscriber
from mcp_proxy_adapter.api.ws_identity import (
    get_identity_from_scope,
    is_allowed_to_subscribe,
)


class WsSubscriber(Subscriber):
    """Subscriber that sends payloads to a WebSocket connection."""

    def __init__(self, websocket: WebSocket) -> None:
        self._ws = websocket

    async def send(self, payload: dict) -> None:
        """Send payload as JSON to the WebSocket. May raise on closed connection."""
        await self._ws.send_json(payload)


async def handle_websocket_job_push(websocket: WebSocket, app: FastAPI) -> None:
    """
    Handle WebSocket connection for job push: subscribe/unsubscribe by job_id.

    Expects JSON messages: {"action": "subscribe", "job_id": "..."} and
    {"action": "unsubscribe", "job_id": "..."}. On disconnect calls
    registry.unsubscribe_all_for(subscriber).
    """
    await websocket.accept()
    logger = get_global_logger()
    registry = getattr(app.state, "job_push_registry", None)
    if registry is None:
        logger.warning("WebSocket /ws: job_push_registry not set, closing")
        await websocket.close(code=1011, reason="Job push not enabled")
        return
    subscriber = WsSubscriber(websocket)
    scope = websocket.scope
    identity = get_identity_from_scope(scope)
    try:
        while True:
            raw = await websocket.receive_text()
            try:
                data = json.loads(raw) if isinstance(raw, str) else raw
            except json.JSONDecodeError:
                continue
            if not isinstance(data, dict):
                continue
            action = data.get("action")
            job_id = data.get("job_id")
            if action == ACTION_SUBSCRIBE and job_id:
                job_id_str = str(job_id).strip()
                if not job_id_str:
                    continue
                if not is_allowed_to_subscribe(identity, job_id_str):
                    await websocket.send_json(
                        {
                            "action": ACTION_SUBSCRIBE_ERROR,
                            "job_id": job_id_str,
                            "error": "Not allowed to subscribe to this job",
                        }
                    )
                    continue
                await registry.subscribe(job_id_str, subscriber)
            elif action == ACTION_UNSUBSCRIBE and job_id:
                job_id_str = str(job_id).strip()
                if job_id_str:
                    await registry.unsubscribe(job_id_str, subscriber)
    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.warning("WebSocket /ws error: %s", e)
    finally:
        await registry.unsubscribe_all_for(subscriber)
