from __future__ import annotations

import collections
import hashlib
import json
import random
import shutil
import warnings
from operator import itemgetter
from pathlib import Path
from typing import TYPE_CHECKING, Any

import jinja2
from docutils import nodes
from docutils.parsers.rst import Directive
from docutils.parsers.rst.directives import flag
from docutils.statemachine import StringList
from sphinx.util.nodes import nested_parse_with_titles

from altair.utils.execeval import eval_block
from tests.examples_arguments_syntax import iter_examples_arguments_syntax
from tests.examples_methods_syntax import iter_examples_methods_syntax

from .utils import (
    create_generic_image,
    create_thumbnail,
    get_docstring_and_rest,
    prev_this_next,
)

if TYPE_CHECKING:
    from docutils.nodes import Node


EXAMPLE_MODULE = "altair.examples"


GALLERY_TEMPLATE = jinja2.Template(
    """
.. This document is auto-generated by the altair-gallery extension. Do not modify directly.

.. _{{ gallery_ref }}:

{{ title }}
{% for char in title %}-{% endfor %}

This gallery contains a selection of examples of the plots Altair can create. Some may seem fairly complicated at first glance, but they are built by combining a simple set of declarative building blocks.

Many draw upon sample datasets compiled by the `Vega <https://vega.github.io/vega/>`_ project.

If you can't find the plots you are looking for here, make sure to check out the :ref:`altair-ecosystem` section, which has links to packages for making e.g. network diagrams and animations.

.. note::

    With the release of Altair 6, the documentation was updated to use
    ``from altair.datasets import data`` instead of ``from vega_datasets import data``.
    This change also introduced updated column names in some datasets (e.g., spaces
    instead of underscores).

{% for grouper, group in examples %}

.. _gallery-category-{{ grouper }}:

{{ grouper }}
{% for char in grouper %}~{% endfor %}

.. raw:: html

   <span class="gallery">
   {% for example in group %}
   <a class="imagegroup" href="{{ example.name }}.html">
   <span
        class="image" alt="{{ example.title }}"
{% if example['use_svg'] %}
        style="background-image: url(..{{ image_dir }}/{{ example.name }}-thumb.svg);"
{% else %}
        style="background-image: url(..{{ image_dir }}/{{ example.name }}-thumb.png);"
{% endif %}
    ></span>

     <span class="image-title">{{ example.title }}</span>
   </a>
   {% endfor %}
   </span>

   <div style='clear:both;'></div>

{% endfor %}


.. toctree::
   :maxdepth: 2
   :caption: Examples
   :hidden:

   Gallery <self>
   Tutorials <../case_studies/index>
"""
)

MINIGALLERY_TEMPLATE = jinja2.Template(
    """
.. raw:: html

    <div id="showcase">
      <div class="examples">
      {% for example in examples %}
      <a
        class="preview" href="{{ gallery_dir }}/{{ example.name }}.html"
{% if example['use_svg'] %}
        style="background-image: url(.{{ image_dir }}/{{ example.name }}-thumb.svg)"
{% else %}
        style="background-image: url(.{{ image_dir }}/{{ example.name }}-thumb.png)"
{% endif %}
      ></a>
      {% endfor %}
      </div>
    </div>
"""
)


EXAMPLE_TEMPLATE = jinja2.Template(
    """
:orphan:
:html_theme.sidebar_secondary.remove:

.. This document is auto-generated by the altair-gallery extension. Do not modify directly.

.. _gallery_{{ name }}:

{{ docstring }}

.. altair-plot::
    {% if code_below %}:remove-code:{% endif %}
    {% if strict %}:strict:{% endif %}

{{ code | indent(4) }}

.. tab-set::

    .. tab-item:: Method syntax
        :sync: method

        .. code:: python

{{ method_code | indent(12) }}

    .. tab-item:: Attribute syntax
        :sync: attribute

        .. code:: python

{{ code | indent(12) }}
"""
)


def save_example_pngs(
    examples: list[dict[str, Any]], image_dir: Path, make_thumbnails: bool = True
) -> None:
    """Save example pngs and (optionally) thumbnails."""
    encoding = "utf-8"

    # store hashes so that we know whether images need to be generated
    hash_file: Path = image_dir / "_image_hashes.json"

    if hash_file.exists():
        with hash_file.open(encoding=encoding) as f:
            hashes = json.load(f)
    else:
        hashes = {}

    for example in examples:
        name: str = example["name"]
        use_svg: bool = example["use_svg"]
        code = example["code"]

        filename = name + (".svg" if use_svg else ".png")
        image_file = image_dir / filename

        example_hash = hashlib.sha256(code.encode()).hexdigest()[:32]
        hashes_match = hashes.get(filename, "") == example_hash

        if hashes_match and image_file.exists():
            print(f"-> using cached {image_file!s}")
        else:
            # the file changed or the image file does not exist. Generate it.
            print(f"-> saving {image_file!s}")
            chart = eval_block(code, strict=True)
            try:
                chart.save(image_file)
                hashes[filename] = example_hash
            except ImportError:
                warnings.warn("Unable to save image: using generic image", stacklevel=1)
                create_generic_image(image_file)

            with hash_file.open("w", encoding=encoding) as f:
                json.dump(hashes, f)

        if make_thumbnails:
            params = example.get("galleryParameters", {})
            if use_svg:
                # Thumbnail for SVG is identical to original image
                shutil.copyfile(image_file, image_dir / f"{name}-thumb.svg")
            else:
                create_thumbnail(image_file, image_dir / f"{name}-thumb.png", **params)

    # Save hashes so we know whether we need to re-generate plots
    with hash_file.open("w", encoding=encoding) as f:
        json.dump(hashes, f)


