"""Redis pub/sub helpers for real-time WebSocket event streaming.

Each project gets its own Redis channel: project:{project_id}:events
Workers publish structured JSON events; the WebSocket handler subscribes and
forwards them to real-time clients.

The subscriber uses short ``get_message`` timeouts instead of ``listen()`` so that
when the API process shuts down (SIGINT / reload), asyncio cancellation can unwind
promptly instead of blocking until the next published event.

Usage in workers (prefer ``publish_project_event_dual`` from ``run_events``; agent lines emit as
``run.progress`` only):
    redis = await get_pubsub_redis()
    try:
        await publish_event(redis, str(project_id), "run.progress", {
            "run_id": "...", "text": "...", "type": "output", "stream": "stdout", "stage": "generating"
        })
    finally:
        await redis.aclose()

Usage in WebSocket handler:
    redis = await get_pubsub_redis()
    async for event in subscribe_project_events(redis, str(project_id)):
        await websocket.send_json(event)
"""
from __future__ import annotations

import json
from datetime import datetime, timezone
from typing import AsyncIterator

from redis import asyncio as aioredis

import structlog

from src.config import get_settings

logger = structlog.get_logger()

_CHANNEL_PREFIX = "project"
_DEVICE_CHANNEL_PREFIX = "device"
# ``listen()`` blocks on the socket with no wakeup until traffic; short polls allow
# graceful shutdown when the server cancels the WebSocket handler task.
_PUBLISH_POLL_TIMEOUT_S = 1.0


async def get_pubsub_redis() -> aioredis.Redis:
    """Create a dedicated Redis connection for pub/sub.

    Pub/sub requires its own connection (not from the shared pool used by arq
    or the FastAPI dependency) because subscribing blocks the connection for
    receiving messages only.
    """
    settings = get_settings()
    return aioredis.from_url(
        settings.redis_url,
        encoding="utf-8",
        decode_responses=True,
    )


async def publish_event(
    redis: aioredis.Redis,
    project_id: str,
    event: str,
    data: dict,
) -> None:
    """Publish a structured event to a project's Redis pub/sub channel."""
    msg = {
        "event": event,
        "project_id": project_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "data": data,
    }
    channel = f"{_CHANNEL_PREFIX}:{project_id}:events"
    try:
        await redis.publish(channel, json.dumps(msg))
    except Exception as exc:
        # Non-fatal: log and continue. WS events are best-effort.
        logger.warning("pubsub_publish_failed", channel=channel, event=event, error=str(exc))


async def subscribe_project_events(
    redis: aioredis.Redis,
    project_id: str,
) -> AsyncIterator[dict]:
    """Subscribe to a project's event channel and yield parsed event dicts.

    Caller is responsible for closing the Redis connection after the iterator
    is exhausted or after catching WebSocketDisconnect.
    """
    channel = f"{_CHANNEL_PREFIX}:{project_id}:events"
    pubsub = redis.pubsub()
    await pubsub.subscribe(channel)
    logger.info("pubsub_subscribed", channel=channel)
    try:
        while True:
            raw_msg = await pubsub.get_message(
                ignore_subscribe_messages=True,
                timeout=_PUBLISH_POLL_TIMEOUT_S,
            )
            if raw_msg is None:
                continue
            if raw_msg.get("type") == "message":
                try:
                    yield json.loads(raw_msg["data"])
                except (json.JSONDecodeError, KeyError, TypeError):
                    continue
    finally:
        try:
            await pubsub.unsubscribe(channel)
            await pubsub.aclose()
        except Exception:
            pass
        logger.info("pubsub_unsubscribed", channel=channel)


async def publish_device_event(
    redis: aioredis.Redis,
    device_id: str,
    event: str,
    data: dict,
) -> None:
    """Publish a structured event to a device's Redis pub/sub channel.

    Mirrors ``publish_event`` above (per-project channel) for the device-side WS channel
    (``/ws/devices/{device_id}``) used to deliver ``run.assigned`` when a device is connected.
    """
    msg = {
        "event": event,
        "device_id": device_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "data": data,
    }
    channel = f"{_DEVICE_CHANNEL_PREFIX}:{device_id}:events"
    try:
        await redis.publish(channel, json.dumps(msg))
    except Exception as exc:
        # Non-fatal: the run stays queued and is picked up via GET /api/devices/{id}/runs.
        logger.warning("pubsub_publish_failed", channel=channel, event=event, error=str(exc))


async def subscribe_device_events(
    redis: aioredis.Redis,
    device_id: str,
) -> AsyncIterator[dict]:
    """Subscribe to a device's event channel and yield parsed event dicts. Mirrors
    ``subscribe_project_events``; caller closes the Redis connection after use."""
    channel = f"{_DEVICE_CHANNEL_PREFIX}:{device_id}:events"
    pubsub = redis.pubsub()
    await pubsub.subscribe(channel)
    logger.info("pubsub_subscribed", channel=channel)
    try:
        while True:
            raw_msg = await pubsub.get_message(
                ignore_subscribe_messages=True,
                timeout=_PUBLISH_POLL_TIMEOUT_S,
            )
            if raw_msg is None:
                continue
            if raw_msg.get("type") == "message":
                try:
                    yield json.loads(raw_msg["data"])
                except (json.JSONDecodeError, KeyError, TypeError):
                    continue
    finally:
        try:
            await pubsub.unsubscribe(channel)
            await pubsub.aclose()
        except Exception:
            pass
        logger.info("pubsub_unsubscribed", channel=channel)