import base64
import struct
import tempfile
from contextlib import contextmanager
from io import BytesIO, StringIO
from typing import Any, BinaryIO, Dict, Generator, Optional, Tuple, Union

import matplotlib.pyplot as plt
from matplotlib.artist import Artist
from matplotlib.figure import Figure

from pybloqs.block.base import BaseBlock
from pybloqs.block.convenience import add_block_types
from pybloqs.html import append_to, parse
from pybloqs.static import Css, JScript
from pybloqs.util import cfg_to_css_string

try:
    import plotly.offline as po
    from plotly.graph_objs import Figure as PlotlyFigure

    _PLOTLY_AVAILABLE = True
except ImportError:
    _PLOTLY_AVAILABLE = False

try:
    try:
        # Bokeh >= 3.1
        from bokeh.resources import Resources

        JSResources = Resources
        CSSResources = Resources
    except ImportError:
        try:
            from bokeh.resources import CSSResources, JSResources
        except ImportError:
            from bokeh.templates import CSSResources, JSResources
    try:
        from bokeh.plotting.figure import Figure as BokehFigure
    except ImportError:
        from bokeh.plotting import figure as BokehFigure
    from bokeh.embed.standalone import components
    from bokeh.io import export_png

    _BOKEH_AVAILABLE = True
except ImportError:
    _BOKEH_AVAILABLE = False


_MIME_TYPES = {"png": "png", "svg": "svg+xml"}

_PLOT_FORMAT: str = "png"
_PLOT_MIME_TYPE = _MIME_TYPES[_PLOT_FORMAT]
_PLOT_DPI: int = 100


@contextmanager
def plot_format(plot_format: Optional[str] = None, dpi: Optional[int] = None) -> Generator[None, None, None]:
    """
    Temporarily set the plot formatting settings

    :param plot_format: The plot format (e.g 'png')
    :type plot_format:  str
    :param dpi:         The DPI of the plots
    :type dpi:          int
    """
    old = get_plot_format()
    set_plot_format(plot_format, dpi)
    yield
    set_plot_format(*old)


def get_plot_format() -> Tuple[str, int]:
    """
    Get the current plot format parameters

    :return: tuple of format and dpi
    """
    return _PLOT_FORMAT, _PLOT_DPI


def set_plot_format(plot_format: Optional[str] = None, plot_dpi: Optional[int] = None) -> None:
    """
    Overwrite the current plot format settings

    :param plot_format: The plot format (e.g. 'png')
    :type  plot_format: str
    :param plot_dpi:    The DPI of the plots
    :type  plot_dpi:    int
    """
    global _PLOT_FORMAT
    global _PLOT_MIME_TYPE
    global _PLOT_DPI
    if plot_format is not None:
        _PLOT_FORMAT = plot_format
        _PLOT_MIME_TYPE = _MIME_TYPES[plot_format]

    if plot_dpi is not None:
        _PLOT_DPI = plot_dpi


