import asyncio
from uuid import UUID

from tinydb import TinyDB

from settings import settings
from proto import Computer, DynamicComputerState

db = TinyDB(settings.database_path)
computers = db.table("computers")


class StateManager:
    def __init__(self):
        self.websockets = set()
        self.connected_computers = set()

        self.current_state = None
        self.update_state()

    def update_state(self):

        self.current_state = {"computers": []}

        for computer in computers.all():

            static = Computer.parse_obj(computer)

            dynamic = DynamicComputerState(
                is_online=static.uuid in self.connected_computers,
            )

            self.current_state["computers"].append(
                {"static": static.dict(), "dynamic": dynamic.dict()}
            )

    async def push_state(self, socket):
        try:
            await socket.send_json(self.current_state)
        except RuntimeError:
            print("dead socket?")
            self.websockets.remove(socket)

    async def on_connect(self, socket):
        self.websockets.add(socket)
        await self.push_state(socket)

    async def on_disconnect(self, socket):
        self.websockets.remove(socket)

    async def on_change(self):
        self.update_state()
        await asyncio.gather(*[self.push_state(socket) for socket in self.websockets])

    async def on_computer_register(self, computer: Computer):

        computers.insert(computer.dict())
        await self.on_change()

    async def on_computer_connect(self, computer_id: UUID):
        self.connected_computers.add(computer_id)
        await self.on_change()

    async def on_computer_disconnect(self, computer_id: UUID):
        self.connected_computers.remove(computer_id)
        await self.on_change()

    async def on_computer_message(self, computer_id: UUID, message):
        print(f"[on_computer_message] UUID: {computer_id} Message: {message!r}")