"""Real-time WebSocket gateway.

Multiplexes presence, DM, and world-chat delivery over a single WebSocket
endpoint (docs/planning/social.md US-2, US-4).

Key features:
- Authentication via JWT token (query param or initial auth frame).
- Scoped subscriptions to DM conversations and world channels.
- Server-derived presence heartbeat (Redis TTL refreshed periodically).
- Broadcast presence changes (online/offline) to mutual friends.
- Fan-out via Redis pub/sub.
- Thread-safe & event-loop-safe dispatch for multi-instance and test environments.
- Graceful reconnect without duplicate delivery or orphaned state.
"""

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager
from datetime import datetime, timezone
import json
import logging
from typing import Any

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status
from jose import JWTError, jwt
from sqlalchemy import or_, select
from sqlalchemy.ext.asyncio import AsyncSession
import redis.asyncio as aioredis

from app.config import settings
from app.database import async_session_factory, get_db
from app.models.chat import ConversationParticipant
from app.models.social import Friendship, FriendshipStatus
from app.models.user import User
from app.models.world import World, WorldVisibility
from app.services.presence import presence_store

logger = logging.getLogger(__name__)

router = APIRouter()


@asynccontextmanager
async def get_db_session():
    """Helper to acquire an async DB session respecting test overrides if present."""
    from app.main import app

    override = app.dependency_overrides.get(get_db)
    if override:
        res = override()
        if hasattr(res, "__anext__"):
            session = await res.__anext__()
            try:
                yield session
            finally:
                try:
                    await res.__anext__()
                except (StopAsyncIteration, Exception):
                    pass
            return
        elif hasattr(res, "__next__"):
            session = next(res)
            try:
                yield session
            finally:
                try:
                    next(res)
                except (StopIteration, Exception):
                    pass
            return

    async with async_session_factory() as session:
        yield session


async def authenticate_token(token: str) -> User | None:
    """Validate JWT token and return authenticated User, or None."""
    try:
        payload = jwt.decode(token, settings.public_pem, algorithms=[settings.ALGORITHM])
        user_id_raw = payload.get("sub")
        if user_id_raw is None:
            return None
        user_id = int(user_id_raw)
    except (JWTError, ValueError, TypeError):
        return None

    async with get_db_session() as db:
        res = await db.execute(select(User).where(User.id == user_id))
        user = res.scalar_one_or_none()
        if user and user.is_active:
            return user
    return None


async def get_friend_ids(user_id: int) -> set[int]:
    """Return set of accepted friend user IDs for user_id."""
    try:
        async with get_db_session() as db:
            res = await db.execute(
                select(Friendship).where(
                    or_(Friendship.requester_id == user_id, Friendship.addressee_id == user_id),
                    Friendship.status == FriendshipStatus.ACCEPTED,
                )
            )
            friendships = res.scalars().all()
            return {
                f.addressee_id if f.requester_id == user_id else f.requester_id
                for f in friendships
            }
    except Exception as exc:
        logger.error("Failed to get friend IDs for user %s: %s", user_id, exc, exc_info=True)
        return set()


async def check_world_access(world_id: int, user_id: int) -> bool:
    """Verify if user has permission to join/view a world."""
    async with get_db_session() as db:
        res = await db.execute(select(World).where(World.id == world_id))
        world = res.scalar_one_or_none()
        if not world:
            return False
        if world.visibility == WorldVisibility.PUBLIC:
            return True
        if world.owner_id == user_id:
            return True
        if world.visibility == WorldVisibility.FRIENDS_ONLY:
            res_fr = await db.execute(
                select(Friendship.id).where(
                    or_(
                        and_(Friendship.requester_id == world.owner_id, Friendship.addressee_id == user_id),
                        and_(Friendship.requester_id == user_id, Friendship.addressee_id == world.owner_id),
                    ),
                    Friendship.status == FriendshipStatus.ACCEPTED,
                ).limit(1)
            )
            return res_fr.scalar_one_or_none() is not None
        return False


async def check_conversation_access(conversation_id: int, user_id: int) -> bool:
    """Verify if user is a participant in the given conversation."""
    async with get_db_session() as db:
        res = await db.execute(
            select(ConversationParticipant.id).where(
                ConversationParticipant.conversation_id == conversation_id,
                ConversationParticipant.user_id == user_id,
            ).limit(1)
        )
        return res.scalar_one_or_none() is not None


async def _async_update_last_seen(user_id: int, now: datetime):
    try:
        async with async_session_factory() as session:
            res = await session.execute(select(User).where(User.id == user_id))
            db_user = res.scalar_one_or_none()
            if db_user:
                db_user.last_seen_at = now
                await session.commit()
    except Exception:
        pass


