import asyncio

from pycrdt import Map

from traitlets import Unicode, Bool
from traitlets.config import LoggingConfigurable
from jupyter_server_documents.kernels.message_cache import KernelMessageCache

class OutputProcessor(LoggingConfigurable):
    
    _file_id = Unicode(default_value=None, allow_none=True)

    use_outputs_service = Bool(
        default_value=True,
        help="Should outputs be routed to the outputs service to minimize the in memory ydoc size."
    ).tag(config=True)

    @property
    def settings(self):
        """A shortcut for the Tornado web app settings."""
        #      self.KernelClient.KernelManager.AsyncMultiKernelManager.ServerApp
        return self.parent.parent.parent.parent.web_app.settings

    @property
    def kernel_client(self):
        """A shortcut to the kernel client this output processor is attached to."""
        return self.parent

    @property
    def outputs_manager(self):
        """A shortcut for the OutputsManager instance."""
        return self.settings["outputs_manager"]
    
    @property
    def session_manager(self):
        """A shortcut for the kernel session manager."""
        return self.settings["session_manager"]

    @property
    def file_id_manager(self):
        """A shortcut for the file id manager."""
        return self.settings["file_id_manager"]
    
    @property
    def yroom_manager(self):
        """A shortcut for the jupyter server ydoc manager."""
        return self.settings["yroom_manager"]

    async def clear_cell_outputs(self, cell_id, room_id):
        """Clears the outputs of a cell in ydoc"""
        room = self.yroom_manager.get_room(room_id)
        if room is None:
            self.log.error(f"YRoom not found: {room_id}")
            return
        notebook = await room.get_jupyter_ydoc()
        self.log.info(f"Notebook: {notebook}")
 
        cell_index, target_cell = notebook.find_cell(cell_id)
        if target_cell is not None:
            target_cell["outputs"].clear()
            self.log.info(f"Cleared outputs for ydoc: {room_id} {cell_index}")

    def clear(self, cell_id):
        """Clear all outputs for a given cell Id."""
        if self._file_id is not None:
            if self.use_outputs_service:
                room_id = f"json:notebook:{self._file_id}"
                asyncio.create_task(self.clear_cell_outputs(cell_id, room_id))
                self.outputs_manager.clear(file_id=self._file_id, cell_id=cell_id)

    # Outgoing messages
    def process_output(self, msg_type: str, cell_id: str, content: dict):
        """Process outgoing messages from the kernel.
        
        This returns the input dmsg if no the message should be sent to
        clients, or None if it should not be sent.

        The dmsg is a deserialized message generated by calling:

        > self.kernel_client.session.deserialize(dmsg, content=False)

        The content has not been deserialized yet as we need to verify we
        should process it.
        """
        asyncio.create_task(self.output_task(msg_type, cell_id, content))
        return None # Don't allow the original message to propagate to the frontend

    async def output_task(self, msg_type, cell_id, content):
        """A coroutine to handle output messages."""
        try:
            # TODO: The session manager may have multiple notebooks connected to the kernel
            # but currently get_session only returns the first. We need to fix this and
            # find the notebook with the right cell_id.
            kernel_session = await self.session_manager.get_session(
                kernel_id=self.parent.parent.kernel_id
            )
        except Exception as e:
            self.log.error(
                f"An exception occurred when processing output for cell {cell_id}"
            )
            self.log.exception(e)
            return
        else:
            path = kernel_session["path"]

        file_id = self.file_id_manager.get_id(path)
        if file_id is None:
            self.log.error(f"Could not find file_id for path: {path}")
            return
        self._file_id = file_id

        display_id = content.get("transient", {}).get("display_id")
        # Convert from the message spec to the nbformat output structure
        if self.use_outputs_service: 
            output = self.transform_output(msg_type, content, ydoc=False)
            output = self.outputs_manager.write(file_id, cell_id, output, display_id)
        else:
            output = self.transform_output(msg_type, content, ydoc=True)

        # Get the notebook ydoc from the room
        room_id = f"json:notebook:{file_id}"
        room = self.yroom_manager.get_room(room_id)
        if room is None:
            self.log.error(f"YRoom not found: {room_id}")
            return
        notebook = await room.get_jupyter_ydoc()
        self.log.info(f"Notebook: {notebook}")
 
        # Write the outputs to the ydoc cell.
        _, target_cell = notebook.find_cell(cell_id)
        if target_cell is not None and output is not None:
            output_index = self.outputs_manager.get_output_index(display_id) if display_id else None
            if output_index is not None:
                target_cell["outputs"][output_index] = output
            else:
                target_cell["outputs"].append(output)
            self.log.info(f"Wrote output to ydoc: {path} {cell_id} {output}")

    
    def transform_output(self, msg_type, content, ydoc=False):
        """Transform output from IOPub messages to the nbformat specification."""
        if ydoc:
            factory = Map
        else:
            factory = lambda x: x
        if msg_type == "stream":
            output = factory({
                "output_type": "stream",
                "text": content["text"],
                "name": content["name"]
            })
        elif msg_type == "display_data" or msg_type == "update_display_data":
            output = factory({
                "output_type": "display_data",
                "data": content["data"],
                "metadata": content["metadata"]
            })
        elif msg_type == "execute_result":
            output = factory({
                "output_type": "execute_result",
                "data": content["data"],
                "metadata": content["metadata"],
                "execution_count": content["execution_count"]
            })
        elif msg_type == "error":
            output = factory({
                "output_type": "error",
                "traceback": content["traceback"],
                "ename": content["ename"],
                "evalue": content["evalue"]
            })
        return output
