"""Sphinx extension for intro_zsh example captions and the examples list.

- Moves code-block captions below the literal block.
- Expands ``.. intro-zsh-examples-list::`` into a linked list of all examples
  (labels use Sphinx ``numfig`` / ``toc_fignumbers``).

HTML and LaTeX render ``Example m.n:`` and chapter numbers natively.
Primary PDF is Sphinx ``latexpdf``; rst2pdf is a vanilla fallback only.

Edits are written back into the environment's doctree pickles during
``env-get-updated`` (after figure numbers are assigned) so later builders
see them.

Requires ``numfig = True`` and ``numfig_format['code-block']`` in conf.py.
"""

from __future__ import annotations

import os
import pickle
from os import path

from docutils import nodes
from sphinx import addnodes
from sphinx.application import Sphinx
from sphinx.environment import BuildEnvironment
from sphinx.util.docutils import SphinxDirective


class examples_list(nodes.General, nodes.Element):
    """Placeholder expanded in env-get-updated."""


class IntroZshExamplesList(SphinxDirective):
    has_content = False

    def run(self) -> list[nodes.Node]:
        return [examples_list("")]


def _is_intro_doc(docname: str) -> bool:
    return docname.startswith("intro_zsh/")


def _code_block_containers(doctree: nodes.document):
    for node in doctree.findall(nodes.container):
        if node.get("literal_block"):
            yield node


def _caption_of(container: nodes.container) -> nodes.caption | None:
    for child in container.children:
        if isinstance(child, nodes.caption):
            return child
    return None


def _move_caption_below(container: nodes.container) -> None:
    caption = _caption_of(container)
    if caption is None:
        return
    others = [c for c in container.children if c is not caption]
    container[:] = others + [caption]


def _format_example_prefix(app: Sphinx, numbers: tuple[int, ...]) -> str:
    fmt = app.config.numfig_format.get("code-block", "Example %s:")
    return (fmt % ".".join(map(str, numbers))).rstrip() + " "


def _strip_example_prefix(text: str) -> str:
    if text.startswith("Example ") and ":" in text:
        return text.split(":", 1)[1].strip()
    return text


def _intro_chapter_docnames(env: BuildEnvironment) -> list[str]:
    index = "intro_zsh/index"
    skip = {"intro_zsh/list_of_examples"}

    try:
        doctree = env.get_doctree(index)
    except Exception:
        return sorted(
            d for d in env.all_docs if _is_intro_doc(d) and d not in skip
        )

    ordered: list[str] = []
    for node in doctree.findall(addnodes.toctree):
        for _title, subdoc in node.get("entries", []):
            if not subdoc or subdoc == "self" or subdoc in skip:
                continue
            if _is_intro_doc(subdoc):
                ordered.append(subdoc)
    if ordered:
        return ordered
    return sorted(d for d in env.all_docs if _is_intro_doc(d) and d not in skip)


def _caption_title(container: nodes.container) -> str:
    caption = _caption_of(container)
    if caption is None:
        return ""
    return _strip_example_prefix(caption.astext())


def _make_example_ref(
    app: Sphinx, docname: str, fig_id: str, label: str
) -> nodes.reference:
    ref = nodes.reference("", label, internal=True)
    ref["refid"] = fig_id
    ref["refdocname"] = docname
    try:
        if app.builder.name not in ("singlehtml", "pdf", "latex"):
            uri = app.builder.get_target_uri(docname)
            ref["refuri"] = f"{uri}#{fig_id}"
    except Exception:
        pass
    return ref


def _build_examples_list(app: Sphinx) -> nodes.bullet_list:
    env = app.env
    blist = nodes.bullet_list()
    for docname in _intro_chapter_docnames(env):
        fignums = env.toc_fignumbers.get(docname, {}).get("code-block", {})
        if not fignums:
            continue
        try:
            doctree = env.get_doctree(docname)
        except Exception:
            continue
        for container in _code_block_containers(doctree):
            if not container.get("ids"):
                continue
            fig_id = container["ids"][0]
            numbers = fignums.get(fig_id)
            if not numbers:
                continue
            title = _caption_title(container)
            prefix = _format_example_prefix(app, numbers).rstrip()
            label = f"{prefix} {title}".strip() if title else prefix
            item = nodes.list_item()
            para = nodes.paragraph()
            para += _make_example_ref(app, docname, fig_id, label)
            item += para
            blist += item
    return blist