class ConnectionManager:
    """Manages active WebSocket connections, subscriptions, presence, and fan-out."""

    def __init__(self):
        # ws -> User
        self._connections: dict[WebSocket, User] = {}
        # ws -> event loop id
        self._ws_loops: dict[WebSocket, int] = {}
        # user_id -> set of active WebSockets
        self._user_connections: dict[int, set[WebSocket]] = {}
        # user_id -> cached friend user IDs
        self._user_friend_ids: dict[int, set[int]] = {}
        # world_id -> set of WebSockets
        self._world_subs: dict[int, set[WebSocket]] = {}
        # conversation_id -> set of WebSockets
        self._dm_subs: dict[int, set[WebSocket]] = {}
        # ws -> asyncio heartbeat task
        self._heartbeat_tasks: dict[WebSocket, asyncio.Task] = {}
        # Per-event-loop Redis pub/sub listener tasks and events
        self._redis_tasks: dict[int, asyncio.Task] = {}
        self._subscribed_events: dict[int, asyncio.Event] = {}

    async def ensure_redis_listener(self):
        """Ensure background Redis pub/sub listener is connected and subscribed for current loop."""
        loop = asyncio.get_running_loop()
        loop_key = id(loop)
        task = self._redis_tasks.get(loop_key)
        if task is None or task.done():
            event = asyncio.Event()
            self._subscribed_events[loop_key] = event
            task = asyncio.create_task(self._redis_listener(loop_key, event))
            self._redis_tasks[loop_key] = task
        else:
            event = self._subscribed_events.get(loop_key)

        if event:
            try:
                await asyncio.wait_for(event.wait(), timeout=2.0)
            except asyncio.TimeoutError:
                pass

    async def connect(self, websocket: WebSocket, user: User):
        """Register an authenticated connection."""
        await self.ensure_redis_listener()
        loop_key = id(asyncio.get_running_loop())

        self._connections[websocket] = user
        self._ws_loops[websocket] = loop_key
        self._user_connections.setdefault(user.id, set()).add(websocket)

        # Cache friend IDs for instant presence dispatch
        friend_ids = await get_friend_ids(user.id)
        self._user_friend_ids[user.id] = friend_ids

        # Mark user online in Redis and DB
        await presence_store.heartbeat(user.id)
        now = datetime.now(timezone.utc)
        asyncio.create_task(_async_update_last_seen(user.id, now))

        # Start periodic heartbeat task (every 20 seconds)
        self._heartbeat_tasks[websocket] = asyncio.create_task(
            self._run_periodic_heartbeat(user.id)
        )

        # Broadcast online presence to friends via Redis
        await self.broadcast_presence(user.id, user.username, is_online=True, friend_ids=friend_ids)

    async def disconnect(self, websocket: WebSocket):
        """Unregister a connection and clean up subscriptions and presence."""
        user = self._connections.pop(websocket, None)
        self._ws_loops.pop(websocket, None)
        task = self._heartbeat_tasks.pop(websocket, None)
        if task:
            task.cancel()

        # Clean subscriptions
        for world_id in list(self._world_subs.keys()):
            self._world_subs[world_id].discard(websocket)
            if not self._world_subs[world_id]:
                self._world_subs.pop(world_id, None)

        for conv_id in list(self._dm_subs.keys()):
            self._dm_subs[conv_id].discard(websocket)
            if not self._dm_subs[conv_id]:
                self._dm_subs.pop(conv_id, None)

        if user is None:
            return

        user_id = user.id
        username = user.username
        user_set = self._user_connections.get(user_id)
        if user_set:
            user_set.discard(websocket)
            if not user_set:
                self._user_connections.pop(user_id, None)

        # If user has no more open connections, mark offline immediately
        if user_id not in self._user_connections:
            await presence_store.clear_heartbeat(user_id)
            now = datetime.now(timezone.utc)
            asyncio.create_task(_async_update_last_seen(user_id, now))
            friend_ids = self._user_friend_ids.pop(user_id, set())
            await self.broadcast_presence(
                user_id, username, is_online=False, last_seen_at=now, friend_ids=friend_ids
            )

    async def _run_periodic_heartbeat(self, user_id: int):
        """Keep Redis TTL alive while connection is open."""
        try:
            while True:
                await asyncio.sleep(20)
                await presence_store.heartbeat(user_id)
        except asyncio.CancelledError:
            pass

    async def subscribe_world(self, websocket: WebSocket, world_id: int) -> bool:
        user = self._connections.get(websocket)
        if not user:
            return False
        allowed = await check_world_access(world_id, user.id)
        if not allowed:
            return False
        self._world_subs.setdefault(world_id, set()).add(websocket)
        return True

    def unsubscribe_world(self, websocket: WebSocket, world_id: int):
        if world_id in self._world_subs:
            self._world_subs[world_id].discard(websocket)

    async def subscribe_dm(self, websocket: WebSocket, conversation_id: int) -> bool:
        user = self._connections.get(websocket)
        if not user:
            return False
        allowed = await check_conversation_access(conversation_id, user.id)
        if not allowed:
            return False
        self._dm_subs.setdefault(conversation_id, set()).add(websocket)
        return True

    def unsubscribe_dm(self, websocket: WebSocket, conversation_id: int):
        if conversation_id in self._dm_subs:
            self._dm_subs[conversation_id].discard(websocket)

    async def send_to_ws(self, websocket: WebSocket, message: dict):
        """Deliver frame to websocket if it belongs to current event loop."""
        cur_loop = id(asyncio.get_running_loop())
        if self._ws_loops.get(websocket) != cur_loop:
            return
        try:
            await websocket.send_text(json.dumps(message, default=str))
        except Exception:
            pass

    async def broadcast_world_chat(
        self, world_id: int, message_payload: dict, room_id: int | None = None
    ):
        """Send world chat message to all connected clients on this loop subscribed to world_id."""
        targets = list(self._world_subs.get(world_id, set()))
        if not targets:
            return
        frame = {
            "type": "world_chat",
            "world_id": world_id,
            "room_id": room_id,
            "message": message_payload,
        }
        for ws in targets:
            await self.send_to_ws(ws, frame)

    async def broadcast_dm_chat(self, conversation_id: int, message_payload: dict):
        """Send DM message to all connected clients on this loop subscribed to conversation_id."""
        targets = list(self._dm_subs.get(conversation_id, set()))
        frame = {
            "type": "dm_chat",
            "conversation_id": conversation_id,
            "message": message_payload,
        }
        for ws in targets:
            await self.send_to_ws(ws, frame)

    async def broadcast_presence(
        self,
        user_id: int,
        username: str,
        is_online: bool,
        last_seen_at: datetime | None = None,
        friend_ids: set[int] | None = None,
    ):
        """Broadcast presence update to mutual friends via Redis."""
        if friend_ids is None:
            friend_ids = await get_friend_ids(user_id)
        if not friend_ids:
            return

        frame = {
            "type": "presence",
            "user_id": user_id,
            "username": username,
            "is_online": is_online,
            "last_seen_at": last_seen_at.isoformat() if last_seen_at else None,
            "friend_ids": list(friend_ids),
        }

        # Redis broadcast for all loops/instances
        try:
            client = aioredis.from_url(settings.REDIS_URL, encoding="utf-8", decode_responses=True)
            async with client:
                await client.publish("presence:events", json.dumps(frame, default=str))
        except Exception as exc:
            logger.error("Failed to publish presence to Redis: %s", exc, exc_info=True)

    async def _redis_listener(self, loop_key: int, event: asyncio.Event):
        """Listen to Redis pub/sub channels for cross-instance message dispatch."""
        while True:
            try:
                client = aioredis.from_url(settings.REDIS_URL, encoding="utf-8", decode_responses=True)
                pubsub = client.pubsub()
                await pubsub.psubscribe("world:*:chat", "dm:*:chat", "presence:events")
                event.set()
                logger.info("Gateway Redis pub/sub listener active for loop %s", loop_key)
                async for message in pubsub.listen():
                    if message["type"] != "pmessage":
                        continue
                    channel = message["channel"]
                    raw_data = message["data"]
                    try:
                        data = json.loads(raw_data)
                    except Exception:
                        continue

                    if channel.startswith("world:") and channel.endswith(":chat"):
                        parts = channel.split(":")
                        world_id = int(parts[1])
                        await self.broadcast_world_chat(world_id, data, room_id=data.get("room_id"))
                    elif channel.startswith("dm:") and channel.endswith(":chat"):
                        parts = channel.split(":")
                        conv_id = int(parts[1])
                        await self.broadcast_dm_chat(conv_id, data)
                    elif channel == "presence:events":
                        logger.info("Redis presence event received: %s. Active users: %s", data, list(self._user_connections.keys()))
                        for friend_id in data.get("friend_ids", []):
                            targets = list(self._user_connections.get(friend_id, set()))
                            logger.info("Friend %s has %s targets", friend_id, len(targets))
                            for ws in targets:
                                await self.send_to_ws(ws, data)
            except asyncio.CancelledError:
                break
            except Exception as exc:
                logger.debug("Redis listener encountered error, retrying in 2s: %s", exc)
                event.clear()
                await asyncio.sleep(2)


