import asyncio
import json
from typing import Optional, List

from ..frame import Frame
from ..log import logger

try:
    import websockets
except ImportError:
    msg = 'Missing dependency websockets, to fix: pip install websockets'
    logger.fatal(msg)
    raise ImportError(msg)


class WebsocketConnection(object):
    def __init__(self, frame_handler=None) -> None:
        super().__init__()
        self._connection = None  # type: ignore
        self._connected = False
        self._endpoint = None  # type: Optional[str]
        self._spaces = set()  # type: set
        self._recv_task = None
        self._auth = None
        self._frame_handler = frame_handler
        logger.debug(f'{self.__class__.__name__} initialized.')

    async def _recv_loop(self):
        assert self._connected, 'Not connected!'
        try:
            async for message in self._connection:
                frame = Frame.from_dict(json.loads(message))
                await self._frame_handler(frame)
        except asyncio.streams.IncompleteReadError:
            logger.error(f'{self.__class__.__name__} disconnected.')

    async def connect(self, endpoint: str, auth: Optional[str] = None) -> None:
        assert endpoint, 'endpoint is required.'
        if self._connected:
            raise ConnectionError('Already connected.')
        if not endpoint.startswith('ws://'):
            raise ValueError(f'Expected endpoint to begin with "ws://".'
                             f'Got: {endpoint!r}')
        _endpoint = endpoint  # endpoint including token
        if not endpoint.endswith('/'):
            _endpoint = endpoint + '/'
        if auth:
            _endpoint = endpoint + f'{auth}/'
        else:
            logger.warning(f'{self.__class__.__name__} to {endpoint} has no password.')
        self._endpoint = endpoint
        self._auth = auth
        logger.debug(f'{self.__class__.__name__} connecting to {endpoint}.')
        self._connection = await websockets.connect(_endpoint)
        self._recv_task = asyncio.get_running_loop().create_task(self._recv_loop())
        self._connected = True
        logger.info(f'{self.__class__.__name__} connected to {endpoint}.')

    async def reconnect(self):
        logger.info(f'{self.__class__.__name__} reconnecting.')
        while True:
            if self._connected:
                logger.debug(f'{self.__class__.__name__} closing existing connection.')
                self.close()
                await asyncio.sleep(1)
            try:
                await self.connect(self._endpoint, self._auth)
                break
            except Exception as e:
                logger.error(f'{self.__class__.__name__} could not connect, trying again...')
                await asyncio.sleep(10)

    def close(self):
        if self._recv_task:
            logger.debug(f'{self.__class__.__name__} canceling recv loop.')
            self._recv_task.cancel()
        if self._connection:
            logger.debug(f'{self.__class__.__name__} closing connection.')
            self._connection.close()
        self._connected = False

    async def join(self, space: str) -> None:
        self._spaces.add(space)
        logger.debug(f'{self.__class__.__name__} joining space {space!r}.')

    async def leave(self, space: str) -> None:
        self._spaces.remove(space)
        logger.debug(f'{self.__class__.__name__} leaving space {space!r}.')

    def spaces(self):
        return [s for s in self._spaces]

    async def send(self, frame: Frame, space: str):
        if not self._connected:
            logger.warning(f'{self.__class__.__name__} skipping {frame!r} to {space!r} as not connected.')
            return
        # if space not in self._spaces:
        #     raise KeyError(f'Join the space {space!r} before sending.')
        frame.meta['space'] = space
        logger.debug(f'{self.__class__.__name__} sending {frame!r} to {space!r}.')
        try:
            await self._connection.send(json.dumps(frame.to_dict()))
        except Exception as e:
            logger.info(f'{self.__class__.__name__} disconnected?.')
            await self.reconnect()

    async def recv(self, wait=True) -> Frame:
        raise NotImplementedError()

    async def broadcast(self, frame: Frame, spaces: List[str] = None):
        # spaces = spaces or self._spaces
        # for space in set(spaces):
        await self.send(frame, '')

import asyncio
import json
from typing import Optional, List

from ..frame import Frame
from ..log import logger

try:
    import websockets
