"""
WebSocket endpoint /ws for job push: typed subscribe/unsubscribe (queue_job, ws_delivery, transfer_session).

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

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

from __future__ import annotations

import json

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

from mcp_proxy_adapter.api.ws_identity import (
    get_identity_from_scope,
    is_allowed_to_subscribe,
)
from mcp_proxy_adapter.core.job_push.events import (
    ACTION_SUBSCRIBE,
    ACTION_SUBSCRIBE_ERROR,
    ACTION_UNSUBSCRIBE,
    SUBSCRIPTION_TYPE_QUEUE_JOB,
    SUBSCRIPTION_TYPE_TRANSFER_SESSION,
    SUBSCRIPTION_TYPE_WS_DELIVERY,
)
from mcp_proxy_adapter.core.job_push.subscriptions import (
    Subscriber,
    SubscriptionCategory,
)
from mcp_proxy_adapter.core.logging import get_global_logger
from mcp_proxy_adapter.transfer.errors import (
    TransferExpiredError,
    TransferSessionNotFoundError,
)


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

    def __init__(self, websocket: WebSocket) -> None:
        """
        Initialize WebSocket subscriber.

        Args:
            websocket: WebSocket connection to send payloads to.
        """
        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 for job push: typed subscribe/unsubscribe.

    Requires subscription_type: queue_job (job_id), ws_delivery (ws_id), or
    transfer_session (transfer_id; session must exist).
    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)
    manager = getattr(app.state, "ws_delivery_manager", None)
    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")
            sub_type = data.get("subscription_type")
            if action == ACTION_SUBSCRIBE:
                if not isinstance(sub_type, str) or not sub_type.strip():
                    await websocket.send_json(
                        {
                            "action": ACTION_SUBSCRIBE_ERROR,
                            "error": "subscription_type required",
                        }
                    )
                    continue
                sub_type = sub_type.strip()
                if sub_type == SUBSCRIPTION_TYPE_QUEUE_JOB:
                    job_id = data.get("job_id")
                    if job_id is None or not str(job_id).strip():
                        await websocket.send_json(
                            {
                                "action": ACTION_SUBSCRIBE_ERROR,
                                "error": "job_id required for queue_job",
                            }
                        )
                        continue
                    job_id_str = str(job_id).strip()
                    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, SubscriptionCategory.QUEUE_JOB, subscriber
                    )
                elif sub_type == SUBSCRIPTION_TYPE_WS_DELIVERY:
                    ws_id = data.get("ws_id")
                    if ws_id is None or not str(ws_id).strip():
                        await websocket.send_json(
                            {
                                "action": ACTION_SUBSCRIBE_ERROR,
                                "error": "ws_id required for ws_delivery",
                            }
                        )
                        continue
                    ws_id_str = str(ws_id).strip()
                    if not is_allowed_to_subscribe(identity, ws_id_str):
                        await websocket.send_json(
                            {
                                "action": ACTION_SUBSCRIBE_ERROR,
                                "ws_id": ws_id_str,
                                "error": "Not allowed to subscribe to this job",
                            }
                        )
                        continue
                    await registry.subscribe(
                        ws_id_str, SubscriptionCategory.WS_DELIVERY, subscriber
                    )
                    if manager is not None:
                        await manager.flush_replay_to_subscriber(ws_id_str, subscriber)
                    else:
                        logger.warning(
                            "WebSocket /ws: ws_delivery_manager missing, "
                            "cannot flush replay"
                        )
                elif sub_type == SUBSCRIPTION_TYPE_TRANSFER_SESSION:
                    raw_tid = data.get("transfer_id")
                    if raw_tid is None or not str(raw_tid).strip():
                        await websocket.send_json(
                            {
                                "action": ACTION_SUBSCRIBE_ERROR,
                                "error": "transfer_id required for transfer_session",
                            }
                        )
                        continue
                    tid_str = str(raw_tid).strip()
                    if not is_allowed_to_subscribe(identity, tid_str):
                        await websocket.send_json(
                            {
                                "action": ACTION_SUBSCRIBE_ERROR,
                                "transfer_id": tid_str,
                                "error": "Not allowed to subscribe to this transfer",
                            }
                        )
                        continue
                    try:
                        from mcp_proxy_adapter.api.handlers import (
                            get_transfer_store,
                        )

                        store = get_transfer_store()
                        store.get_session(tid_str)
                    except TransferSessionNotFoundError:
                        await websocket.send_json(
                            {
                                "action": ACTION_SUBSCRIBE_ERROR,
                                "transfer_id": tid_str,
                                "error": "transfer_id not found",
                            }
                        )
                        continue
                    except TransferExpiredError:
                        await websocket.send_json(
                            {
                                "action": ACTION_SUBSCRIBE_ERROR,
                                "transfer_id": tid_str,
                                "error": "transfer session expired",
                            }
                        )
                        continue
                    await registry.subscribe(
                        tid_str, SubscriptionCategory.TRANSFER_SESSION, subscriber
                    )
                else:
                    await websocket.send_json(
                        {
                            "action": ACTION_SUBSCRIBE_ERROR,
                            "error": "unknown subscription_type",
                        }
                    )
            elif action == ACTION_UNSUBSCRIBE:
                if not isinstance(sub_type, str) or not sub_type.strip():
                    await websocket.send_json(
                        {
                            "action": ACTION_SUBSCRIBE_ERROR,
                            "error": "subscription_type required",
                        }
                    )
                    continue
                sub_type = sub_type.strip()
                if sub_type == SUBSCRIPTION_TYPE_QUEUE_JOB:
                    job_id = data.get("job_id")
                    if job_id is not None and str(job_id).strip():
                        await registry.unsubscribe(
                            str(job_id).strip(),
                            SubscriptionCategory.QUEUE_JOB,
                            subscriber,
                        )
                elif sub_type == SUBSCRIPTION_TYPE_WS_DELIVERY:
                    ws_id = data.get("ws_id")
                    if ws_id is not None and str(ws_id).strip():
                        await registry.unsubscribe(
                            str(ws_id).strip(),
                            SubscriptionCategory.WS_DELIVERY,
                            subscriber,
                        )
                elif sub_type == SUBSCRIPTION_TYPE_TRANSFER_SESSION:
                    tid = data.get("transfer_id")
                    if tid is not None and str(tid).strip():
                        await registry.unsubscribe(
                            str(tid).strip(),
                            SubscriptionCategory.TRANSFER_SESSION,
                            subscriber,
                        )
    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.warning("WebSocket /ws error: %s", e)
    finally:
        await registry.unsubscribe_all_for(subscriber)