connection_manager = ConnectionManager()


# ──────────────────────────────── WebSocket Route ────────────────────────────


@router.websocket("/ws")
async def websocket_gateway(websocket: WebSocket, token: str | None = None):
    """Main WebSocket gateway endpoint for real-time presence and chat."""
    user: User | None = None
    # 1. Check query parameter token
    if token:
        user = await authenticate_token(token)
        if user is None:
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
            return

    await websocket.accept()

    # 2. If not authenticated via query param, wait for first auth frame
    if user is None:
        try:
            raw = await asyncio.wait_for(websocket.receive_text(), timeout=3.0)
            data = json.loads(raw)
            if data.get("type") == "auth" and data.get("token"):
                user = await authenticate_token(data["token"])
        except (asyncio.TimeoutError, json.JSONDecodeError, WebSocketDisconnect):
            pass

        if user is None:
            try:
                await websocket.send_text(json.dumps({"type": "error", "message": "Authentication failed"}))
                await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
            except Exception:
                pass
            return

    # Mark user online in Redis immediately upon successful auth
    await presence_store.heartbeat(user.id)

    # Send auth_ok handshake immediately upon successful authentication
    await websocket.send_text(
        json.dumps({
            "type": "auth_ok",
            "user_id": user.id,
            "username": user.username,
        })
    )

    # Register with connection manager for presence and real-time event delivery
    await connection_manager.connect(websocket, user)

    try:
        while True:
            text = await websocket.receive_text()
            try:
                msg = json.loads(text)
            except json.JSONDecodeError:
                await websocket.send_text(json.dumps({"type": "error", "message": "Invalid JSON"}))
                continue

            msg_type = msg.get("type")
            if msg_type == "ping":
                await websocket.send_text(json.dumps({"type": "pong"}))
            elif msg_type == "heartbeat":
                await presence_store.heartbeat(user.id)
                await websocket.send_text(json.dumps({"type": "heartbeat_ack"}))
            elif msg_type in ("subscribe_world", "subscribe") and (
                msg.get("world_id") or msg.get("channel") == "world"
            ):
                world_id = msg.get("world_id") or msg.get("id")
                if isinstance(world_id, int):
                    ok = await connection_manager.subscribe_world(websocket, world_id)
                    if ok:
                        await websocket.send_text(
                            json.dumps({"type": "subscribed", "channel": "world", "id": world_id})
                        )
                    else:
                        await websocket.send_text(
                            json.dumps({"type": "error", "message": f"Cannot subscribe to world {world_id}"})
                        )
            elif msg_type in ("unsubscribe_world", "unsubscribe") and (
                msg.get("world_id") or msg.get("channel") == "world"
            ):
                world_id = msg.get("world_id") or msg.get("id")
                if isinstance(world_id, int):
                    connection_manager.unsubscribe_world(websocket, world_id)
                    await websocket.send_text(
                        json.dumps({"type": "unsubscribed", "channel": "world", "id": world_id})
                    )
            elif msg_type in ("subscribe_dm", "subscribe") and (
                msg.get("conversation_id") or msg.get("channel") == "dm"
            ):
                conv_id = msg.get("conversation_id") or msg.get("id")
                if isinstance(conv_id, int):
                    ok = await connection_manager.subscribe_dm(websocket, conv_id)
                    if ok:
                        await websocket.send_text(
                            json.dumps({"type": "subscribed", "channel": "dm", "id": conv_id})
                        )
                    else:
                        await websocket.send_text(
                            json.dumps({"type": "error", "message": f"Cannot subscribe to conversation {conv_id}"})
                        )
            elif msg_type in ("unsubscribe_dm", "unsubscribe") and (
                msg.get("conversation_id") or msg.get("channel") == "dm"
            ):
                conv_id = msg.get("conversation_id") or msg.get("id")
                if isinstance(conv_id, int):
                    connection_manager.unsubscribe_dm(websocket, conv_id)
                    await websocket.send_text(
                        json.dumps({"type": "unsubscribed", "channel": "dm", "id": conv_id})
                    )
            else:
                await websocket.send_text(json.dumps({"type": "error", "message": f"Unknown action: {msg_type}"}))

    except WebSocketDisconnect:
        pass
    finally:
        await connection_manager.disconnect(websocket)