class ImgBlock(BaseBlock):
    def __init__(
        self,
        data: Union[BaseBlock, bytes],
        mime_type: str = "png",
        width: Optional[str] = None,
        height: Optional[str] = None,
        img_style: Optional[Dict] = None,
        **kwargs,
    ) -> None:
        """
        Create a block containing an image. The dimensions can be sniffed from GIF
        and PNG data, for other formats, the `width` and `height` parameters can be
        used to specify the size.

        :param data: Either a block to convert to an image, or raw image data.
        :mime_type: The mime type of the image (e.g. "png" or "svg+xml")
        :param width: Image width override as string, e.g. "50px". If unspecified, the actual
                      image width is used.
        :param width: Image height override as string, e.g. "50px". If unspecified, the actual
                      image height is used.
        :param kwargs: Optional styling arguments. The `style` keyword argument has special
                       meaning in that it allows styling to be grouped as one argument.
                       It is also useful in case a styling parameter name clashes with a standard
                       block parameter.
        """
        # Wrapping a block in an image will render it
        if hasattr(data, "_write_block"):
            img_file = data.save(fmt=mime_type)
            with open(img_file, "rb") as f:
                data = f.read()

        self._img_data = base64.b64encode(data)
        self._mime_type = mime_type

        if width is None and height is None:
            if mime_type.lower() == "png":
                if struct.unpack("ccc", data[1:4]) != (b"P", b"N", b"G"):
                    raise ValueError("Image type is not png and does not match mime type")
                x, y = struct.unpack(">ii", data[16:24])
            elif mime_type.lower() == "gif":
                x, y = struct.unpack("<HH", data[6:10])
            else:
                raise ValueError(f"Can't determine image dimensions for mime type {mime_type}")

            width, height = (f"{x}px", f"{y}px")

        if img_style is None:
            img_styles = {}
        else:
            img_styles = img_style

        if width is not None:
            img_styles["width"] = width
        if height is not None:
            img_styles["height"] = height

        self._img_styles = img_styles

        super().__init__(**kwargs)

    def _write_contents(self, container, *args, **kwargs) -> None:
        src = StringIO()
        mime = f"data:image/{self._mime_type};base64,"
        src.write(mime)
        src.write(self._img_data.decode())

        img = append_to(container, "img", src=src.getvalue())

        if len(self._img_styles) > 0:
            img["style"] = cfg_to_css_string(self._img_styles)

    @staticmethod
    def from_file(img_file: Union[str, BinaryIO], **kwargs) -> "ImgBlock":
        """
        Load an image block from a file.

        :param img_file: File path or file-like object.
        :param kwargs: Arguments to pass to the `ImgBlock` constructor.
        :return: ImgBlock instance.
        """
        close_file = False

        if isinstance(img_file, str):
            img_file = open(img_file, "rb")  # noqa: SIM115
            close_file = True

        try:
            return ImgBlock(img_file.read(), **kwargs)
        finally:
            # Close the file in case it was opened in this function
            if close_file:
                img_file.close()


class PlotBlock(ImgBlock):
    def __init__(
        self,
        plot,
        close_plot: bool = True,
        bbox_inches="tight",
        width: Optional[str] = None,
        height: Optional[str] = None,
        **kwargs,
    ) -> None:
        """
        Create a block containing a matplotlib figure

        :param plot: A matplotlib figure, axes or artist object.
        :close_plot: Optional (default=True). Set to True to close the plot after it is
                     captured into an image and avoid lingering plot windows.
        :bbox_inches: Optional bounding box parameter for 'figure.savefig'.
        :param kwargs: Optional styling arguments. The `style` keyword argument has special
                       meaning in that it allows styling to be grouped as one argument.
                       It is also useful in case a styling parameter name clashes with a standard
                       block parameter.
        """
        if not isinstance(plot, Artist):
            raise ValueError("PlotBlock contents must be matplotlib Artist")

        if isinstance(plot, Figure):
            figure = plot
        elif isinstance(plot, Artist):
            figure = plot.get_figure()
        else:
            raise ValueError("Unexpected plot object type %s", type(plot))

        img_data = BytesIO()

        legends = []

        for ax in figure.get_axes():
            legend = ax.get_legend()
            if legend is not None:
                # Patch Legend get_window_extent since there seems to be a bug where
                # it is passed an unexpected renderer instance.
                _orig_get_window_extent = legend.get_window_extent

                def _patched_get_window_extent(*_) -> Any:
                    return _orig_get_window_extent()

                legend.get_window_extent = _patched_get_window_extent
                legends.append(legend)

        if len(figure.axes) == 0:
            # empty plot, disable bbox_inches to that savefig still works
            bbox_inches = None

        figure.savefig(
            img_data, dpi=_PLOT_DPI, format=_PLOT_FORMAT, bbox_extra_artists=legends, bbox_inches=bbox_inches
        )

        plt_width, plt_height = figure.get_size_inches()

        width = width or f"{plt_width:0.3f}in"
        height = height or f"{plt_height:0.3f}in"

        if close_plot:
            plt.close(figure)

        super().__init__(img_data.getvalue(), _PLOT_MIME_TYPE, width=width, height=height, **kwargs)

    def _to_static(self) -> BaseBlock:
        # Convert to a basic image block in case we contain 'dynamic' svg content
        return ImgBlock(self) if self._mime_type == "svg" else super()._to_static()


