"""Distributed notification queue for background task events (SEP-1686).

Enables distributed Docket workers to send MCP notifications to clients
without holding session references. Workers push to a Redis queue,
the MCP server process subscribes and forwards to the client's session.

Pattern: Fire-and-forward with retry
- One queue per session_id
- LPUSH/BRPOP for reliable ordered delivery
- Retry up to 3 times on delivery failure, then discard
- TTL-based expiration for stale messages

Note: Docket's execution.subscribe() handles task state/progress events via
Redis Pub/Sub. This module handles elicitation-specific notifications that
require reliable delivery (input_required prompts, cancel signals).
"""

from __future__ import annotations

import asyncio
import json
import logging
import weakref
from contextlib import suppress
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any

import mcp.types

if TYPE_CHECKING:
    from docket import Docket
    from mcp.server.session import ServerSession

    from fastmcp.server.server import FastMCP

logger = logging.getLogger(__name__)

# Redis key patterns
NOTIFICATION_QUEUE_KEY = "fastmcp:notifications:{session_id}"
NOTIFICATION_ACTIVE_KEY = "fastmcp:notifications:{session_id}:active"

# Configuration
NOTIFICATION_TTL_SECONDS = 300  # 5 minute message TTL (elicitation response window)
MAX_DELIVERY_ATTEMPTS = 3  # Retry failed deliveries before discarding
SUBSCRIBER_TIMEOUT_SECONDS = 30  # BRPOP timeout (also heartbeat interval)


async def push_notification(
    session_id: str,
    notification: dict[str, Any],
    docket: Docket,
) -> None:
    """Push notification to session's queue (called from Docket worker).

    Used for elicitation-specific notifications (input_required, cancel)
    that need reliable delivery across distributed processes.

    Args:
        session_id: Target session's identifier
        notification: MCP notification dict (method, params, _meta)
        docket: Docket instance for Redis access
    """
    key = docket.key(NOTIFICATION_QUEUE_KEY.format(session_id=session_id))
    message = json.dumps(
        {
            "notification": notification,
            "attempt": 0,
            "enqueued_at": datetime.now(timezone.utc).isoformat(),
        }
    )
    async with docket.redis() as redis:
        await redis.lpush(key, message)
        await redis.expire(key, NOTIFICATION_TTL_SECONDS)


async def notification_subscriber_loop(
    session_id: str,
    session: ServerSession,
    docket: Docket,
    fastmcp: FastMCP,
) -> None:
    """Subscribe to notification queue and forward to session.

    Runs in the MCP server process. Bridges distributed workers to clients.

    This loop:
    1. Maintains a heartbeat (active subscriber marker for debugging)
    2. Blocks on BRPOP waiting for notifications
    3. Forwards notifications to the client's session
    4. Retries failed deliveries, then discards (no dead-letter queue)

    Args:
        session_id: Session identifier to subscribe to
        session: MCP ServerSession for sending notifications
        docket: Docket instance for Redis access
        fastmcp: FastMCP server instance (for elicitation relay)
    """
    queue_key = docket.key(NOTIFICATION_QUEUE_KEY.format(session_id=session_id))
    active_key = docket.key(NOTIFICATION_ACTIVE_KEY.format(session_id=session_id))

    logger.debug("Starting notification subscriber for session %s", session_id)

    while True:
        try:
            async with docket.redis() as redis:
                # Heartbeat: mark subscriber as active (for distributed debugging)
                await redis.set(active_key, "1", ex=SUBSCRIBER_TIMEOUT_SECONDS * 2)

                # Blocking wait for notification (timeout refreshes heartbeat)
                # Using BRPOP (right pop) for FIFO order with LPUSH (left push)
                result = await redis.brpop(
                    [queue_key], timeout=SUBSCRIBER_TIMEOUT_SECONDS
                )
                if not result:
                    continue  # Timeout - refresh heartbeat and retry

                _, message_bytes = result
                message = json.loads(message_bytes)
                notification_dict = message["notification"]
                attempt = message.get("attempt", 0)

                try:
                    # Reconstruct and send MCP notification
                    await _send_mcp_notification(
                        session, notification_dict, session_id, docket, fastmcp
                    )
                    logger.debug(
                        "Delivered notification to session %s (attempt %d)",
                        session_id,
                        attempt + 1,
                    )
                except Exception as send_error:
                    # Delivery failed - retry or discard
                    if attempt < MAX_DELIVERY_ATTEMPTS - 1:
                        # Re-queue with incremented attempt (back of queue)
                        message["attempt"] = attempt + 1
                        message["last_error"] = str(send_error)
                        await redis.lpush(queue_key, json.dumps(message))
                        logger.debug(
                            "Requeued notification for session %s (attempt %d): %s",
                            session_id,
                            attempt + 2,
                            send_error,
                        )
                    else:
                        # Discard after max attempts (session likely disconnected)
                        logger.warning(
                            "Discarding notification for session %s after %d attempts: %s",
                            session_id,
                            MAX_DELIVERY_ATTEMPTS,
                            send_error,
                        )

        except asyncio.CancelledError:
            # Graceful shutdown - leave pending messages in queue for reconnect
            logger.debug("Notification subscriber cancelled for session %s", session_id)
            break
        except Exception as e:
            logger.debug(
                "Notification subscriber error for session %s: %s", session_id, e
            )
            await asyncio.sleep(1)  # Backoff on error


