from atexit import register
import struct
import json
from time import time
import pydantic
import typing
import asyncio
import threading

from core.test_base import _assert, DEVICE_SOCKET_PORT


class LedStripDescription(pydantic.BaseModel):
    id: int
    name: str
    description: str
    location: str


class BoardDescription(pydantic.BaseModel):
    name: str
    description: str
    ledStrips: typing.List[LedStripDescription] = []
    sensors: typing.List = []


PAYLOAD_LENGTH_SIZE = 4


class Board(object):
    def __init__(self, description: BoardDescription):
        self.description = description
        self.thread = threading.Thread(
            target=lambda: asyncio.run(self.run()), daemon=True)
        self.queue = asyncio.Queue()
        self.on_register_response = None
        self.failure = None
        self.is_registered = False

    def set_on_register_response(self, callback):
        self.on_register_response = callback

    def get_description(self) -> BoardDescription:
        return self.description

    def expect_request(self, request_name: str, callback: typing.Callable = None):
        self.queue.put_nowait(
            (request_name, callback)
        )

    def start(self):
        self.thread.start()

    def wait_for_finish(self):
        self.thread.join()
        _assert(self.failure is None,
                f"Fake board client failed with error: {type(self.failure)}/{str(self.failure)}")

    # internals
    async def run(self):
        try:
            self.reader, self.writer = await asyncio.open_connection("127.0.0.1", DEVICE_SOCKET_PORT)

            register_request = {
                "name": "RegisterBoardRequest",
                "body": self.description.dict()
            }

            await self.send(register_request)
            response = await asyncio.wait_for(self.read(), timeout=1)

            if self.on_register_response:
                self.on_register_response(response["name"], response["body"])

            self.is_registered = True

            while not self.queue.empty():
                expected_request, callback = await self.queue.get()
                message = await asyncio.wait_for(self.read(), timeout=5)

                _assert(
                    expected_request == message["name"], message="Received unexpected request: " + message["name"])

                if callback is not None:
                    response = callback(message["body"])

                    if type(response) is dict:
                        if "conversation-id" in message:
                            response["conversation-id"] = message["conversation-id"]
                        await self.send(response)
                        await asyncio.sleep(0.1)

                await asyncio.sleep(0.1)

            self.writer.close()
            await self.writer.wait_closed()
        except Exception as e:
            self.failure = e

    async def send(self, message: dict):
        message = json.dumps(message)
        length = struct.pack("i", len(message))

        self.writer.write(length)
        self.writer.write(message.encode())

        await self.writer.drain()

    async def read(self) -> dict:
        length = await self.reader.read(PAYLOAD_LENGTH_SIZE)
        length = struct.unpack("i", length)

        payload = await self.reader.read(*length)
        return json.loads(payload.decode())
