import websockets
import asyncio
import json
import logging

logging.basicConfig(level=logging.INFO)

class JsonWebsocketConnection(object):

    def __init__(self,
                 websocket: websockets.WebSocketServerProtocol):
        
        logging.info("incoming connection")
        self._websocket = websocket
        self._is_closed = False;
    
    def is_closed(self):
        return self._is_closed
    
    async def close(self):
        logging.info("closing connection")
        try:
            if self._is_closed:
                return
            self._websocket.close()
            self._is_closed = True
        except Exception as e:
            logging.warning("error closing connection: %s", str(e))
    
    async def send(self, message: dict):
        string_message = json.dumps(message)
        logging.debug("sending message: %s", string_message)
        try:
            await self._websocket.send(string_message)
        except Exception as e:
            logging.warning("error sending message: %s", str(e))
    
    async def handle_message(self, message: dict):
        pass # override this function

    async def run(self):
        try:
            async for message in self._websocket:
                try:
                    json_message = json.loads(message)
                    await self.handle_message(json_message) 
                except ValueError as e:
                    logging.warning("received unprocessable message %s", str(e))
                
        
        except Exception as e:
                logging.warning("error in websocket connection: %s", str(e))
                self._is_closed = True
        finally:
            self._is_closed = True
        
        



class JsonWebsocketServer(object):
    def __init__(self, handler_class: JsonWebsocketConnection, host:str = 'localhost', port:int = 8765, ssl_context = None):
        self._host = host
        self._port = port
        self._handler_class = handler_class
        self._ssl_context = ssl_context

    def run(self):
        async def main(websocket: websockets.WebSocketServerProtocol,
                       path: str):

            connection = self._handler_class(websocket)

            await connection.run()

        if (self._ssl_context is not None):
            start_server = websockets.serve(main, self._host, self._port, ssl=self._ssl_context)
        else:
            start_server = websockets.serve(main, self._host, self._port)

        asyncio.get_event_loop().run_until_complete(start_server)
        asyncio.get_event_loop().run_forever()