"""
Web Channel Adapter — WebSocket endpoint and message normalizer.

The web channel lets portal users chat with AI employees directly from
the Konstruct portal UI. Messages flow through the same agent pipeline
as Slack and WhatsApp — the only difference is the transport layer.

Message flow:
  1. Browser opens WebSocket at /chat/ws/{conversation_id}
  2. Client sends {"type": "auth", "userId": ..., "role": ..., "tenantId": ...}
     NOTE: Browsers cannot set custom HTTP headers on WebSocket connections,
     so auth credentials are sent as the first JSON message (Pitfall 1).
  3. For each user message (type="message"):
     a. Server immediately sends {"type": "typing"} to client (CHAT-05)
     b. normalize_web_event() converts to KonstructMessage (channel=WEB)
     c. User message saved to web_conversation_messages
     d. handle_message.delay(msg | extras) dispatches to Celery pipeline
     e. Server subscribes to Redis pub-sub channel for the response
     f. When orchestrator publishes the response:
        - Save assistant message to web_conversation_messages
        - Send {"type": "response", "text": ..., "conversation_id": ...} to client
  4. On disconnect: unsubscribe and close all Redis connections

Design notes:
  - thread_id = conversation_id — scopes agent memory to one conversation (Pitfall 3)
  - Redis pub-sub connections closed in try/finally to prevent leaks (Pitfall 2)
  - DB access uses configure_rls_hook + current_tenant_id context var per project pattern
  - WebSocket is a long-lived connection; each message/response cycle is synchronous
    within the connection but non-blocking for other connections

Push notifications:
  - Connected users are tracked in _connected_users (in-memory dict)
  - When the WebSocket send for "done" raises (client disconnected mid-stream),
    a push notification is fired so the user sees the response on their device.
"""

from __future__ import annotations

import asyncio
import json
import logging
import uuid
from datetime import datetime, timezone
from typing import Any

import redis.asyncio as aioredis
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from sqlalchemy import delete, select, text

from orchestrator.agents.builder import build_messages_with_memory, build_system_prompt
from orchestrator.agents.runner import run_agent_streaming
from orchestrator.memory.short_term import get_recent_messages, append_message
from orchestrator.tasks import handle_message, embed_and_store
from shared.config import settings
from shared.db import async_session_factory, engine
from shared.models.chat import WebConversation, WebConversationMessage
from shared.models.message import ChannelType, KonstructMessage, MessageContent, SenderInfo
from shared.models.tenant import Agent
from shared.redis_keys import webchat_response_key
from shared.rls import configure_rls_hook, current_tenant_id

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Router — mounted in gateway/main.py
# ---------------------------------------------------------------------------
web_chat_router = APIRouter(tags=["web-chat"])

# Timeout for waiting for an agent response via Redis pub-sub (seconds)
_RESPONSE_TIMEOUT_SECONDS = 180

# ---------------------------------------------------------------------------
# Connected user tracking — used to decide whether to send push notifications
# ---------------------------------------------------------------------------
# Maps user_id -> set of conversation_ids with active WebSocket connections.
# When a user disconnects, their entry is removed. If the agent response
# finishes after disconnect, a push notification is sent.
_connected_users: dict[str, set[str]] = {}


def _mark_connected(user_id: str, conversation_id: str) -> None:
    """Record that user_id has an active WebSocket for conversation_id."""
    if user_id not in _connected_users:
        _connected_users[user_id] = set()
    _connected_users[user_id].add(conversation_id)


def _mark_disconnected(user_id: str, conversation_id: str) -> None:
    """Remove the active WebSocket record for user_id + conversation_id."""
    if user_id in _connected_users:
        _connected_users[user_id].discard(conversation_id)
        if not _connected_users[user_id]:
            del _connected_users[user_id]


def is_user_connected(user_id: str) -> bool:
    """Return True if the user has any active WebSocket connection."""
    return user_id in _connected_users and bool(_connected_users[user_id])