class PlotlyPlotBlock(BaseBlock):
    def __init__(self, contents, plotly_kwargs: Optional[Dict[str, Any]] = None, static_kwargs=None, **kwargs) -> None:
        """
        Writes out the content as raw text or HTML.

        :param contents: Plotly graphics object figure.
        :param plotly_kwargs: Kwargs that are passed to plotly plot function.
        :param static_kwargs: Kwargs that are passed to plotly function write_image() for static output.
        :param kwargs: Optional styling arguments. The `style` keyword argument has special
                       meaning in that it allows styling to be grouped as one argument.
                       It is also useful in case a styling parameter name clashes with a standard
                       block parameter.
        """
        self.resource_deps = [JScript(script_string=po.offline.get_plotlyjs(), name="plotly")]

        super().__init__(**kwargs)

        if not isinstance(contents, PlotlyFigure):
            raise ValueError("Expected plotly.graph_objs.graph_objs.Figure type but got %s", type(contents))

        plotly_kwargs = plotly_kwargs or {}
        self.static_kwargs = static_kwargs or {}
        prefix = "<script>if (typeof require !== 'undefined') {var Plotly=require('plotly')}</script>"
        self._fig = contents
        self._contents = prefix + po.plot(contents, include_plotlyjs=False, output_type="div", **plotly_kwargs)

    def _write_contents(self, container, *args, **kwargs) -> None:
        container.append(parse(self._contents))

    def _repr_html_(self) -> str:
        return self.render_html()

    def _to_static(self) -> ImgBlock:
        # Create a static png image for use e.g. in an email body
        with tempfile.NamedTemporaryFile(suffix=".png") as f:
            self._fig.write_image(f.name, **self.static_kwargs)
            static_block = ImgBlock.from_file(f)
        return static_block


class BokehPlotBlock(BaseBlock):
    def __init__(self, contents, static_kwargs=None, **kwargs) -> None:
        """
        Writes out the content as raw text or HTML.

        :param contents: Bokeh plotting figure.
        :param static_kwargs: Kwargs that are passed to Bokeh function export_png() for static output.
        :param kwargs: Optional styling arguments. The `style` keyword argument has special
                       meaning in that it allows styling to be grouped as one argument.
                       It is also useful in case a styling parameter name clashes with a standard
                       block parameter.
        """
        self.resource_deps = [JScript(script_string=s, name="bokeh_js") for s in JSResources().js_raw]
        self.resource_deps += [Css(css_string=s, name="bokeh_css") for s in CSSResources().css_raw]

        super().__init__(**kwargs)

        if not isinstance(contents, BokehFigure):
            raise ValueError("Expected bokeh.plotting.figure.Figure type but got %s", type(contents))

        self._fig = contents
        self.static_kwargs = static_kwargs or {}
        script, div = components(contents)
        self._contents = script + div

    def _write_contents(self, container, *args, **kwargs) -> None:
        container.append(parse(self._contents))

    def _to_static(self) -> ImgBlock:
        # Create a static png image for use e.g. in an email body
        with tempfile.NamedTemporaryFile(suffix=".png") as f:
            export_png(self._fig, filename=f.name, **self.static_kwargs)
            static_block = ImgBlock.from_file(f)
        return static_block


add_block_types(Artist, PlotBlock)
# If Plotly or Bokeh are not installed skip registration
if _PLOTLY_AVAILABLE:
    add_block_types(PlotlyFigure, PlotlyPlotBlock)

if _BOKEH_AVAILABLE:
    add_block_types(BokehFigure, BokehPlotBlock)
