from __future__ import annotations

import asyncio
import json
import logging
import socket
import urllib.parse
from enum import IntEnum
from typing import TYPE_CHECKING

from ably.sync.http.httputils import HttpUtils
from ably.sync.types.connectiondetails import ConnectionDetails
from ably.sync.util.eventemitter import EventEmitter
from ably.sync.util.exceptions import AblyException
from ably.sync.util.helper import Timer, unix_time_ms

try:
    # websockets 15+ preferred imports
    from websockets import ClientConnection as WebSocketClientProtocol
    from websockets import connect as ws_connect
except ImportError:
    # websockets 14 and earlier fallback
    from websockets.client import WebSocketClientProtocol
    from websockets.client import connect as ws_connect

from websockets.exceptions import ConnectionClosedOK, WebSocketException

if TYPE_CHECKING:
    from ably.sync.realtime.connection import ConnectionManager

log = logging.getLogger(__name__)


class ProtocolMessageAction(IntEnum):
    HEARTBEAT = 0
    CONNECTED = 4
    DISCONNECTED = 6
    CLOSE = 7
    CLOSED = 8
    ERROR = 9
    ATTACH = 10
    ATTACHED = 11
    DETACH = 12
    DETACHED = 13
    MESSAGE = 15
    AUTH = 17