async def _send_push_notification(
    user_id: str,
    title: str,
    body: str,
    conversation_id: str | None = None,
) -> None:
    """
    Fire-and-forget push notification delivery.

    Queries push_subscriptions for the user and calls pywebpush directly.
    Deletes stale (410 Gone) subscriptions automatically.
    Silently ignores errors — push is best-effort.
    """
    from shared.models.push import PushSubscription
    from shared.api.push import _send_push

    try:
        user_uuid = uuid.UUID(user_id)
        payload = {
            "title": title,
            "body": body,
            "data": {"conversationId": conversation_id},
        }

        async with async_session_factory() as session:
            result = await session.execute(
                select(PushSubscription).where(PushSubscription.user_id == user_uuid)
            )
            subscriptions = result.scalars().all()

            if not subscriptions:
                return

            stale_endpoints: list[str] = []
            for sub in subscriptions:
                try:
                    ok = await _send_push(sub, payload)
                    if not ok:
                        stale_endpoints.append(sub.endpoint)
                except Exception as exc:
                    logger.warning("Push delivery failed for user=%s: %s", user_id, exc)

            if stale_endpoints:
                await session.execute(
                    delete(PushSubscription).where(
                        PushSubscription.user_id == user_uuid,
                        PushSubscription.endpoint.in_(stale_endpoints),
                    )
                )
                await session.commit()

    except Exception as exc:
        logger.warning("Push notification send error for user=%s: %s", user_id, exc)


def normalize_web_event(event: dict[str, Any]) -> KonstructMessage:
    """
    Normalize a web channel event dict into a KonstructMessage.

    The web channel normalizer sets thread_id = conversation_id so that
    the agent memory pipeline scopes context to this conversation (Pitfall 3).

    Args:
        event: Dict with keys: text, tenant_id, agent_id, user_id,
               display_name, conversation_id.

    Returns:
        KonstructMessage with channel=WEB, thread_id=conversation_id.
    """
    tenant_id: str = event.get("tenant_id", "") or ""
    user_id: str = event.get("user_id", "") or ""
    display_name: str = event.get("display_name", "Portal User")
    conversation_id: str = event.get("conversation_id", "") or ""
    text_content: str = event.get("text", "") or ""

    return KonstructMessage(
        id=str(uuid.uuid4()),
        tenant_id=tenant_id,
        channel=ChannelType.WEB,
        channel_metadata={
            "portal_user_id": user_id,
            "tenant_id": tenant_id,
            "conversation_id": conversation_id,
        },
        sender=SenderInfo(
            user_id=user_id,
            display_name=display_name,
        ),
        content=MessageContent(
            text=text_content,
        ),
        timestamp=datetime.now(timezone.utc),
        thread_id=conversation_id,
        reply_to=None,
        context={},
    )


