import asyncio
import json
import uuid

from tornado.websocket import websocket_connect

from .exceptions import CommandTimeoutError
from .events import RESPONSE_EVENT


class WebSocket(object):
    def __init__(self, url, on_open_callback=None, on_message_callback=None):
        self.url = url
        self.on_open_callback = on_open_callback
        self.on_message_callback = on_message_callback

        self.ws = None
        self.request_futures = {}

        websocket_connect(url, callback=self._on_socket_open, on_message_callback=self._on_socket_message)

    def _on_socket_open(self, future):
        if future.exception() is None:
            self.ws = future.result()
        else:
            raise future.exception()

        if self.on_open_callback:
            self.on_open_callback()

    def _on_socket_message(self, message):
        if message is None:
            # This happens when the connection is closed
            return
        message = json.loads(message)
        event = message['headers']['event']

        # Command responses are handled separately
        if event == RESPONSE_EVENT:
            self._handle_command_response(message)
        elif self.on_message_callback:
            self.on_message_callback(message)

    def _handle_command_response(self, message):
        request_id = message['headers']['request']
        future = self.request_futures.pop(request_id, None)
        if future:
            future.set_result(message['body'])

    async def send_command(self, command, payload_body={}, return_data=False):
        request_id = str(uuid.uuid4())
        payload = {
            'headers': {
                'request_id': request_id,
                'command': command
            },
            'body': payload_body
        }

        future = asyncio.Future()
        self.request_futures[request_id] = future

        await self.ws.write_message(json.dumps(payload))
        try:
            response = await asyncio.wait_for(future, timeout=30)
        except asyncio.TimeoutError:
            raise CommandTimeoutError

        # Return only error by default
        if return_data:
            return response
        return response['error']

    def close(self):
        if self.ws is not None:
            self.ws.close()
