"""
ThreadUpdatesConsumer

WebSocket consumer for real-time thread/conversation updates.
Clients subscribe to this to receive:
- Agent response streaming tokens
- Tool call notifications
- Response completion events
- Error notifications

This consumer is read-only - it only broadcasts updates from Celery tasks.
The actual agent responses are generated by the generate_agent_response task.

Query Parameters:
    conversation_id: Required GraphQL ID for the conversation to watch
"""

from __future__ import annotations

import json
import logging
import urllib.parse
import uuid

from channels.db import database_sync_to_async
from channels.generic.websocket import AsyncWebsocketConsumer
from graphql_relay import from_global_id

from config.ratelimit.decorators import check_ws_rate_limit
from config.websocket.auth_handshake import AuthHandshakeMixin
from config.websocket.middleware import WS_CLOSE_RATE_LIMITED
from config.websocket.utils.auth_helpers import check_auth_and_close_if_failed
from opencontractserver.conversations.models import Conversation
from opencontractserver.corpuses.models import Corpus
from opencontractserver.documents.models import Document

logger = logging.getLogger(__name__)


def get_thread_channel_group(conversation_id: int) -> str:
    """Get the channel group name for a conversation/thread."""
    return f"thread_{conversation_id}"


class ThreadUpdatesConsumer(AuthHandshakeMixin, AsyncWebsocketConsumer):
    """
    WebSocket consumer for subscribing to thread updates.

    Clients connect with a conversation_id and receive real-time updates
    about agent responses being generated in that thread.

    This consumer does NOT handle sending messages - that's done via GraphQL.
    It only receives and broadcasts updates from Celery tasks.
    """

    conversation: Conversation | None = None
    conversation_id: int | None = None
    session_id: str | None = None
    user_id: int | None = None

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.consumer_id = uuid.uuid4()
        logger.debug(f"[ThreadUpdates {self.consumer_id}] __init__ called.")

    # -------------------------------------------------------------------------
    #  WebSocket lifecycle
    # -------------------------------------------------------------------------

    async def connect(self) -> None:
        """
        Authenticate, extract conversation_id, validate access, and subscribe.
        """
        self.session_id = str(uuid.uuid4())
        logger.debug(
            f"[ThreadUpdates {self.consumer_id} | Session {self.session_id}] "
            f"connect() called. Path: {self.scope['path']}"
        )

        # Rate limit new connections (skip JSON message — connection
        # is about to be closed so the client won't see it)
        if await check_ws_rate_limit(self, "WS_CONNECT", send_message=False):
            await self.close(code=WS_CLOSE_RATE_LIMITED)
            return

        # Authenticate user (set by auth middleware)
        if await check_auth_and_close_if_failed(self, self.session_id):
            return

        user = self.scope.get("user")
        self.user_id = user.pk

        # Parse query parameters
        query_string = self.scope.get("query_string", b"").decode("utf-8")
        params = dict(urllib.parse.parse_qsl(query_string))

        raw_conversation_id = params.get("conversation_id")
        if not raw_conversation_id:
            logger.warning(
                f"[ThreadUpdates {self.consumer_id}] No conversation_id provided"
            )
            await self.close(code=4002)
            return

        # Decode GraphQL ID
        try:
            _, pk_str = from_global_id(raw_conversation_id)
            self.conversation_id = int(pk_str)
        except Exception:
            try:
                self.conversation_id = int(raw_conversation_id)
            except ValueError:
                logger.warning(
                    f"[ThreadUpdates {self.consumer_id}] Invalid conversation_id: {raw_conversation_id}"
                )
                await self.close(code=4002)
                return

        # Validate access to conversation
        has_access = await self._check_conversation_access(
            user, cache_conversation=True
        )
        if not has_access:
            logger.warning(
                f"[ThreadUpdates {self.consumer_id}] Access denied for conversation {self.conversation_id}"
            )
            await self.close(code=4003)
            return

        # Subscribe to thread channel group
        self.room_group_name = get_thread_channel_group(self.conversation_id)
        await self.channel_layer.group_add(self.room_group_name, self.channel_name)

        await self.accept_with_auth()
        logger.info(
            f"[ThreadUpdates {self.consumer_id}] Subscribed to {self.room_group_name}"
        )

        # Send connection confirmation
        await self.send(
            text_data=json.dumps(
                {
                    "type": "CONNECTED",
                    "conversation_id": str(self.conversation_id),
                    "session_id": self.session_id,
                }
            )
        )

    async def disconnect(self, close_code: int) -> None:
        """Leave the thread channel group on disconnect."""
        await self.cleanup_auth_handshake()
        if hasattr(self, "room_group_name") and self.room_group_name:
            await self.channel_layer.group_discard(
                self.room_group_name, self.channel_name
            )
            logger.info(
                f"[ThreadUpdates {self.consumer_id}] Disconnected from {self.room_group_name}"
            )

    async def receive(self, text_data: str) -> None:
        """
        Handle incoming messages from client.

        This consumer is primarily for receiving broadcasts, but we handle
        a few client-initiated message types for connection management.
        """
        try:
            data = json.loads(text_data)
        except json.JSONDecodeError:
            logger.warning(f"[ThreadUpdates {self.consumer_id}] Invalid JSON received")
            return

        if isinstance(data, dict) and data.get("type") == "AUTH":
            await self.handle_auth_message(data)
            return

        if await check_ws_rate_limit(self, "WS_HEARTBEAT"):
            return

        msg_type = data.get("type", "") if isinstance(data, dict) else ""

        if msg_type == "ping":
            await self.send(text_data=json.dumps({"type": "pong"}))

        elif msg_type == "heartbeat":
            await self.send(
                text_data=json.dumps(
                    {
                        "type": "heartbeat_ack",
                        "session_id": self.session_id,
                    }
                )
            )

        else:
            logger.debug(
                f"[ThreadUpdates {self.consumer_id}] Unknown message type: {msg_type}"
            )

    # -------------------------------------------------------------------------
    #  Channel layer message handlers (from Celery tasks)
    # -------------------------------------------------------------------------

    async def agent_stream_start(self, event: dict) -> None:
        """Handle agent starting to generate a response."""
        await self.send(
            text_data=json.dumps(
                {
                    "type": "AGENT_STREAM_START",
                    "message_id": event.get("message_id"),
                    "agent_id": event.get("agent_id"),
                    "agent_name": event.get("agent_name"),
                    "agent_slug": event.get("agent_slug"),
                }
            )
        )

    async def agent_stream_token(self, event: dict) -> None:
        """Handle streaming token from agent."""
        await self.send(
            text_data=json.dumps(
                {
                    "type": "AGENT_STREAM_TOKEN",
                    "message_id": event.get("message_id"),
                    "token": event.get("token"),
                }
            )
        )

    async def agent_tool_call(self, event: dict) -> None:
        """Handle agent calling a tool."""
        await self.send(
            text_data=json.dumps(
                {
                    "type": "AGENT_TOOL_CALL",
                    "message_id": event.get("message_id"),
                    "tool": event.get("tool"),
                    "args": event.get("args"),
                }
            )
        )

    async def agent_stream_complete(self, event: dict) -> None:
        """Handle agent completing response generation."""
        await self.send(
            text_data=json.dumps(
                {
                    "type": "AGENT_STREAM_COMPLETE",
                    "message_id": event.get("message_id"),
                    "content": event.get("content"),
                    "sources": event.get("sources", []),
                    "timeline": event.get("timeline", []),
                }
            )
        )

    async def agent_stream_error(self, event: dict) -> None:
        """Handle agent error during response generation."""
        await self.send(
            text_data=json.dumps(
                {
                    "type": "AGENT_STREAM_ERROR",
                    "message_id": event.get("message_id"),
                    "error": event.get("error"),
                }
            )
        )

    # -------------------------------------------------------------------------
    #  Permission check (shared between connect() and AuthHandshakeMixin refresh)
    # -------------------------------------------------------------------------

    @database_sync_to_async
    def _check_conversation_access(self, user, *, cache_conversation: bool) -> bool:
        """
        Check if ``user`` may participate in ``self.conversation_id``.

        For conversations with BOTH chat_with_corpus AND chat_with_document set
        (doc-in-corpus threads), the user must have access to BOTH the corpus
        AND the document (AND logic). For single-context conversations the user
        only needs access to that one resource.

        ``cache_conversation`` controls whether the resolved Conversation is
        attached to ``self.conversation`` for downstream use; the connect path
        wants this side effect, the refresh path does not.
        """
        try:
            conversation = Conversation.objects.get(pk=self.conversation_id)
        except Conversation.DoesNotExist:
            return False

        if cache_conversation:
            self.conversation = conversation

        if conversation.creator_id == user.pk or user.is_superuser:
            return True

        has_corpus = (
            Corpus.objects.visible_to_user(user)
            .filter(pk=conversation.chat_with_corpus_id)
            .exists()
            if conversation.chat_with_corpus_id
            else None
        )
        has_document = (
            Document.objects.visible_to_user(user)
            .filter(pk=conversation.chat_with_document_id)
            .exists()
            if conversation.chat_with_document_id
            else None
        )

        if has_corpus is not None and has_document is not None:
            return has_corpus and has_document
        if has_corpus is not None:
            return has_corpus
        if has_document is not None:
            return has_document
        return False

    # -------------------------------------------------------------------------
    #  AuthHandshakeMixin override
    # -------------------------------------------------------------------------

    async def _validate_resource_permissions(self, user) -> bool:
        """Re-validate conversation access on token refresh."""
        if self.conversation_id is None:
            return True
        return await self._check_conversation_access(user, cache_conversation=False)
