from i3ipc.aio import Connection
from i3ipc import Event
from .workspace_tree import WorkspaceTree
from .utils import async_debounce
import asyncio


class SwayConnection:

    last_output_set = None

    def __init__(self, workspace_tree: WorkspaceTree):
        self.connection = Connection(auto_reconnect=True)
        self.workspace_tree = workspace_tree
        self.on_change_callbacks = []

    async def on_workspace(self, connection=None, _=None):
        new_data = {ws.name: ws for ws in await self.connection.get_workspaces()}
        self.workspace_tree.update_workspaces(new_data)
        for callback in self.on_change_callbacks:
            await callback(self.workspace_tree)

    @async_debounce(0.1)
    async def on_output(self, connection, event):
        """On output event, update the workspace tree"""
        print("Output event received", flush=True)
        outputs = await connection.get_outputs()
        if self.last_output_set:
            if len(outputs) == len(self.last_output_set):
                # If the number of outputs is the same, we can assume that the outputs are the same
                return
        self.last_output_set = outputs
        await self.workspace_tree.update_context(connection)

    async def on_mode(self, event):
        """On mode change event, do something. Not sure what yet."""
        pass

    async def on_shutdown(self, event):
        """On shutdown event, close the connection and exit"""
        await self.connection.close()
        asyncio.get_event_loop().stop()

    async def run(self):
        # Connect to sway
        await self.connection.connect()

        # Register event listeners
        self.connection.on(Event.WORKSPACE, self.on_workspace)
        self.connection.on(Event.OUTPUT, self.on_output)
        self.connection.on(Event.MODE, self.on_mode)
        self.connection.on(Event.SHUTDOWN, self.on_shutdown)

        # Get initial output set
        self.last_output_set = await self.connection.get_outputs()

        # Update the workspace tree
        await self.workspace_tree.update_context(self.connection)
        await self.on_workspace()

        # Wait for events
        await self.connection.main()

    def on_change(self, callback):
        """Register a callback to be called when the workspace tree changes. The callback should be a coroutine, and should take the workspace tree as an argument."""
        self.on_change_callbacks.append(callback)

    async def close(self):
        await self.connection.close()