from .jsonutil import json_clean
from nbformat.v4 import output_from_msg
from typing import Dict, List, Any, Optional

from jupyter_client.client import KernelClient


class OutputWidget:
    """This class mimics a front end output widget"""
    def __init__(
            self,
            comm_id: str,
            state: Dict[str, Any],
            kernel_client: KernelClient,
            executor) -> None:

        self.comm_id: str = comm_id
        self.state: Dict[str, Any] = state
        self.kernel_client: KernelClient = kernel_client
        self.executor = executor
        self.topic: bytes = ('comm-%s' % self.comm_id).encode('ascii')
        self.outputs: List = self.state['outputs']
        self.clear_before_next_output: bool = False

    def clear_output(
            self,
            outs: List,
            msg: Dict,
            cell_index: int) -> None:

        self.parent_header = msg['parent_header']
        content = msg['content']
        if content.get('wait'):
            self.clear_before_next_output = True
        else:
            self.outputs = []
            # sync back the state to the kernel
            self.sync_state()
            if hasattr(self.executor, 'widget_state'):
                # sync the state to the nbconvert state as well, since that is used for testing
                self.executor.widget_state[self.comm_id]['outputs'] = self.outputs

    def sync_state(self) -> None:
        state = {'outputs': self.outputs}
        msg = {'method': 'update', 'state': state, 'buffer_paths': []}
        self.send(msg)

    def _publish_msg(
            self,
            msg_type: str,
            data: Optional[Dict] = None,
            metadata: Optional[Dict] = None,
            buffers: Optional[List] = None,
            **keys) -> None:
        """Helper for sending a comm message on IOPub"""
        data = {} if data is None else data
        metadata = {} if metadata is None else metadata
        content = json_clean(dict(data=data, comm_id=self.comm_id, **keys))
        msg = self.kernel_client.session.msg(msg_type, content=content, parent=self.parent_header,
                                             metadata=metadata)
        self.kernel_client.shell_channel.send(msg)

    def send(
            self,
            data: Optional[Dict] = None,
            metadata: Optional[Dict] = None,
            buffers: Optional[List] = None) -> None:

        self._publish_msg('comm_msg', data=data, metadata=metadata, buffers=buffers)

    def output(
            self,
            outs: List,
            msg: Dict,
            display_id: str,
            cell_index: int) -> None:

        if self.clear_before_next_output:
            self.outputs = []
            self.clear_before_next_output = False
        self.parent_header = msg['parent_header']
        output = output_from_msg(msg)

        if self.outputs:
            # try to coalesce/merge output text
            last_output = self.outputs[-1]
            if (last_output['output_type'] == 'stream'
                    and output['output_type'] == 'stream'
                    and last_output['name'] == output['name']):
                last_output['text'] += output['text']
            else:
                self.outputs.append(output)
        else:
            self.outputs.append(output)
        self.sync_state()
        if hasattr(self.executor, 'widget_state'):
            # sync the state to the nbconvert state as well, since that is used for testing
            self.executor.widget_state[self.comm_id]['outputs'] = self.outputs

    def set_state(self, state: Dict) -> None:
        if 'msg_id' in state:
            msg_id = state.get('msg_id')
            if msg_id:
                self.executor.register_output_hook(msg_id, self)
                self.msg_id = msg_id
            else:
                self.executor.remove_output_hook(self.msg_id, self)
                self.msg_id = msg_id

    def handle_msg(self, msg: Dict) -> None:
        content = msg['content']
        comm_id = content['comm_id']
        assert comm_id == self.comm_id
        data = content['data']
        if 'state' in data:
            self.set_state(data['state'])