except ImportError:
    msg = 'Missing dependency websockets, to fix: pip install websockets'
    logger.fatal(msg)
    raise ImportError(msg)


class WebsocketConnection(object):
    def __init__(self, frame_handler=None) -> None:
        super().__init__()
        self._connection = None  # type: ignore
        self._connected = False
        self._endpoint = None  # type: Optional[str]
        self._spaces = set()  # type: set
        self._recv_task = None
        self._auth = None
        self._frame_handler = frame_handler
        logger.debug(f'{self.__class__.__name__} initialized.')

    async def _recv_loop(self):
        assert self._connected, 'Not connected!'
        try:
            async for message in self._connection:
                frame = Frame.from_dict(json.loads(message))
                await self._frame_handler(frame)
        except asyncio.streams.IncompleteReadError:
            logger.error(f'{self.__class__.__name__} disconnected.')

    async def connect(self, endpoint: str, auth: Optional[str] = None) -> None:
        assert endpoint, 'endpoint is required.'
        if self._connected:
            raise ConnectionError('Already connected.')
        if not endpoint.startswith('ws://'):
            raise ValueError(f'Expected endpoint to begin with "ws://".'
                             f'Got: {endpoint!r}')
        _endpoint = endpoint  # endpoint including token
        if not endpoint.endswith('/'):
            _endpoint = endpoint + '/'
        if auth:
            _endpoint = endpoint + f'{auth}/'
        else:
            logger.warning(f'{self.__class__.__name__} to {endpoint} has no password.')
        self._endpoint = endpoint
        self._auth = auth
        logger.debug(f'{self.__class__.__name__} connecting to {endpoint}.')
        self._connection = await websockets.connect(_endpoint)
        self._recv_task = asyncio.get_running_loop().create_task(self._recv_loop())
        self._connected = True
        logger.info(f'{self.__class__.__name__} connected to {endpoint}.')

    async def reconnect(self):
        logger.info(f'{self.__class__.__name__} reconnecting.')
        while True:
            if self._connected:
                logger.debug(f'{self.__class__.__name__} closing existing connection.')
                self.close()
                await asyncio.sleep(1)
            try:
                await self.connect(self._endpoint, self._auth)
                break
            except Exception as e:
                logger.error(f'{self.__class__.__name__} could not connect, trying again...')
                await asyncio.sleep(10)

    def close(self):
        if self._recv_task:
            logger.debug(f'{self.__class__.__name__} canceling recv loop.')
            self._recv_task.cancel()
        if self._connection:
            logger.debug(f'{self.__class__.__name__} closing connection.')
            self._connection.close()
        self._connected = False

    async def join(self, space: str) -> None:
        self._spaces.add(space)
        logger.debug(f'{self.__class__.__name__} joining space {space!r}.')

    async def leave(self, space: str) -> None:
        self._spaces.remove(space)
        logger.debug(f'{self.__class__.__name__} leaving space {space!r}.')

    def spaces(self):
        return [s for s in self._spaces]

    async def send(self, frame: Frame, space: str):
        if not self._connected:
            logger.warning(f'{self.__class__.__name__} skipping {frame!r} to {space!r} as not connected.')
            return
        # if space not in self._spaces:
        #     raise KeyError(f'Join the space {space!r} before sending.')
        frame.meta['space'] = space
        logger.debug(f'{self.__class__.__name__} sending {frame!r} to {space!r}.')
        try:
            await self._connection.send(json.dumps(frame.to_dict()))
        except Exception as e:
            logger.info(f'{self.__class__.__name__} disconnected?.')
            await self.reconnect()

    async def recv(self, wait=True) -> Frame:
        raise NotImplementedError()

    async def broadcast(self, frame: Frame, spaces: List[str] = None):
        # spaces = spaces or self._spaces
        # for space in set(spaces):
        await self.send(frame, '')

# type: ignore # type: Optional[str] # type: set # endpoint including token # if space not in self._spaces: #     raise KeyError(f'Join the space {space!r} before sending.') # spaces = spaces or self._spaces # for space in set(spaces):