class WebSocketTransport(EventEmitter):
    def __init__(self, connection_manager: ConnectionManager, host: str, params: dict):
        self.websocket: WebSocketClientProtocol | None = None
        self.read_loop: asyncio.Task | None = None
        self.connect_task: asyncio.Task | None = None
        self.ws_connect_task: asyncio.Task | None = None
        self.connection_manager = connection_manager
        self.options = self.connection_manager.options
        self.is_connected = False
        self.idle_timer = None
        self.last_activity = None
        self.max_idle_interval = None
        self.is_disposed = False
        self.host = host
        self.params = params
        super().__init__()

    def connect(self):
        headers = HttpUtils.default_headers()
        query_params = urllib.parse.urlencode(self.params)
        ws_url = (f'wss://{self.host}?{query_params}')
        log.info(f'connect(): attempting to connect to {ws_url}')
        self.ws_connect_task = asyncio.create_task(self.ws_connect(ws_url, headers))
        self.ws_connect_task.add_done_callback(self.on_ws_connect_done)

    def on_ws_connect_done(self, task: asyncio.Task):
        try:
            exception = task.exception()
        except asyncio.CancelledError as e:
            exception = e
        if exception is None or isinstance(exception, ConnectionClosedOK):
            return
        log.info(
            f'WebSocketTransport.on_ws_connect_done(): exception = {exception}'
        )

    def ws_connect(self, ws_url, headers):
        try:
            # Use additional_headers for websockets 15+, fallback to extra_headers for older versions
            try:
                with ws_connect(ws_url, additional_headers=headers) as websocket:
                    self._handle_websocket_connection(ws_url, websocket)
            except TypeError:
                # Fallback for websockets 14 and earlier
                with ws_connect(ws_url, extra_headers=headers) as websocket:
                    self._handle_websocket_connection(ws_url, websocket)
        except (WebSocketException, socket.gaierror) as e:
            exception = AblyException(f'Error opening websocket connection: {e}', 400, 40000)
            log.exception(f'WebSocketTransport.ws_connect(): Error opening websocket connection: {exception}')
            self._emit('failed', exception)
            raise exception from e

    def _handle_websocket_connection(self, ws_url, websocket):
        log.info(f'ws_connect(): connection established to {ws_url}')
        self._emit('connected')
        self.websocket = websocket
        self.read_loop = self.connection_manager.options.loop.create_task(self.ws_read_loop())
        self.read_loop.add_done_callback(self.on_read_loop_done)
        try:
            self.read_loop
        except WebSocketException as err:
            if not self.is_disposed:
                self.dispose()
                self.connection_manager.deactivate_transport(err)

    def on_protocol_message(self, msg):
        self.on_activity()
        log.debug(f'WebSocketTransport.on_protocol_message(): received protocol message: {msg}')
        action = msg.get('action')
        if action == ProtocolMessageAction.CONNECTED:
            connection_id = msg.get('connectionId')
            connection_details = ConnectionDetails.from_dict(msg.get('connectionDetails'))

            error = msg.get('error')
            exception = None
            if error:
                exception = AblyException.from_dict(error)

            max_idle_interval = connection_details.max_idle_interval
            if max_idle_interval:
                self.max_idle_interval = max_idle_interval + self.options.realtime_request_timeout
                self.on_activity()
            self.is_connected = True
            if self.host != self.options.get_realtime_host():  # RTN17e
                self.options.fallback_realtime_host = self.host
            self.connection_manager.on_connected(connection_details, connection_id, reason=exception)
        elif action == ProtocolMessageAction.DISCONNECTED:
            error = msg.get('error')
            exception = None
            if error is not None:
                exception = AblyException.from_dict(error)
            self.connection_manager.on_disconnected(exception)
        elif action == ProtocolMessageAction.AUTH:
            try:
                self.connection_manager.ably.auth.authorize()
            except Exception as exc:
                log.exception(f"WebSocketTransport.on_protocol_message(): An exception \
                                occurred during reauth: {exc}")
        elif action == ProtocolMessageAction.CLOSED:
            if self.ws_connect_task:
                self.ws_connect_task.cancel()
            self.connection_manager.on_closed()
        elif action == ProtocolMessageAction.ERROR:
            error = msg.get('error')
            exception = AblyException.from_dict(error)
            self.connection_manager.on_error(msg, exception)
        elif action == ProtocolMessageAction.HEARTBEAT:
            id = msg.get('id')
            self.connection_manager.on_heartbeat(id)
        elif action in (
            ProtocolMessageAction.ATTACHED,
            ProtocolMessageAction.DETACHED,
            ProtocolMessageAction.MESSAGE
        ):
            self.connection_manager.on_channel_message(msg)

    def ws_read_loop(self):
        if not self.websocket:
            raise AblyException('ws_read_loop started with no websocket', 500, 50000)
        try:
            for raw in self.websocket:
                msg = json.loads(raw)
                task = asyncio.create_task(self.on_protocol_message(msg))
                task.add_done_callback(self.on_protcol_message_handled)
        except ConnectionClosedOK:
            return

    def on_protcol_message_handled(self, task):
        try:
            exception = task.exception()
        except Exception as e:
            exception = e
        if exception is not None:
            log.exception(f"WebSocketTransport.on_protocol_message_handled(): uncaught exception: {exception}")

    def on_read_loop_done(self, task: asyncio.Task):
        try:
            exception = task.exception()
        except asyncio.CancelledError as e:
            exception = e
        if isinstance(exception, ConnectionClosedOK):
            return

    def dispose(self):
        self.is_disposed = True
        if self.read_loop:
            self.read_loop.cancel()
        if self.ws_connect_task:
            self.ws_connect_task.cancel()
        if self.idle_timer:
            self.idle_timer.cancel()
        if self.websocket:
            try:
                self.websocket.close()
            except asyncio.CancelledError:
                return

    def close(self):
        self.send({'action': ProtocolMessageAction.CLOSE})

    def send(self, message: dict):
        if self.websocket is None:
            raise Exception()
        raw_msg = json.dumps(message)
        log.info(f'WebSocketTransport.send(): sending {raw_msg}')
        self.websocket.send(raw_msg)

    def set_idle_timer(self, timeout: float):
        if not self.idle_timer:
            self.idle_timer = Timer(timeout, self.on_idle_timer_expire)

    def on_idle_timer_expire(self):
        self.idle_timer = None
        since_last = unix_time_ms() - self.last_activity
        time_remaining = self.max_idle_interval - since_last
        msg = f"No activity seen from realtime in {since_last} ms; assuming connection has dropped"
        if time_remaining <= 0:
            log.error(msg)
            self.disconnect(AblyException(msg, 408, 80003))
        else:
            self.set_idle_timer(time_remaining + 100)

    def on_activity(self):
        if not self.max_idle_interval:
            return
        self.last_activity = unix_time_ms()
        self.set_idle_timer(self.max_idle_interval + 100)

    def disconnect(self, reason=None):
        self.dispose()
        self.connection_manager.deactivate_transport(reason)