def populate_examples(**kwds: Any) -> list[dict[str, Any]]:
    """Iterate through Altair examples and extract code."""
    examples = sorted(iter_examples_arguments_syntax(), key=itemgetter("name"))
    method_examples = {x["name"]: x for x in iter_examples_methods_syntax()}

    for example in examples:
        docstring, category, code, lineno = get_docstring_and_rest(example["filename"])
        if example["name"] in method_examples:
            _, _, method_code, _ = get_docstring_and_rest(
                method_examples[example["name"]]["filename"]
            )
        else:
            method_code = code
            code += (
                "# No channel encoding options are specified in this chart\n"
                "# so the code is the same as for the method-based syntax.\n"
            )
        example.update(kwds)
        if category is None:
            msg = f"The example {example['name']} is not assigned to a category"
            raise Exception(msg)
        example.update(
            {
                "docstring": docstring,
                "title": docstring.strip().split("\n")[0],
                "code": code,
                "method_code": method_code,
                "category": category.title(),
                "lineno": lineno,
            }
        )

    return examples


def _indices(x: str, /) -> list[int]:
    return [int(idx) for idx in x.split()]


class AltairMiniGalleryDirective(Directive):
    has_content = False

    option_spec = {
        "size": int,
        "names": str,
        "indices": _indices,
        "shuffle": flag,
        "seed": int,
        "titles": bool,
        "width": str,
    }

    def run(self) -> list[Node]:
        size = self.options.get("size", 15)
        names = [name.strip() for name in self.options.get("names", "").split(",")]
        indices = self.options.get("indices", [])
        shuffle = "shuffle" in self.options
        seed = self.options.get("seed", 42)
        titles = self.options.get("titles", False)
        width = self.options.get("width", None)

        env = self.state.document.settings.env
        app = env.app

        gallery_dir = app.builder.config.altair_gallery_dir

        examples = populate_examples()

        if names:
            if len(names) < size:
                msg = (
                    "altair-minigallery: if names are specified, "
                    "the list must be at least as long as size."
                )
                raise ValueError(msg)
            mapping = {example["name"]: example for example in examples}
            examples = [mapping[name] for name in names]
        else:
            if indices:
                examples = [examples[i] for i in indices]
            if shuffle:
                random.seed(seed)
                random.shuffle(examples)
            if size:
                examples = examples[:size]

        include = MINIGALLERY_TEMPLATE.render(
            image_dir="/_static",
            gallery_dir=gallery_dir,
            examples=examples,
            titles=titles,
            width=width,
        )

        # parse and return documentation
        result = StringList()
        for line in include.split("\n"):
            result.append(line, "<altair-minigallery>")
        node = nodes.paragraph()
        node.document = self.state.document
        nested_parse_with_titles(self.state, result, node)

        return node.children


def main(app) -> None:
    src_dir = Path(app.builder.srcdir)
    target_dir: Path = src_dir / Path(app.builder.config.altair_gallery_dir)
    image_dir: Path = src_dir / "_images"

    gallery_ref = app.builder.config.altair_gallery_ref
    gallery_title = app.builder.config.altair_gallery_title
    examples = populate_examples(gallery_ref=gallery_ref, code_below=True, strict=False)

    target_dir.mkdir(parents=True, exist_ok=True)
    image_dir.mkdir(exist_ok=True)

    examples = sorted(examples, key=itemgetter("title"))
    examples_toc = collections.OrderedDict(
        {
            "Simple Charts": [],
            "Bar Charts": [],
            "Line Charts": [],
            "Area Charts": [],
            "Circular Plots": [],
            "Scatter Plots": [],
            "Uncertainties And Trends": [],
            "Distributions": [],
            "Tables": [],
            "Maps": [],
            "Interactive Charts": [],
            "Advanced Calculations": [],
            "Case Studies": [],
        }
    )
    for d in examples:
        examples_toc[d["category"]].append(d)

    encoding = "utf-8"

    # Write the gallery index file
    fp = target_dir / "index.rst"
    fp.write_text(
        GALLERY_TEMPLATE.render(
            title=gallery_title,
            examples=examples_toc.items(),
            image_dir="/_static",
            gallery_ref=gallery_ref,
        ),
        encoding=encoding,
    )

    # save the images to file
    save_example_pngs(examples, image_dir)

    # Write the individual example files
    for prev_ex, example, next_ex in prev_this_next(examples):
        if prev_ex:
            example["prev_ref"] = "gallery_{name}".format(**prev_ex)
        if next_ex:
            example["next_ref"] = "gallery_{name}".format(**next_ex)
        fp = target_dir / "".join((example["name"], ".rst"))
        fp.write_text(EXAMPLE_TEMPLATE.render(example), encoding=encoding)


def setup(app) -> None:
    app.connect("builder-inited", main)
    app.add_css_file("altair-gallery.css")
    app.add_config_value("altair_gallery_dir", "gallery", "env")
    app.add_config_value("altair_gallery_ref", "example-gallery", "env")
    app.add_config_value("altair_gallery_title", "Example Gallery", "env")
    app.add_directive_to_domain("py", "altair-minigallery", AltairMiniGalleryDirective)
