# -*- coding: utf-8 -*-
# Copyright (c) Vispy Development Team. All Rights Reserved.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.
"""Scraper for sphinx-gallery.

This is used to collect screenshots from the examples when executed via
sphinx-gallery. This can be included in any project wanting to take
advantage of this by adding the following to your sphinx ``conf.py``:

.. code-block:: python

    sphinx_gallery_conf = {
        ...
        'image_scrapers': ('vispy',)
    }

The scraper is provided to sphinx-gallery via the
``vispy._get_sg_image_scraper()`` function.

"""

from __future__ import annotations

import os
import time
import shutil
from vispy.io import imsave
from vispy.gloo.util import _screenshot
from vispy.scene import SceneCanvas

from sphinx_gallery.scrapers import optipng, figure_rst


class VisPyGalleryScraper:
    """Custom sphinx-gallery scraper to save the current Canvas to an image."""

    def __repr__(self):
        return self.__class__.__name__

    def __call__(self, block, block_vars, gallery_conf):
        """Scrape VisPy Canvases and applications.

        Parameters
        ----------
        block : tuple
            A tuple containing the (label, content, line_number) of the block.
        block_vars : dict
            Dict of block variables.
        gallery_conf : dict
            Contains the configuration of Sphinx-Gallery

        Returns
        -------
        rst : str
            The ReSTructuredText that will be rendered to HTML containing
            the images. This is often produced by
            :func:`sphinx_gallery.scrapers.figure_rst`.

        """
        example_fn = block_vars["src_file"]
        frame_num_list = self._get_frame_list_from_source(example_fn)
        image_path_iterator = block_vars['image_path_iterator']
        canvas_or_widget = get_canvaslike_from_globals(block_vars["example_globals"])
        if not frame_num_list:
            image_paths = []
        elif isinstance(frame_num_list[0], str):
            # example produces an image/animation as output
            image_paths = []
            for frame_image, image_path in zip(frame_num_list, image_path_iterator):
                image_path = os.path.splitext(image_path)[0] + os.path.splitext(frame_image)[1]
                shutil.move(frame_image, image_path)
                image_paths.append(image_path)
        else:
            image_paths = self._save_example_to_files(
                canvas_or_widget, frame_num_list, gallery_conf, image_path_iterator)

        fig_titles = ""  # alt text
        # FUTURE: Handle non-images (ex. MP4s) with raw HTML
        return figure_rst(image_paths, gallery_conf['src_dir'], fig_titles)

    def _save_example_to_files(self, canvas_or_widget, frame_num_list, gallery_conf, image_path_iterator):
        image_path = next(image_path_iterator)
        frame_grabber = FrameGrabber(canvas_or_widget, frame_num_list)
        frame_grabber.collect_frames()
        if len(frame_num_list) > 1:
            # let's make an animation
            # FUTURE: mp4 with imageio?
            image_path = os.path.splitext(image_path)[0] + ".gif"
            frame_grabber.save_animation(image_path)
        else:
            frame_grabber.save_frame(image_path)
        frame_grabber.cleanup()
        if 'images' in gallery_conf['compress_images']:
            optipng(image_path, gallery_conf['compress_images_args'])
        return [image_path]

    def _get_frame_list_from_source(self, filename):
        lines = open(filename, 'rb').read().decode('utf-8').splitlines()
        for line in lines[:10]:
            if not line.startswith("# vispy:"):
                continue
            if "gallery-exports" in line:
                _frames = line.split('gallery-exports')[1].split(',')[0].strip()
                frames = self._frame_exports_to_list(_frames)
                break
            if "gallery " in line:
                # Get what frames to grab
                _frames = line.split('gallery')[1].split(',')[0].strip()
                frames = self._frame_specifier_to_list(_frames)
                break
        else:
            # no frame number hint - don't grab any frames
            frames = []
        return frames

    def _frame_specifier_to_list(self, frame_specifier):
        _frames = frame_specifier or '0'
        frames = [int(i) for i in _frames.split(':')]
        if not frames:
            frames = [5]
        if len(frames) > 1:
            frames = list(range(*frames))
        return frames

    def _frame_exports_to_list(self, frame_specifier):
        frames = frame_specifier.split(" ")
        frame_paths = []
        for frame_fn in frames:
            # existing image file created by the example
            if not os.path.isfile(frame_fn):
                raise FileNotFoundError(
                    "Example gallery frame specifier must be a frame number, "
                    "frame range, or relative filename produced by the example.")
            frame_paths.append(frame_fn)
        return frame_paths