async def _send_mcp_notification(
    session: ServerSession,
    notification_dict: dict[str, Any],
    session_id: str,
    docket: Docket,
    fastmcp: FastMCP,
) -> None:
    """Reconstruct MCP notification from dict and send to session.

    For input_required notifications with elicitation metadata, also sends
    a standard elicitation/create request to the client and relays the
    response back to the worker via Redis.

    Args:
        session: MCP ServerSession
        notification_dict: Notification as dict (method, params, _meta)
        session_id: Session identifier (for elicitation relay)
        docket: Docket instance (for notification delivery)
        fastmcp: FastMCP server instance (for elicitation relay)
    """
    method = notification_dict.get("method", "notifications/tasks/status")
    if method != "notifications/tasks/status":
        raise ValueError(f"Unsupported notification method for subscriber: {method}")

    notification = mcp.types.TaskStatusNotification.model_validate(
        {
            "method": "notifications/tasks/status",
            "params": notification_dict.get("params", {}),
            "_meta": notification_dict.get("_meta"),
        }
    )
    server_notification = mcp.types.ServerNotification(notification)

    await session.send_notification(server_notification)

    # If this is an input_required notification with elicitation metadata,
    # relay the elicitation to the client via standard elicitation/create
    params = notification_dict.get("params", {})
    if params.get("status") == "input_required":
        meta = notification_dict.get("_meta", {})
        related_task = meta.get("io.modelcontextprotocol/related-task", {})
        elicitation = related_task.get("elicitation")
        if elicitation:
            task_id = params.get("taskId")
            if not task_id:
                logger.warning(
                    "input_required notification missing taskId, skipping relay"
                )
                return
            if "task_scope" not in related_task:
                logger.warning(
                    "input_required notification for task %s missing task_scope "
                    "metadata, skipping elicitation relay",
                    task_id,
                )
                return
            task_scope = related_task["task_scope"]
            from fastmcp.server.tasks.elicitation import relay_elicitation

            task = asyncio.create_task(
                relay_elicitation(session, task_scope, task_id, elicitation, fastmcp),
                name=f"elicitation-relay-{task_id[:8]}",
            )
            _background_tasks.add(task)
            task.add_done_callback(_background_tasks.discard)


# =============================================================================
# Subscriber Management
# =============================================================================

# Strong references to fire-and-forget relay tasks (prevent GC mid-flight)
_background_tasks: set[asyncio.Task[None]] = set()

# Registry of active subscribers per session (prevents duplicates)
# Uses weakref to session to detect disconnects
_active_subscribers: dict[
    str, tuple[asyncio.Task[None], weakref.ref[ServerSession]]
] = {}


async def ensure_subscriber_running(
    session_id: str,
    session: ServerSession,
    docket: Docket,
    fastmcp: FastMCP,
) -> None:
    """Start notification subscriber if not already running (idempotent).

    Subscriber is created on first task submission and cleaned up on disconnect.
    Safe to call multiple times for the same session.

    Args:
        session_id: Session identifier
        session: MCP ServerSession
        docket: Docket instance
        fastmcp: FastMCP server instance (for elicitation relay)
    """
    # Check if subscriber already running for this session
    if session_id in _active_subscribers:
        task, session_ref = _active_subscribers[session_id]
        # Check if task is still running AND session is still alive
        if not task.done() and session_ref() is not None:
            return  # Already running

        # Task finished or session dead - clean up
        if not task.done():
            task.cancel()
            with suppress(asyncio.CancelledError):
                await task
        del _active_subscribers[session_id]

    # Start new subscriber task
    task = asyncio.create_task(
        notification_subscriber_loop(session_id, session, docket, fastmcp),
        name=f"notification-subscriber-{session_id[:8]}",
    )
    _active_subscribers[session_id] = (task, weakref.ref(session))
    logger.debug("Started notification subscriber for session %s", session_id)


async def stop_subscriber(session_id: str) -> None:
    """Stop notification subscriber for a session.

    Called when session disconnects. Pending messages remain in queue
    for delivery if client reconnects (with TTL expiration).

    Args:
        session_id: Session identifier
    """
    if session_id not in _active_subscribers:
        return

    task, _ = _active_subscribers.pop(session_id)
    if not task.done():
        task.cancel()
        with suppress(asyncio.CancelledError):
            await task
    logger.debug("Stopped notification subscriber for session %s", session_id)


def get_subscriber_count() -> int:
    """Get number of active subscribers (for monitoring)."""
    return len(_active_subscribers)
