import asyncio
from pyee.asyncio import AsyncIOEventEmitter
from datetime import datetime
import json
import os
import websockets
from gumband_python_sdk.constants import SocketEventType
from gumband_python_sdk.logging_handler import get_local_logger
from gumband_python_sdk.network_handler import NetworkHandler

GB_MUTE_HEARTBEAT = os.environ.get("GB_MUTE_HEARTBEAT", "false").lower() == "true"
local_logger = get_local_logger()

RECONNECT_DELAY = 5  # seconds


class WebSocketService(AsyncIOEventEmitter):
    def __init__(self, network_handler: NetworkHandler, exhibit_id: int):
        super().__init__()
        self.network_handler: NetworkHandler = network_handler
        self.base_url = self.network_handler.endpoint_manager.get_ws_endpoint()
        self.exhibit_id = exhibit_id
        self.wsClient = None
        self.last_ping = None
        self.event_handlers = {}

    async def connect(self):
        try:
            if self.wsClient:
                await self.wsClient.close()
                self.wsClient = None

            headers = {"Authorization": f"Bearer {self.network_handler.jwt}"}
            self.wsClient = await websockets.connect(uri=self.base_url, additional_headers=headers)
            await self._on_open()

            # Start the message and ping listener
            self.listener_task = asyncio.create_task(self._listen_for_messages(self.wsClient))
            self.health_check_task = asyncio.create_task(self._websocket_health_check())
            self.last_ping = int(datetime.now().timestamp() * 1000)
        except websockets.ConnectionClosed as e:
            local_logger.warning(f"WebSocket closed: {e.code} - {e.reason}")
            await self._on_close()
        except Exception as e:
            print(e)

    async def _on_open(self):
        local_logger.info("WebSocket connection opened")
        # Send handshake message
        handshake_message = {
            "action": SocketEventType.HANDSHAKE.value,
            "data": {"jwt": self.network_handler.jwt},
        }
        await self.wsClient.send(json.dumps(handshake_message))

    async def _listen_for_messages(self, websocket):
        try:
            async for message in websocket:
                await self._on_message(message)
        except websockets.ConnectionClosed as e:
            local_logger.warning(f"WebSocket closed: {e.code} - {e.reason}")
            await self._on_close()

    async def _on_message(self, message):
        try:
            info = json.loads(message)
            action = info.get("action")
            data = info.get("data")
            self.emit(action, action=action, data=data)

            if action == SocketEventType.HANDSHAKE_SUCCESSFUL.value:
                local_logger.info("Handshake successful")
                subscribe_message = {
                    "action": SocketEventType.SUBSCRIBE.value,
                    "data": {"exhibitId": self.exhibit_id},
                }
                await self.send(subscribe_message)
        except Exception as e:
            error_message = f"Error parsing WebSocket JSON: {str(e)}"
            local_logger.error(error_message)
            self.emit(SocketEventType.ERROR.value, error_message)

    async def _on_close(self):
        local_logger.warning("WebSocket connection closed")
        self.listener_task.cancel()
        self.health_check_task.cancel()

        local_logger.info(f"Attempting to reconnect in {RECONNECT_DELAY} seconds...")
        await asyncio.sleep(RECONNECT_DELAY)
        await self.connect()

    async def send(self, message: dict):
        if self.wsClient:
            await self.wsClient.send(json.dumps(message))
        else:
            local_logger.error("WebSocket is not connected.")

    async def _websocket_health_check(self):
        """
        Periodically checks for the latest ping from the server, and attempts a reconnection
        if this client hasn't received a ping for a while.
        """
        # while True:
        #     try:
        #         if int(datetime.now().timestamp() * 1000) - self.last_ping > SOCKETS["SDK_OFFLINE_MS"]:
        #             await self.wsClient.close()
        #             await self._on_close()
        #         await asyncio.sleep(5)
        #     except Exception as e:
        #         local_logger.error(f"Error during websocket health check: {e}")
        #         break
