import asyncio

from pycrdt import Map

from traitlets import Unicode, Bool, Set
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)
    _pending_clear_output_cells = Set(default_value=set())

    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 get_jupyter_ydoc(self, file_id):
        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
        ydoc = await room.get_jupyter_ydoc()
        
        return ydoc

    async def _clear_ydoc_outputs(self, cell_id):
        """Clears the outputs of a cell in ydoc"""
        
        if not self._file_id:
            return
        
        notebook = await self.get_jupyter_ydoc(self._file_id)
        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 {self._file_id=}, {cell_index=}")

    async def clear_cell_outputs(self, cell_id):
        """Clears all outputs for a cell on disk and in ydoc."""

        if self._file_id is not None:
            await self._clear_ydoc_outputs(cell_id)
            self._pending_clear_output_cells.discard(cell_id)
            
            if self.use_outputs_service:
                self.outputs_manager.clear(file_id=self._file_id, cell_id=cell_id)
            

    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.
        """
        if msg_type == "clear_output":
            asyncio.create_task(self.clear_output_task(cell_id, content))
        else:
            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 clear_output_task(self, cell_id, content):
        """A courotine to handle clear_output messages"""

        wait = content.get("wait", False)
        if wait:
            self._pending_clear_output_cells.add(cell_id)
        else:
            await self.clear_cell_outputs(cell_id)

    async def output_task(self, msg_type, cell_id, content):
        """A coroutine to handle output messages."""

        # Check for pending clear_output before processing output
        if cell_id in self._pending_clear_output_cells:
            await self.clear_cell_outputs(cell_id)

        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=file_id,
                cell_id=cell_id,
                output=output,
                display_id=display_id
            )
        else:
            output = self.transform_output(msg_type, content, ydoc=True)

        notebook = await self.get_jupyter_ydoc(file_id)
        if not notebook:
            return
        
        # 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
