# @File: backend/routers/whatsapp_client_ws.py
# @Desc: WebSocket endpoint for the client WhatsApp chat page. Auth is done
#        via a `token` query parameter (JWT) since browsers can't set
#        Authorization headers on WebSocket handshakes.
#
# URL: /api/client/whatsapp/ws?token=<jwt>
#
# Inbound client events:
#   { "type": "typing", "contact_id": 123, "typing": true }
#   { "type": "ping" }
#
# Outbound server events:
#   { "type": "snapshot", "presence": [...], "typing_contact_ids": [...] }
#   { "type": "presence", "contact_id": 123, "online": true, "last_seen_at": "..." }
#   { "type": "typing",   "contact_id": 123, "typing": true }
#   { "type": "message",  "message": { ... whatsapp_messages row ... } }
#   { "type": "pong" }
from __future__ import annotations

import json
import logging
from typing import Optional

from core.auth import AccessTokenError, decode_access_token
from core.database import db_manager
from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect, status
from services.whatsapp.chat import _get_active_subscription  # internal helper, reused
from services.whatsapp.presence import presence_hub
from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)

router = APIRouter(tags=["whatsapp-client-ws"])


async def _resolve_subscription_id(user_id: str) -> Optional[int]:
    """Look up the caller's active WhatsApp subscription id."""
    try:
        if not db_manager.async_session_maker:
            await db_manager.ensure_initialized()
        if not db_manager.async_session_maker:
            return None
        async with db_manager.async_session_maker() as db:  # type: AsyncSession
            subscription = await _get_active_subscription(db, user_id)
            return subscription.id if subscription else None
    except Exception:  # noqa: BLE001
        logger.exception("failed to resolve whatsapp subscription for ws handshake")
        return None


@router.websocket("/api/client/whatsapp/ws")
async def whatsapp_client_ws(
    websocket: WebSocket,
    token: Optional[str] = Query(default=None),
) -> None:
    # --- auth ---
    if not token:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    try:
        payload = decode_access_token(token)
    except AccessTokenError:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    user_id = payload.get("sub")
    if not user_id:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    subscription_id = await _resolve_subscription_id(str(user_id))
    if not subscription_id:
        await websocket.accept()
        try:
            await websocket.send_json(
                {
                    "type": "error",
                    "detail": "no_active_subscription",
                }
            )
        finally:
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    # --- session ---
    await websocket.accept()
    await presence_hub.register(subscription_id, websocket)
    try:
        while True:
            raw = await websocket.receive_text()
            await _handle_inbound(subscription_id, websocket, raw)
    except WebSocketDisconnect:
        pass
    except Exception:  # noqa: BLE001
        logger.exception("whatsapp ws session error")
    finally:
        await presence_hub.unregister(subscription_id, websocket)


async def _handle_inbound(
    subscription_id: int, websocket: WebSocket, raw: str
) -> None:
    try:
        data = json.loads(raw) if raw else {}
    except json.JSONDecodeError:
        return
    if not isinstance(data, dict):
        return

    event_type = data.get("type")
    if event_type == "ping":
        try:
            await websocket.send_json({"type": "pong"})
        except Exception:  # noqa: BLE001
            pass
        return

    if event_type == "typing":
        contact_id = data.get("contact_id")
        if not isinstance(contact_id, int):
            return
        typing = bool(data.get("typing", False))
        # When the operator types, mark THEM as typing for the *selected*
        # contact thread. The other side of the conversation (the WhatsApp
        # contact) produces typing events via upstream webhooks.
        await presence_hub.publish_typing(
            subscription_id, contact_id, typing=typing
        )
        return