import multiprocessing

import slideflow as sf

from typing import Optional

from ._renderer import Renderer, CapturedException

#----------------------------------------------------------------------------

class AsyncRenderManager:

    """Manager to assist with rendering tile-level model predictions."""

    def __init__(self):

        self._closed        = False

        self._is_async      = False

        self._cur_args      = None

        self._cur_result    = None

        self._cur_stamp     = 0

        self._renderer_obj  = None

        self._args_queue    = None

        self._result_queue  = None

        self._process       = None

        self._model_path    = None

        self._live_updates  = False

        self.tile_px        = None

        self.extract_px     = None

        self._addl_render   = []

        self._set_device()

    def _set_device(self) -> None:

        """Set the device for the renderer."""

        if sf.util.torch_available:

            from slideflow.model import torch_utils

            self.device = torch_utils.get_device()

        else:

            self.device = None

    def close(self) -> None:

        """Close the renderer."""

        self._closed = True

        self._renderer_obj = None

        if self._process is not None:

            self._process.terminate()

        self._process = None

        self._args_queue = None

        self._result_queue = None

    @property

    def is_async(self) -> bool:

        """Return whether the renderer is in asynchronous mode.

        Returns:

            bool: Whether the renderer is in asynchronous mode.

        """

        return self._is_async

    def set_renderer(self, renderer_class: type, **kwargs) -> None:

        """Set the renderer class for the renderer.

        Args:

            renderer_class (type): Renderer class to use.

        """

        assert not self._closed

        if self.is_async:

            self._set_args_async(set_renderer=(renderer_class, kwargs))

        else:

            self._renderer_obj = renderer_class(device=self.device, **kwargs)

            for _renderer in self._addl_render:

                self._renderer_obj.add_renderer(_renderer)

    def close_renderer(self) -> None:

        if self.is_async:

            self._set_args_async(close_renderer=True)

        else:

            self._renderer_obj = None

    def add_to_render_pipeline(self, renderer: Renderer) -> None:

        """Add a renderer to the rendering pipeline.

        Args:

            renderer (Renderer): Renderer to add to the pipeline.

                This renderer will be triggered before the main renderer.

        Raises:

            ValueError: If the renderer is in asynchronous mode.

        """

        if self.is_async:

            raise ValueError("Cannot add to rendering pipeline when in "

                             "asynchronous mode.")

        self._addl_render += [renderer]

        if self._renderer_obj is not None:

            self._renderer_obj.add_renderer(renderer)

    def remove_from_render_pipeline(self, renderer: Renderer) -> None:

        """Remove a renderer from the rendering pipeline.

        Args:

            renderer (Renderer): Renderer to remove from the pipeline.

        Raises:

            ValueError: If the renderer is in asynchronous mode.

        """

        if self.is_async:

            raise ValueError("Cannot remove rendering pipeline when in "

                             "asynchronous mode.")

        idx = self._addl_render.index(renderer)

        del self._addl_render[idx]

        if self._renderer_obj is not None:

            self._renderer_obj.remove_renderer(renderer)

    def set_async(self, is_async):

        """Set the renderer to synchronous or asynchronous mode.

        Args:

            is_async (bool): Whether to set the renderer to asynchronous mode.

        """

        self._is_async = is_async

    def set_args(self, **args):

        """Set the arguments for the renderer."""

        assert not self._closed

        if args != self._cur_args or self._live_updates:

            if self._is_async:

                self._set_args_async(**args)

            else:

                self._set_args_sync(**args)

            if not self._live_updates:

                self._cur_args = args

    def _set_args_async(self, **args):

        """Set the arguments for the renderer in asynchronous mode."""

        if self._process is None:

            ctx = multiprocessing.get_context('spawn')

            self._args_queue = ctx.Queue()

            self._result_queue = ctx.Queue()

            self._process = ctx.Process(target=self._process_fn,

                                        args=(self._args_queue,

                                              self._result_queue,

                                              self._model_path,

                                              self._live_updates),

                                        daemon=True)

            self._process.start()

        self._args_queue.put([args, self._cur_stamp])

    def _set_args_sync(self, **args):

        """Set the arguments for the renderer in synchronous mode."""

        if self._renderer_obj is None:

            self._renderer_obj = Renderer(device=self.device)

            for _renderer in self._addl_render:

                self._renderer_obj.add_renderer(_renderer)

            self._renderer_obj._model = self._model

            self._renderer_obj._saliency = self._saliency

        self._cur_result = self._renderer_obj.render(**args)

    def get_result(self):

        """Get the result of the renderer.

        Returns:

            EasyDict: The result of the renderer.

        """

        assert not self._closed

        if self._result_queue is not None:

            while self._result_queue.qsize() > 0:

                result, stamp = self._result_queue.get()

                if stamp == self._cur_stamp:

                    self._cur_result = result

        return self._cur_result

    def clear_result(self):

        """Clear the result of the renderer."""

        assert not self._closed

        self._cur_args = None

        self._cur_result = None

        self._cur_stamp += 1

    def load_model(self, model_path: str) -> None:

        """Load a model for the renderer.

        Args:

            model_path (str): Path to the model.

        """

        if self._is_async:

            self._set_args_async(load_model=model_path)

        elif model_path != self._model_path:

            self._model_path = model_path

            if self._renderer_obj is None:

                self._renderer_obj = Renderer(device=self.device)

                for _renderer in self._addl_render:

                    self._renderer_obj.add_renderer(_renderer)

            self._renderer_obj.load_model(model_path, device=self.device)

    def clear_model(self):

        """Clear the model for the renderer."""

        self._model_path = None

        if self._renderer_obj is not None:

            self._renderer_obj._umap_encoders = None

            self._renderer_obj._model = None

            self._renderer_obj._saliency = None

    @property

    def _model(self):

        if self._renderer_obj is not None:

            return self._renderer_obj._model

        else:

            return None

    @property

    def _saliency(self):

        if self._renderer_obj is not None:

            return self._renderer_obj._saliency

        else:

            return None

    @property

    def _umap_encoders(self):

        if self._renderer_obj is not None:

            return self._renderer_obj._umap_encoders

        else:

            return None

    @staticmethod

    def _process_fn(

        args_queue: multiprocessing.Queue,

        result_queue: multiprocessing.Queue,

        model_path: Optional[str] = None,

        live_updates: bool = False

    ):

        if sf.util.torch_available:

            from slideflow.model import torch_utils

            device = torch_utils.get_device()

        else:

            device = None

        renderer_obj = Renderer(device=device)

        if model_path:

            renderer_obj.load_model(model_path, device=device)

        while True:

            while args_queue.qsize() > 0:

                args, stamp = args_queue.get()

                if 'close_renderer' in args:

                    renderer_obj = Renderer(device=device)

                if 'set_renderer' in args:

                    renderer_class, kwargs = args['set_renderer']

                    renderer_obj = renderer_class(**kwargs)

                if 'load_model' in args:

                    renderer_obj.load_model(args['load_model'], device=device)

                if 'quit' in args:

                    return

            if (live_updates and not result_queue.qsize()):

                result = renderer_obj.render(**args)

                if 'error' in result:

                    result.error = CapturedException(result.error)

                result_queue.put([result, stamp])