async def _handle_websocket_connection(
    websocket: WebSocket,
    conversation_id: str,
) -> None:
    """
    Core WebSocket connection handler — separated for testability.

    Lifecycle:
      1. Accept connection
      2. Wait for auth message (browser cannot send custom headers)
      3. Loop: receive messages → type indicator → Celery dispatch → Redis subscribe → response

    Args:
        websocket:       The FastAPI WebSocket connection.
        conversation_id: The conversation UUID from the URL path.
    """
    await websocket.accept()

    # -------------------------------------------------------------------------
    # Step 1: Auth handshake
    # Browsers cannot send custom HTTP headers on WebSocket connections.
    # Auth credentials are sent as the first JSON message.
    # -------------------------------------------------------------------------
    try:
        auth_msg = await websocket.receive_json()
    except WebSocketDisconnect:
        return

    if auth_msg.get("type") != "auth":
        await websocket.send_json({"type": "error", "message": "First message must be auth"})
        await websocket.close(code=4001)
        return

    user_id_str: str = auth_msg.get("userId", "") or ""
    user_role: str = auth_msg.get("role", "") or ""
    tenant_id_str: str = auth_msg.get("tenantId", "") or ""

    if not user_id_str or not tenant_id_str:
        await websocket.send_json({"type": "error", "message": "Missing userId or tenantId in auth"})
        await websocket.close(code=4001)
        return

    # Validate UUID format
    try:
        uuid.UUID(user_id_str)
        tenant_uuid = uuid.UUID(tenant_id_str)
    except (ValueError, AttributeError):
        await websocket.send_json({"type": "error", "message": "Invalid UUID format in auth"})
        await websocket.close(code=4001)
        return

    logger.info(
        "WebSocket auth: user=%s role=%s tenant=%s conversation=%s",
        user_id_str, user_role, tenant_id_str, conversation_id,
    )

    # Track this user as connected (for push notification gating)
    _mark_connected(user_id_str, conversation_id)

    # -------------------------------------------------------------------------
    # Step 2: Message loop
    # -------------------------------------------------------------------------
    try:
        while True:
            try:
                msg_data = await websocket.receive_json()
            except (WebSocketDisconnect, Exception):
                break

            if msg_data.get("type") != "message":
                continue

            text_content: str = msg_data.get("text", "") or ""
            agent_id_str: str = msg_data.get("agentId", "") or ""
            msg_conversation_id: str = msg_data.get("conversationId", conversation_id) or conversation_id
            display_name: str = msg_data.get("displayName", "Portal User")

            # -------------------------------------------------------------------
            # a. Send typing indicator IMMEDIATELY — before any DB or Celery work
            # -------------------------------------------------------------------
            await websocket.send_json({"type": "typing"})

            # -------------------------------------------------------------------
            # b. Save user message to web_conversation_messages
            # -------------------------------------------------------------------
            configure_rls_hook(engine)
            rls_token = current_tenant_id.set(tenant_uuid)
            saved_conversation_id = msg_conversation_id

            try:
                async with async_session_factory() as session:
                    # Look up the conversation to get tenant-scoped context
                    conv_stmt = select(WebConversation).where(
                        WebConversation.id == uuid.UUID(msg_conversation_id)
                    )
                    conv_result = await session.execute(conv_stmt)
                    conversation = conv_result.scalar_one_or_none()

                    if conversation is not None:
                        # Save user message
                        user_msg = WebConversationMessage(
                            conversation_id=uuid.UUID(msg_conversation_id),
                            tenant_id=tenant_uuid,
                            role="user",
                            content=text_content,
                        )
                        session.add(user_msg)

                        # Update conversation timestamp
                        await session.execute(
                            text(
                                "UPDATE web_conversations SET updated_at = NOW() WHERE id = :conv_id"
                            ),
                            {"conv_id": str(msg_conversation_id)},
                        )
                        await session.commit()
                        saved_conversation_id = msg_conversation_id
            except Exception:
                logger.exception(
                    "Failed to save user message for conversation=%s", msg_conversation_id
                )
            finally:
                current_tenant_id.reset(rls_token)

            # -------------------------------------------------------------------
            # c. Build KonstructMessage and stream LLM response DIRECTLY
            #
            # Bypasses Celery entirely for web chat — calls the LLM pool's
            # streaming endpoint from the WebSocket handler. This eliminates
            # ~5-10s of Celery queue + Redis pub-sub round-trip overhead.
            # Slack/WhatsApp still use Celery (async webhook pattern).
            # -------------------------------------------------------------------
            event = {
                "text": text_content,
                "tenant_id": tenant_id_str,
                "agent_id": agent_id_str,
                "user_id": user_id_str,
                "display_name": display_name,
                "conversation_id": saved_conversation_id,
            }
            normalized_msg = normalize_web_event(event)

            # Load agent for this tenant
            agent: Agent | None = None
            rls_token3 = current_tenant_id.set(tenant_uuid)
            try:
                async with async_session_factory() as session:
                    from sqlalchemy import select as sa_select
                    agent_stmt = sa_select(Agent).where(
                        Agent.tenant_id == tenant_uuid,
                        Agent.is_active == True,
                    ).limit(1)
                    agent_result = await session.execute(agent_stmt)
                    agent = agent_result.scalar_one_or_none()
            finally:
                current_tenant_id.reset(rls_token3)

            if agent is None:
                await websocket.send_json({
                    "type": "done",
                    "text": "No active AI employee is configured for this workspace.",
                    "conversation_id": saved_conversation_id,
                })
                continue

            # Build memory-enriched messages (Redis sliding window only — fast)
            redis_mem = aioredis.from_url(settings.redis_url)
            try:
                recent_messages = await get_recent_messages(
                    redis_mem, tenant_id_str, str(agent.id), user_id_str
                )
            finally:
                await redis_mem.aclose()

            enriched_messages = build_messages_with_memory(
                agent=agent,
                current_message=text_content,
                recent_messages=recent_messages,
                relevant_context=[],
                channel="web",
            )

            # Stream LLM response directly to WebSocket — no Celery, no pub-sub
            response_text = ""
            ws_disconnected_during_stream = False
            try:
                async for token in run_agent_streaming(
                    msg=normalized_msg,
                    agent=agent,
                    messages=enriched_messages,
                ):
                    response_text += token
                    try:
                        await websocket.send_json({"type": "chunk", "text": token})
                    except Exception:
                        ws_disconnected_during_stream = True
                        break  # Client disconnected
            except Exception:
                logger.exception("Direct streaming failed for conversation=%s", saved_conversation_id)
                if not response_text:
                    response_text = "I encountered an error processing your message. Please try again."

            # Save to Redis sliding window (fire-and-forget, non-blocking)
            redis_mem2 = aioredis.from_url(settings.redis_url)
            try:
                await append_message(redis_mem2, tenant_id_str, str(agent.id), user_id_str, "user", text_content)
                if response_text:
                    await append_message(redis_mem2, tenant_id_str, str(agent.id), user_id_str, "assistant", response_text)
            finally:
                await redis_mem2.aclose()

            # Fire-and-forget embedding for long-term memory
            try:
                embed_and_store.delay({
                    "tenant_id": tenant_id_str,
                    "agent_id": str(agent.id),
                    "user_id": user_id_str,
                    "role": "user",
                    "content": text_content,
                })
                if response_text:
                    embed_and_store.delay({
                        "tenant_id": tenant_id_str,
                        "agent_id": str(agent.id),
                        "user_id": user_id_str,
                        "role": "assistant",
                        "content": response_text,
                    })
            except Exception:
                pass  # Non-fatal — memory will rebuild over time

            # -------------------------------------------------------------------
            # e. Save assistant message and send final "done" to client
            # -------------------------------------------------------------------
            if response_text:
                rls_token2 = current_tenant_id.set(tenant_uuid)
                try:
                    async with async_session_factory() as session:
                        assistant_msg = WebConversationMessage(
                            conversation_id=uuid.UUID(saved_conversation_id),
                            tenant_id=tenant_uuid,
                            role="assistant",
                            content=response_text,
                        )
                        session.add(assistant_msg)
                        await session.execute(
                            text(
                                "UPDATE web_conversations SET updated_at = NOW() WHERE id = :conv_id"
                            ),
                            {"conv_id": str(saved_conversation_id)},
                        )
                        await session.commit()
                except Exception:
                    logger.exception(
                        "Failed to save assistant message for conversation=%s", saved_conversation_id
                    )
                finally:
                    current_tenant_id.reset(rls_token2)

                # If user disconnected during streaming, send push notification
                if ws_disconnected_during_stream or not is_user_connected(user_id_str):
                    agent_name = agent.name if hasattr(agent, "name") and agent.name else "Your AI Employee"
                    preview = response_text[:100] + ("..." if len(response_text) > 100 else "")
                    asyncio.create_task(
                        _send_push_notification(
                            user_id=user_id_str,
                            title=f"{agent_name} replied",
                            body=preview,
                            conversation_id=saved_conversation_id,
                        )
                    )
                    if ws_disconnected_during_stream:
                        break  # Stop the message loop — WS is gone

                # Signal stream completion to the client
                try:
                    await websocket.send_json({
                        "type": "done",
                        "text": response_text,
                        "conversation_id": saved_conversation_id,
                    })
                except Exception:
                    pass  # Client already disconnected
            else:
                logger.warning(
                    "No response received for conversation=%s", saved_conversation_id,
                )
                try:
                    await websocket.send_json({
                        "type": "error",
                        "message": "I'm having trouble responding right now. Please try again.",
                    })
                except Exception:
                    pass  # Client already disconnected

    finally:
        # Always untrack this user when connection ends
        _mark_disconnected(user_id_str, conversation_id)


@web_chat_router.websocket("/chat/ws/{conversation_id}")
async def chat_websocket(websocket: WebSocket, conversation_id: str) -> None:
    """
    WebSocket endpoint for web chat.

    URL: /chat/ws/{conversation_id}

    Protocol:
      1. Connect
      2. Send: {"type": "auth", "userId": "...", "role": "...", "tenantId": "..."}
      3. Send: {"type": "message", "text": "...", "agentId": "...", "conversationId": "..."}
      4. Receive: {"type": "typing"}
      5. Receive: {"type": "response", "text": "...", "conversation_id": "..."}

    Closes with code 4001 on auth failure.
    """
    try:
        await _handle_websocket_connection(websocket, conversation_id)
    except WebSocketDisconnect:
        logger.info("WebSocket disconnected for conversation=%s", conversation_id)
    except Exception:
        logger.exception("Unhandled error in WebSocket handler for conversation=%s", conversation_id)