def get_canvaslike_from_globals(globals_dict):
    qt_widget = _get_qt_top_parent(globals_dict)
    if qt_widget is not None:
        return qt_widget

    # Get canvas
    if "canvas" in globals_dict:
        return globals_dict["canvas"]
    if "Canvas" in globals_dict:
        return globals_dict["Canvas"]()
    if "fig" in globals_dict:
        return globals_dict["fig"]
    return None


def _get_qt_top_parent(globals_dict):
    if "QWidget" not in globals_dict and "QMainWindow" not in globals_dict and "QtWidgets" not in globals_dict:
        return None

    qtwidgets = globals_dict.get("QtWidgets")
    qmainwindow = globals_dict.get("QMainWindow", getattr(qtwidgets, "QMainWindow", None))
    qwidget = globals_dict.get("QWidget", getattr(qtwidgets, "QWidget", qmainwindow))
    all_qt_widgets = [widget for widget in globals_dict.values()
                      if isinstance(widget, qwidget) and widget is not None]
    all_qt_mains = [widget for widget in all_qt_widgets if isinstance(widget, qmainwindow)]
    if all_qt_mains:
        return all_qt_mains[0]
    if all_qt_widgets:
        return all_qt_widgets[0]
    return None


class FrameGrabber:
    """Helper to grab a series of screenshots from the current Canvas-like object."""

    def __init__(self, canvas_obj, frame_grab_list: list[int]):
        self._canvas = canvas_obj
        self._done = False
        self._current_frame = -1
        self._collected_images = []
        self._frames_to_grab = frame_grab_list[:]  # copy so original list is preserved

    def cleanup(self):
        from PyQt5.QtWidgets import QApplication
        for child_widget in QApplication.allWidgets():
            if hasattr(child_widget, 'close'):
                child_widget.close()
        QApplication.processEvents()

    def on_draw(self, _):
        if self._done:
            return  # Grab only once
        self._current_frame += 1
        if self._current_frame in self._frames_to_grab:
            self._frames_to_grab.remove(self._current_frame)
            if isinstance(self._canvas, SceneCanvas):
                im = self._canvas.render(alpha=True)
            else:
                im = _screenshot()
            self._collected_images.append(im)
        if not self._frames_to_grab or self._current_frame > self._frames_to_grab[0]:
            self._done = True

    def collect_frames(self):
        """Show current Canvas and render and collect all frames requested."""
        if self._is_qt_widget():
            self._grab_qt_screenshot()
        else:
            self._grab_vispy_screenshots()

    def _is_qt_widget(self):
        try:
            from PyQt5.QtWidgets import QWidget
        except ImportError:
            return False
        return isinstance(self._canvas, QWidget)

    def _grab_qt_screenshot(self):
        from PyQt5.QtWidgets import QApplication
        from PyQt5.QtCore import QTimer
        self._canvas.show()
        # Qt is going to grab from the screen so we need the window on top
        self._canvas.raise_()
        # We need to give the GUI event loop and OS time to draw everything
        time.sleep(1.5)
        QApplication.processEvents()
        QTimer.singleShot(1000, self._grab_widget_screenshot)
        time.sleep(1.5)
        QApplication.processEvents()

    def _grab_widget_screenshot(self):
        from PyQt5.QtWidgets import QApplication
        screen = QApplication.screenAt(self._canvas.pos())
        screenshot = screen.grabWindow(int(self._canvas.windowHandle().winId()))
        arr = self._qpixmap_to_ndarray(screenshot)
        self._collected_images.append(arr)

    @staticmethod
    def _qpixmap_to_ndarray(pixmap):
        from PyQt5 import QtGui
        import numpy as np
        im = pixmap.toImage().convertToFormat(QtGui.QImage.Format.Format_RGB32)
        size = pixmap.size()
        width = size.width()
        height = size.height()
        im_bits = im.constBits()
        im_bits.setsize(height * width * 4)
        # Convert 0xffRRGGBB buffer -> (B, G, R, 0xff) -> (R, G, B)
        return np.array(im_bits).reshape((height, width, 4))[:, :, 2::-1]

    def _grab_vispy_screenshots(self):
        os.environ['VISPY_IGNORE_OLD_VERSION'] = 'true'
        self._canvas.events.draw.connect(self.on_draw, position='last')
        with self._canvas as c:
            self._collect_frames(c)

    def _collect_frames(self, canvas, limit=10000):
        n = 0
        while not self._done and n < limit:
            canvas.update()
            canvas.app.process_events()
            n += 1
        if n >= limit or len(self._frames_to_grab) > 0:
            raise RuntimeError("Could not collect any images")

    def save_frame(self, filename, frame_index=0):
        imsave(filename, self._collected_images[frame_index])

    def save_animation(self, filename):
        import imageio  # multiple gif not properly supported yet
        imageio.mimsave(filename, self._collected_images)