websocket server

import asyncio
import json
from asyncio import Queue

import websockets

from aw.msg_enum import TOPIC_UDS, KEY_RESP
from aw.aw_upload.blf_upload import RepeatingTimer

from conn.message import Message
from conn.msg_handler import MsgHandler

from func.wcar_log import printi, printe, printd

from data.product_data.socket_msg_log import SocketMsgLogDb


class SocketServer(object):
    PORT = 28657
    CLIENT_UI = 'ui'
    CLIENT_UDS = 'uds'

    def __init__(self):
        self.running = False
        self.topic_hanlder: dict = {}
        self.send_queue = {}
        self.uds_queue = None
        self.loop = None
        self.web_client_dict: dict = {}
        self.web_client_max = 20
        self.status_check_timer = RepeatingTimer(interval=10, function=self.__check_socket_status)

    def start_socket_server(self):
        if self.running:
            printe('SocketServer, socket has already started!')
            return
        self.running = True
        printi('SocketServer, Start socket server')
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        self.loop.create_task(self.__server())
        self.status_check_timer.start()
        try:
            self.loop.run_forever()
        finally:
            self.loop.close()
            self.loop = None
        printi('SocketServer, Exit socket server loop')

    def stop_socket_server(self):
        printi('SocketServer, ready to stop socket server')
        self.running = False
        self.status_check_timer.cancel()
        tasks = asyncio.all_tasks(loop=self.loop)
        printi('SocketServer, All tasks size: {}'.format(len(tasks)))
        for t in tasks:
            printi('cancel task result: {}'.format(t.cancel()))
        printi('SocketServer, stop asyncio running event loop')
        self.loop.call_soon_threadsafe(self.loop.stop)
        printi('SocketServer, Finished stop event loop')

    def send_msg(self, message):
        if isinstance(message, Message):
            message = json.dumps(message.__dict__)
        printd("Send Socket Message: {}".format(message))
        for hands_ret in self.web_client_dict:
            if self.send_queue[hands_ret].qsize() > 100:
                self.send_queue[hands_ret].get_nowait()
            self.send_queue[hands_ret].put_nowait(message)

    def register_msg_handler(self, topic, handler):
        printi('SocketServer, register_msg_handler, topic: {}, handler: {}'.format(topic, handler))
        self.topic_hanlder.update({topic: handler})

    def is_uds_connected(self):
        return True if self.web_client_dict.get(self.CLIENT_UDS) else False

    async def __server(self):
        self.send_queue = {}
        await websockets.serve(self.__handle_conn, "localhost", SocketServer.PORT)
        printi('SocketServer, Run a server task')

    async def __handle_conn(self, websocket):
        printi('SocketServer, Start handle_conn')
        if len(self.web_client_dict) >= self.web_client_max:
            printi("there has already {len(self.web_client_max)} clients online, please reconnect later.")
            return
        hands_ret = await self.__server_hands(websocket)
        if not hands_ret:
            return
        printi('SocketServer, hand shake success, hands_ret: {}'.format(hands_ret))
        soc_client_data = self.web_client_dict.get(hands_ret)
        if soc_client_data and soc_client_data.tasks:
            for task in soc_client_data.tasks:
                task.cancel()
        await asyncio.sleep(0.1, loop=self.loop)
        printi('SocketServer, All tasks size: {}'.format(len(asyncio.all_tasks(loop=self.loop))))
        if self.CLIENT_UI in hands_ret:
            tasks = [asyncio.create_task(self.__dispatch(websocket)), asyncio.create_task(self.__send(websocket))]
        else:
            self.uds_queue = Queue()
            tasks = [asyncio.create_task(self.__uds_msg_handle(websocket))]
        soc_client_data = SocClientData(websocket, tasks)
        send_queue = Queue(maxsize=500)
        self.web_client_dict.update({hands_ret: soc_client_data})
        self.send_queue.update({hands_ret: send_queue})
        await websocket.send("connection success")
        await asyncio.gather(*tasks)
        # await asyncio.gather(*[self.__dispatch(websocket), self.__send(websocket)])

    async def __server_hands(self, websocket):
        async for msg in websocket:
            printd("SocketServer, recv server_hands msg: {}".format(msg))
            if isinstance(msg, str) and msg.startswith('Hello, I am '):
                printi("SocketServer, Websocket Connected")
                if msg.endswith(self.CLIENT_UI):
                    return (self.CLIENT_UI, websocket.remote_address)
                    #return self.CLIENT_UI
                elif msg.endswith(self.CLIENT_UDS):
                    return self.CLIENT_UDS
            printe('SocketServer, server_hands failed!!')
            return None

    async def __dispatch(self, websocket):
        try:
            async for msg in websocket:
                if msg != 'ping':
                    print('-----msg start --------------------------------')
                    printi(msg)
                    print('-----msg end----------------------------------')
                if msg == 'ping':
                    hands_ret = (self.CLIENT_UI, websocket.remote_address)
                    websocket_client = self.web_client_dict.get(hands_ret)
                    websocket_client.running_status = True
                    await websocket.send('pong')
                else:
                    printi("Get Socket Message: {}".format(msg))
                    message = json.loads(msg, object_hook=Message)
                    topic_handler: MsgHandler = self.topic_hanlder.get(message.topic)
                    if topic_handler:
                        SocketMsgLogDb().fwk_db_write_socket_msg(msg_dict=message.__dict__)
                        topic_handler.handler(message)
                    else:
                        printe('SocketServer, dispatch, can not find handler')
                        topic_handler = self.topic_hanlder.get("DEAULT")
                        topic_handler.handler(message)
        except Exception as e:
            printe("SocketServer, dispatch msg error: {}".format(e))

    async def __send(self, websocket):
        hands_ret = (self.CLIENT_UI, websocket.remote_address)
        while self.running:
            if self.send_queue[hands_ret].empty():
                await asyncio.sleep(0.1, loop=self.loop)
                continue
            msg = self.send_queue[hands_ret].get_nowait()
            try:
                await websocket.send(msg)
            except websockets.ConnectionClosed:
                websocket_client = self.web_client_dict.get(hands_ret)
                websocket_client.running_status = False
                printi(f"client {hands_ret} connection closed.")
            except Exception as ex:
                print(ex)
        printi('SocketServer, send_task, already stop server')

    def __check_socket_status(self):
        """
        定时器返回send_queue 状态
        :return:
        """
        for hands_ret in list(self.web_client_dict.keys()):
            websocket_client = self.web_client_dict.get(hands_ret)
            if not websocket_client.running_status:
                printi(f"remove clent {hands_ret}")
                for task in websocket_client.tasks:
                    task.cancel()
                self.send_queue.pop(hands_ret)
                self.web_client_dict.pop(hands_ret)

    def send_msg_for_uds_key(self, message):
        if isinstance(message, Message):
            message = json.dumps(message.__dict__)
        self.uds_queue.put_nowait(message)

    async def __uds_msg_handle(self, websocket):
        while self.running:
            if self.uds_queue.empty():

                await asyncio.sleep(1, loop=self.loop)
                continue
            msg = self.uds_queue.get_nowait()
            try:
                printi('SocketServer, Start send req msg')
                await websocket.send(msg)
                printi('SocketServer, send req msg finished')
                async for msg in websocket:
                    message = json.loads(msg, object_hook=Message)
                    printi('SocketServer, receive uds msg: {}'.format(msg))
                    if not (message.topic == TOPIC_UDS and message.name == KEY_RESP):
                        printe('SocketServer, wrong msg')
                        continue
                    topic_handler: MsgHandler = self.topic_hanlder.get(message.topic)
                    if topic_handler:
                        topic_handler.handler(message)
                    printi('SocketServer, finished unlock ecu, req_key_ret:')
                    break
            except Exception as e:
                printe("SocketServer, dispatch msg error: {}".format(e))
        printi('SocketServer, send_task, already stop server')


class SocClientData(object):

    def __init__(self, web_soc_client, tasks: list):
        self.web_soc_client = web_soc_client
        self.tasks = tasks
        self.running_status = True