# clients/websocket_manager.py

import asyncio
import json
from typing import Dict, Optional

import websockets

from common.logging import get_logger, span
from common.platform.twitch.model import (
    WebSocketMessage,
    EventSubSubscription,
    EventSubSubscriptionRequest,
    EventSubTransport,
)
from common.platform.twitch.api import TwitchApi

logger = get_logger(__name__)


class WebSocketManager:
    def __init__(
        self,
        twitch_api: TwitchApi,
        event_handler,
        websocket_url: str = "wss://eventsub.wss.twitch.tv/ws",
        keepalive_timeout_seconds: int = 10,
    ):
        """
        Initialize the WebSocketManager.

        Args:
            twitch_api (TwitchApi): Instance to manage EventSub subscriptions.
            event_handler (Callable[[Dict[str, Any]], None]): Function to handle incoming events.
            websocket_url (str, optional): Twitch EventSub WebSocket URL. Defaults to "wss://eventsub.wss.twitch.tv/ws".
            keepalive_timeout_seconds (int, optional): Keepalive timeout in seconds. Defaults to 30.
        """
        self.twitch_api = twitch_api
        self.event_handler = event_handler
        self.websocket_url = websocket_url
        self.keepalive_timeout_seconds = keepalive_timeout_seconds
        self.connections: Dict[str, websockets.ClientConnection] = {}
        self.tasks: Dict[str, asyncio.Task] = {}
        self.session_ids: Dict[str, str] = {}
        self.welcome_events: Dict[str, asyncio.Event] = {}
        self.keep_running = True
        self.connection_closed_events: Dict[str, asyncio.Event] = {}

    def __aenter__(self):
        return self

    def __aexit__(self, exc_type, exc_val, exc_tb):
        return self.stop_all()

    async def connect_shard(self, shard_id: str, url: Optional[str] = None):
        """
        Establish a WebSocket connection for a specific shard.

        Args:
            shard_id (str): The shard identifier.
            url (Optional[str], optional): The WebSocket URL to connect to. Defaults to None.
        """
        attributes = {"shard_id": shard_id, "custom_url": url is not None}

        with span(logger, "connect_websocket_shard", attributes):
            if url is None:
                url = self.websocket_url
                logger.debug(f"Using default WebSocket URL: {url}")
            else:
                logger.debug(f"Using custom WebSocket URL: {url}")

            try:
                uri = (
                    url + f"?keepalive_timeout_seconds={self.keepalive_timeout_seconds}"
                )
                logger.info(f"Connecting to WebSocket for shard {shard_id} at {uri}")

                async with websockets.connect(uri) as websocket:
                    logger.info(
                        f"Successfully connected to WebSocket for shard {shard_id}"
                    )
                    self.connections[shard_id] = websocket

                    # Start handling messages
                    logger.debug(f"Beginning message handling for shard {shard_id}")
                    await self.handle_messages(shard_id, websocket)

                    logger.info(
                        f"WebSocket connection closed normally for shard {shard_id}"
                    )
            except websockets.exceptions.WebSocketException as e:
                logger.error(f"WebSocket protocol error for shard {shard_id}: {e}")
                if shard_id in self.connection_closed_events:
                    logger.debug(
                        f"Setting connection closed event for shard {shard_id}"
                    )
                    self.connection_closed_events[shard_id].set()
                await self.stop_shard(shard_id)
            except Exception as e:
                logger.error(
                    f"Unexpected error in WebSocket connection for shard {shard_id}: {e}"
                )
                if shard_id in self.connection_closed_events:
                    logger.debug(
                        f"Setting connection closed event for shard {shard_id}"
                    )
                    self.connection_closed_events[shard_id].set()
                await self.stop_shard(shard_id)

    async def handle_messages(
        self, shard_id: str, websocket: websockets.ClientConnection
    ):
        """
        Listen and handle incoming messages for a specific shard.

        Args:
            shard_id (str): The shard identifier.
            websocket (websockets.WebSocketClientProtocol): The WebSocket connection.
        """
        logger.info(f"Listening for messages on shard {shard_id}")
        async for message in websocket:
            logger.debug(f"Shard {shard_id} received message: {message}")
            try:
                data = json.loads(message)
                ws_message = WebSocketMessage(**data)
                asyncio.create_task(
                    self.process_message(shard_id, ws_message)
                ).add_done_callback(
                    lambda _: logger.debug(f"Shard {shard_id} message processed.")
                )
                # logger.debug(f"Shard {shard_id} message processed.")
            except json.JSONDecodeError as e:
                logger.error(f"JSON decode error on shard {shard_id}: {e}")
            except Exception as e:
                logger.error(f"Error processing message on shard {shard_id}: {e}")
        logger.debug(f"Shard {shard_id} WebSocket connection closed.")

    async def process_message(self, shard_id: str, message: WebSocketMessage):
        """
        Process a single WebSocket message.

        Args:
            shard_id (str): The shard identifier.
            message (WebSocketMessage): The parsed WebSocket message.
        """
        message_type = message.metadata.message_type

        logger.debug(f"Processing message type: {message_type} for shard {shard_id}")

        if message_type == "session_welcome":
            session_id = message.payload.session.id
            logger.info(
                f"Shard {shard_id} received welcome message with session ID: {session_id}"
            )
            self.session_ids[shard_id] = session_id

            if shard_id in self.welcome_events:
                logger.info(f"Setting welcome event for shard {shard_id}")
                self.welcome_events[shard_id].set()
                logger.debug(f"Welcome event set for shard {shard_id}")
            else:
                logger.warning(f"No welcome event found for shard {shard_id}")

            # Assign subscriptions to this shard
            # await self.assign_subscriptions_to_shard(shard_id)

        elif message_type == "session_keepalive":
            # Keepalive messages are frequent, so we use debug level
            logger.debug(f"Shard {shard_id} received keepalive message")
            # Handle keepalive if needed
            pass

        elif message_type == "notification":
            # Extract useful information from the notification
            subscription_type = message.payload.subscription.type
            event_data = message.payload.event

            logger.info(f"Shard {shard_id} received {subscription_type} notification")
            logger.debug(f"Notification event data: {event_data}")

            if event_data:
                logger.debug("Forwarding notification to event handler")
                await self.event_handler(message.payload)
            else:
                logger.warning("Received notification with empty event data")

        elif message_type == "session_reconnect":
            new_url = message.payload.session.reconnect_url
            logger.info(
                f"Shard {shard_id} received reconnect request with URL: {new_url}"
            )
            logger.debug(f"Initiating reconnection for shard {shard_id}")
            await self.reconnect_shard(shard_id, new_url)

        elif message_type == "revocation":
            reason = message.payload.subscription.status
            subscription_id = message.payload.subscription.id
            subscription_type = message.payload.subscription.type

            logger.warning(
                f"Shard {shard_id} subscription revoked: {subscription_type} (ID: {subscription_id}), Reason: {reason}"
            )
            # Optionally, handle revocation (e.g., remove subscription)

        elif message_type == "session_close":
            logger.info(f"Shard {shard_id} session closed by Twitch")
            logger.debug(
                f"Initiating reconnection for shard {shard_id} after session close"
            )
            await self.reconnect_shard(shard_id)

        else:
            logger.warning(
                f"Shard {shard_id} received unhandled message type: {message_type}"
            )
            logger.debug(f"Unhandled message content: {message.model_dump()}")

    async def assign_subscriptions_to_shard(self, shard_id: str):
        """
        Assign all subscriptions managed by TwitchApi to the specified shard.

        Args:
            shard_id (str): The shard identifier.
        """
        # Fetch all subscriptions and assign to this shard based on hashing or other logic
        subscriptions = await self.twitch_api.list_eventsub_subscriptions()
        for subscription in subscriptions:
            transport = subscription.transport
            if transport.get("method") != "websocket":
                continue  # Only handle WebSocket transports here
            # Assign to shard based on a hashing mechanism or round-robin
            # For simplicity, assign all to the connected shard
            try:
                await self.subscribe(shard_id, subscription)
            except Exception as e:
                logger.error(
                    f"Failed to subscribe shard {shard_id} to subscription {subscription.id}: {e}"
                )

    async def subscribe(self, shard_id: str, subscription: EventSubSubscription):
        """
        Subscribe to an EventSub subscription on a specific shard.

        Args:
            shard_id (str): The shard identifier.
            subscription (EventSubSubscription): The subscription to assign.
        """
        logger.info(f"Subscribing shard {shard_id} to subscription {subscription.id}")
        welcome_event = self.welcome_events.get(shard_id)
        if not welcome_event:
            logger.error(f"No welcome event found for shard {shard_id}")
            return
        try:
            # Wait for the event to be set (with a timeout to prevent indefinite waiting)
            await asyncio.wait_for(welcome_event.wait(), timeout=10)
        except asyncio.TimeoutError:
            logger.error(f"Timeout waiting for welcome message for shard {shard_id}")
            return
        websocket = self.connections.get(shard_id)
        session_id = self.session_ids.get(shard_id)
        if not websocket:
            logger.error(f"No active WebSocket connection for shard {shard_id}")
            return

        try:
            await self.twitch_api.create_eventsub_subscription(
                EventSubSubscriptionRequest(
                    type=subscription.type,
                    version=subscription.version,
                    condition=subscription.condition,
                    transport=EventSubTransport(
                        session_id=session_id, method="websocket"
                    ),
                )
            )
        except Exception as e:
            logger.error(
                f"Failed to subscribe shard {shard_id} to subscription {subscription.id}: {e}"
            )

    async def reconnect_shard(self, shard_id: str, new_url: Optional[str] = None):
        """
        Reconnect a specific shard, optionally using a new URL.

        Args:
            shard_id (str): The shard identifier.
            new_url (Optional[str], optional): New WebSocket URL to connect to. Defaults to None.
        """
        logger.info(f"Reconnecting shard {shard_id}...")
        # await asyncio.sleep(5)  # Backoff before reconnecting
        self.tasks[shard_id] = asyncio.create_task(
            self.connect_shard(shard_id, new_url)
        )

    async def start_shard(self, shard_id: str):
        """
        Start managing a specific shard.

        Args:
            session_id: the session ID to use for the shard
            shard_id (str): The shard identifier.
        """
        logger.info(f"Starting WebSocket shard {shard_id}")
        self.welcome_events[shard_id] = asyncio.Event()
        self.connection_closed_events[shard_id] = asyncio.Event()
        if shard_id in self.tasks:
            logger.warning(f"Shard {shard_id} is already managed.")
            logger.warning(f"Stopping existing task for shard {shard_id}")
            await self.stop_shard(shard_id)
        self.tasks[shard_id] = asyncio.create_task(self.connect_shard(shard_id))

    async def stop_shard(self, shard_id: str):
        """
        Stop managing a specific shard.

        Args:
            shard_id (str): The shard identifier.
        """
        task = self.tasks.get(shard_id)
        if task:
            try:
                task.cancel()
            except asyncio.CancelledError:
                logger.info(f"Shard {shard_id} task cancelled.")
        if shard_id in self.connections:
            await self.connections[shard_id].close()
            del self.connections[shard_id]
            logger.info(f"Shard {shard_id} connection closed.")

    async def stop_all(self):
        """
        Stop managing all shards and close all connections.
        """
        self.keep_running = False
        tasks = list(self.tasks.keys())
        for shard_id in tasks:
            await self.stop_shard(shard_id)
        logger.info("All WebSocket shards stopped.")


__all__ = ["WebSocketManager"]