def _examples_list_replacement(app: Sphinx) -> nodes.Node:
    blist = _build_examples_list(app)
    if len(blist) == 0:
        return nodes.paragraph("", nodes.Text("(No examples found.)"))
    return blist


def _expand_examples_lists(app: Sphinx, doctree: nodes.document) -> None:
    for node in list(doctree.findall(examples_list)):
        node.replace_self(_examples_list_replacement(app))


def _rebuild_list_of_examples_page(app: Sphinx, doctree: nodes.document) -> None:
    """Always refresh the examples TOC from current chapter ids.

    After the first build the ``examples_list`` placeholder is gone (replaced by
    a bullet_list and pickled). Without a forced rebuild, renamed ``:name:``
    values leave dangling destinations.
    """
    replacement = _examples_list_replacement(app)

    placeholders = list(doctree.findall(examples_list))
    if placeholders:
        for node in placeholders:
            node.replace_self(replacement)
        return

    for section in doctree.findall(nodes.section):
        for child in list(section.children):
            if isinstance(child, nodes.bullet_list):
                child.replace_self(replacement)
                return
            if isinstance(child, nodes.paragraph) and child.astext().startswith(
                "(No examples"
            ):
                child.replace_self(replacement)
                return

    for section in doctree.findall(nodes.section):
        section += replacement
        return
    doctree += replacement


def _save_doctree(env: BuildEnvironment, docname: str, doctree: nodes.document) -> None:
    """Persist a mutated doctree so later get_doctree() / builders see it."""
    settings_env = getattr(doctree.settings, "env", None)
    reporter = getattr(doctree, "reporter", None)
    try:
        if hasattr(doctree.settings, "env"):
            doctree.settings.env = None
        doctree.reporter = None
        serialised = pickle.dumps(doctree, pickle.HIGHEST_PROTOCOL)
    finally:
        if settings_env is not None:
            doctree.settings.env = settings_env
        if reporter is not None:
            doctree.reporter = reporter

    env._pickled_doctree_cache[docname] = serialised
    filename = path.join(env.doctreedir, docname + ".doctree")
    os.makedirs(path.dirname(filename), exist_ok=True)
    with open(filename, "wb") as f:
        f.write(serialised)


def _process_intro_doctrees(app: Sphinx) -> None:
    env = app.env
    if not env.config.numfig:
        return

    chapters = [
        d for d in _intro_chapter_docnames(env) if d in env.all_docs
    ]
    list_doc = "intro_zsh/list_of_examples"
    ordered = chapters + ([list_doc] if list_doc in env.all_docs else [])

    for docname in ordered:
        try:
            doctree = env.get_doctree(docname)
        except Exception:
            continue
        for container in _code_block_containers(doctree):
            _move_caption_below(container)
        if docname == list_doc:
            _rebuild_list_of_examples_page(app, doctree)
        elif any(True for _ in doctree.findall(examples_list)):
            _expand_examples_lists(app, doctree)
        _save_doctree(env, docname, doctree)


def on_env_get_updated(app: Sphinx, env: BuildEnvironment) -> list[str]:
    # Runs after TocTreeCollector.assign_figure_numbers (default priority 500).
    _process_intro_doctrees(app)
    return []


def _visit_examples_list_skip(self, node):
    raise nodes.SkipNode


def setup(app: Sphinx):
    app.add_node(
        examples_list,
        html=(_visit_examples_list_skip, None),
        latex=(_visit_examples_list_skip, None),
        text=(_visit_examples_list_skip, None),
        man=(_visit_examples_list_skip, None),
        texinfo=(_visit_examples_list_skip, None),
    )
    app.add_directive("intro-zsh-examples-list", IntroZshExamplesList)
    app.connect("env-get-updated", on_env_get_updated, priority=900)
    return {
        "version": "1.5",
        "parallel_read_safe": True,
        "parallel_write_safe": True,
    }
