""" Web Socket Client """
import asyncio
import ssl
from asyncio import AbstractEventLoop
from typing import Callable, Optional, Awaitable

import aiohttp
from absl import logging
from yarl import URL


class WSClientSocketError(Exception):
    """ Basic Socket error """
    pass


class WSClient:
    """ WebSocket Client """
    def __init__(self, io_loop: AbstractEventLoop, scheme: str, host: str,
                 path: str, user: str, password: str,
                 on_data_handler: Callable, ca_file: Optional[str] = None):
        self.io_loop = io_loop
        self.scheme = scheme
        self.host = host
        self.path = path
        self.user = user
        self.password = password
        self.ca_file = ca_file
        self.main_job: Optional[Awaitable] = None

        self.on_data_handler = on_data_handler
        self.url = URL.build(scheme=self.scheme, host=self.host,
                             path=self.path, user=self.user,
                             password=self.password)
        # Create SSL Context
        self.ssl_context = ssl.create_default_context()
        if self.ca_file is not None:
            self.ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
            self.ssl_context.load_verify_locations(cafile=self.ca_file)

        self.is_running = False

    def start(self) -> None:
        """ Start client """

        def start_callback(o) -> None:
            """ start callback """
            o.is_running = True
            o.main_job = o.io_loop.create_task(o.run())

        self.io_loop.call_soon_threadsafe(start_callback, self)

    def stop(self) -> None:
        """ Stop client """

        async def finisher(obj: "WSClient") -> None:
            """ Finishes the execution and clears the objects params """
            if obj.main_job is not None:
                obj.main_job.cancel()
            obj.io_loop.stop()

        def stop_callback(obj: "WSClient") -> None:
            """ start callback """
            obj.is_running = False
            obj.io_loop.create_task(finisher(obj))

        self.io_loop.call_soon_threadsafe(stop_callback, self)

    async def run(self):
        """ start web socket connection """
        while self.is_running:
            ws = None
            session = aiohttp.ClientSession()
            # noinspection PyBroadException
            try:
                async with session.ws_connect(self.url,
                                              ssl=self.ssl_context) as ws:
                    async for msg in ws:
                        # noinspection PyUnresolvedReferences
                        if msg.type == aiohttp.WSMsgType.BINARY:
                            # noinspection PyUnresolvedReferences
                            self.on_data_handler(msg.data)
                        elif msg.type == aiohttp.WSMsgType.CONTINUATION:
                            pass
                        elif msg.type == aiohttp.WSMsgType.TEXT:
                            pass  # do nothing
                        elif msg.type == aiohttp.WSMsgType.PING:
                            # TODO(s1z): Check if it's handled on the low level
                            pass
                        elif msg.type == aiohttp.WSMsgType.PONG:
                            # TODO(s1z): Check if it's handled on the low level
                            pass
                        elif msg.type == aiohttp.WSMsgType.CLOSE:
                            raise WSClientSocketError("CLOSE message received")
                        elif msg.type == aiohttp.WSMsgType.CLOSING:
                            raise WSClientSocketError("Socket is CLOSING")
                        elif msg.type == aiohttp.WSMsgType.CLOSED:
                            raise WSClientSocketError("Socket is CLOSED")
                        elif msg.type == aiohttp.WSMsgType.ERROR:
                            raise WSClientSocketError("Socket ERROR")
                        await asyncio.sleep(0)
                await asyncio.sleep(0)
            except WSClientSocketError as ex:
                exc_info = (ex.__class__, ex, ex.__traceback__)
                logging.warning("WS error.", exc_info=exc_info)
            except Exception as ex:
                exc_info = (ex.__class__, ex, ex.__traceback__)
                logging.warning("Unknown Error received:", exc_info=exc_info)
            finally:
                if ws is not None:
                    await ws.close()
                if session is not None:
                    await session.close()
            await asyncio.sleep(0)
