# ruff:file-ignore[undocumented-public-module, undocumented-public-function, docstring-missing-returns, implicit-namespace-package]
# Configuration file for the Sphinx documentation builder.
#
# For the full list of built-in configuration values, see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information

import inspect
import operator
import os
import re
import sys
import typing
from datetime import date
from pathlib import Path
from typing import Any

import jaxtyping
from docutils import nodes
from docutils.parsers.rst.states import Inliner
from sphinx.addnodes import pending_xref
from sphinx.application import Sphinx
from sphinx.environment import BuildEnvironment
from sphinx.ext.intersphinx import missing_reference

from differt import __version__

project = "DiffeRT"
copyright = f"2023-{date.today().year}, Jérome Eertmans"  # ruff:ignore[builtin-variable-shadowing, call-date-today]
author = "Jérome Eertmans"
version = __version__
git_ref = os.environ.get("READTHEDOCS_GIT_COMMIT_HASH", "main")
conf_dir = Path(__file__).absolute().parent
root_dir = conf_dir.parent.parent
sys.path.append(str(conf_dir / "_extensions"))

RTD = "READTHEDOCS" in os.environ

# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

extensions = [
    # Built-in
    "sphinx.ext.autodoc",
    "sphinx.ext.autosummary",
    "sphinx.ext.intersphinx",
    "sphinx.ext.githubpages",
    "sphinx.ext.linkcode",
    "sphinx.ext.mathjax",
    "sphinx.ext.napoleon",
    "sphinx.ext.todo",
    # Additional
    "matplotlib.sphinxext.plot_directive",
    "myst_nb",
    "sphinxcontrib.bibtex",
    "sphinxext.opengraph",
    "sphinx_autodoc_typehints",
    "sphinx_copybutton",
    "sphinx_design",
    "sphinx_plotly_directive",
    "sphinx_remove_toctrees",
    # Custom
    "itu_materials_table",
]

templates_path = ["_templates"]
exclude_patterns = []

suppress_warnings = ["config.cache", "mystnb.unknown_mime_type"]

add_module_names = False
add_function_parentheses = False

nitpicky = True
nitpick_ignore = (
    ("py:class", "Array"),
    ("py:class", "differt.plotting._utils._Dispatcher"),
    ("py:class", "differt.utils.TypeVarTuple"),
    ("py:class", "jax._src.typing.SupportsDType"),
    ("py:class", "ndarray"),  # From ArrayLike
    ("py:obj", "differt.geometry._paths._M"),
    ("py:obj", "differt.utils._T"),
    ("py:obj", "differt.geometry._utils._T"),
    ("py:obj", "None.ArrayType"),
    ("py:class", "setup.<locals>.ArrayType"),
    ("py:class", "equinox._module._better_abstract.AbstractVar"),
)

nitpick_ignore_regex = [
    (r"py:.*", r"differt\.rt\..*"),
    (r"py:.*", r"differt\.scene\..*"),
    (r"py:.*", r"differt_core\.rt\..*"),
    (r"py:.*", r"differt_core\.scene\..*"),
    (r"py:class", r".*Remove all items from D\.*"),
    (r"py:class", r".*a shallow copy of D.*"),
    (r"py:class", r".*providing a view on D's.*"),
]

linkcheck_ignore = ["https://doi.org/10.1002/2015RS005659"]
linkcheck_report_timeouts_as_broken = False  # Default value in Sphinx >= 8

# -- MathJax settings

mathjax3_config = {
    "loader": {"load": ["[tex]/boldsymbol"]},
    "tex": {"packages": {"[+]": ["boldsymbol"]}},
}

numfig = True

# -- Intersphinx mapping

intersphinx_mapping = {
    "jax": ("https://jax.readthedocs.io/en/latest/", None),
    "jaxtyping": ("https://docs.kidger.site/jaxtyping/", None),
    "matplotlib": ("https://matplotlib.org/stable/", None),
    "mitsuba": ("https://mitsuba.readthedocs.io/en/stable/", None),
    "numpy": ("https://numpy.org/doc/stable/", None),
    "optax": ("https://optax.readthedocs.io/en/latest/", None),
    "plotly": ("https://plotly.com/python-api-reference/", None),
    "python": ("https://docs.python.org/3/", None),
    "requests": ("https://requests.readthedocs.io/en/latest/", None),
    "scipy": ("https://docs.scipy.org/doc/scipy/", None),
    "sionna": ("https://nvlabs.github.io/sionna/", None),
    "vispy": ("https://vispy.org", None),
    "warp": ("https://nvidia.github.io/warp/stable/", None),
}

# -- OpenGraph settings

ogp_site_url = "https://eertmans.be/DiffeRT/"
ogp_use_first_image = True

# -- Sphinx autodoc typehints settings

always_document_param_types = False
always_use_bars_union = True
autodoc_preserve_defaults = True


def fix_format_jaxtyping_annotation(annotation: Any, _config: Any) -> str | None:
    if not (
        isinstance(annotation, type) and issubclass(annotation, jaxtyping.AbstractArray)
    ):
        return None

    dtype = annotation.dtype.__name__
    atype = annotation.array_type.__name__
    shape = annotation.dim_str

    return f':py:class:`{dtype}[{atype}, "{shape}"] <jaxtyping.{dtype}>`'


typehints_formatter = fix_format_jaxtyping_annotation

# -- MyST-nb settings
myst_url_schemes = {
    "http": None,
    "https": None,
    "mailto": None,
    "ftp": None,
    "wiki": "https://en.wikipedia.org/wiki/{{path}}#{{fragment}}",
    "doi": "https://doi.org/{{path}}",
    "gh-pr": {
        "url": "https://github.com/jeertmans/DiffeRT/pull/{{path}}#{{fragment}}",
        "title": "PR #{{path}}",
        "classes": ["github"],
    },
    "gh-issue": {
        "url": "https://github.com/jeertmans/DiffeRT/issues/{{path}}#{{fragment}}",
        "title": "Issue #{{path}}",
        "classes": ["github"],
    },
    "gh-user": {
        "url": "https://github.com/{{path}}",
        "title": "@{{path}}",
        "classes": ["github"],
    },
    "ext-gh-issue": {
        "url": "https://github.com/{{path}}/issues/{{fragment}}",
        "title": "Issue {{path}}#{{fragment}}",
        "classes": ["github"],
    },
}

myst_heading_anchors = 3

myst_enable_extensions = [
    "amsmath",
    "attrs_inline",
    "colon_fence",
    "dollarmath",
    "html_admonition",
]

nb_execution_mode = "off" if os.environ.get("NB_OFF") else "auto"
nb_execution_timeout = (
    600  # So cells can take a long time, especially when downloading sionna scenes
)
nb_merge_streams = True

# By default, MyST-nb chooses the Widget output instead of the 2D snapshot
# so we need to change priorities, because the widget cannot work if Python is
# not actively running.

nb_mime_priority_overrides = [
    ("*", "text/html", 0),
]

# We cannot access log files on RTD, so we print to stderr

nb_execution_show_tb = RTD

# -- Bibtex

bibtex_bibfiles = ["references.bib"]

# Patch for Plotly from https://github.com/spatialaudio/nbsphinx/issues/128#issuecomment-1158712159

html_js_files = [
    "https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.4/require.min.js",
]

# -- Matplotlib directive

plot_pre_code = """
import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
import numpy as np
"""
plot_include_source = True
plot_html_show_source_link = False
plot_html_show_formats = False

# -- Plotly directive

plotly_pre_code = """
import jax
import jax.numpy as jnp
import jax.random as jr
import numpy as np
import plotly
import plotly.express as px
import plotly.graph_objects as go
"""
plotly_include_source = True
plotly_html_show_source_link = False
plotly_html_show_formats = False

# -- Linkcode settings


def linkcode_resolve(domain: str, info: dict[str, str]) -> str | None:
    if domain != "py":
        return None

    if not info["module"]:
        return None
    if not info["fullname"]:
        return None
    if info["module"].split(".", 1)[0] not in {"differt", "differt_core"}:
        return None

    try:  # ruff:ignore[too-many-statements-in-try-clause]
        mod = sys.modules.get(info["module"])
        obj = operator.attrgetter(info["fullname"])(mod)
        if isinstance(obj, property):
            obj: Any = obj.fget
        obj = inspect.unwrap(obj)
        filename = inspect.getsourcefile(obj)
        source, lineno = inspect.getsourcelines(obj)
    except (AttributeError, TypeError):
        return None

    if filename is None:
        return None

    filename = os.path.relpath(filename, start=root_dir)
    lines = f"#L{lineno}-L{lineno + len(source) - 1}" if lineno else ""

    return f"https://github.com/jeertmans/DiffeRT/blob/{git_ref}/{filename}{lines}"


# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output


html_baseurl = os.environ.get("READTHEDOCS_CANONICAL_URL", "/")
html_theme = "sphinx_book_theme"
html_static_path = ["_static"]

html_theme_options = {
    "show_toc_level": 2,
    "repository_url": "https://github.com/jeertmans/DiffeRT",
    "repository_branch": git_ref,
    "path_to_docs": "docs/source",
    "use_edit_page_button": True,
    "use_source_button": True,
    "use_issues_button": True,
    "use_repository_button": True,
    "navigation_with_keys": False,
    "launch_buttons": {"colab_url": "https://colab.research.google.com"},
}

html_logo = "_static/logo_250px.png"
html_favicon = "_static/favicon.png"

autosummary_generate = True
napolean_use_rtype = False

# Patches


def fix_sionna_folder(_app: Sphinx, obj: Any, _bound_method: bool) -> None:
    """Rename the default folder to a more readeable name."""
    module = getattr(obj, "__module__", None)
    if module and module.rsplit(".", maxsplit=1)[-1] == "_sionna":
        sig = inspect.signature(obj)
        parameters = []

        for param_name, parameter in sig.parameters.items():
            if param_name == "folder":
                parameter = parameter.replace(default="<path-to-differt>/scene/scenes")  # ruff:ignore[redefined-loop-name]

            parameters.append(parameter)

        obj.__signature__ = sig.replace(parameters=parameters)


def fix_reference(
    app: Sphinx, env: BuildEnvironment, node: pending_xref, contnode: nodes.TextElement
) -> nodes.reference | None:
    """Fix some intersphinx references that are broken."""
    if node["refdomain"] == "py":
        target = node["reftarget"]
        if target.startswith("equinox"):  # Sphinx fails to find them in the inventory
            if target.endswith("Module"):
                uri = (
                    "https://docs.kidger.site/equinox/api/module/module/#equinox.Module"
                )
            elif target.endswith("tree_at"):
                uri = (
                    "https://docs.kidger.site/equinox/api/manipulation/#equinox.tree_at"
                )
            elif node["reftype"] == "mod":
                uri = "https://docs.kidger.site/equinox/"
            else:
                return None

            newnode = nodes.reference(
                "", "", internal=False, refuri=uri, reftitle="(in equinox)"
            )
            newnode.append(contnode)

            return newnode

        if target.startswith(("differt.rt.", "differt.scene.")):
            new_target = target.replace("differt.rt.", "differt.geometry.").replace(
                "differt.scene.", "differt.geometry."
            )
            node["reftarget"] = new_target
            py_domain = env.get_domain("py")
            return py_domain.resolve_xref(
                env,
                node.get("refdoc", env.docname),
                app.builder,
                node["reftype"],
                new_target,
                node,
                contnode,
            )

        if target.startswith("jaxtyping"):  # Sphinx fails to find them in the inventory
            if node["reftype"] == "class":
                if "Array" in target or "ndarray" in target or "PyTree" in target:
                    uri = "https://docs.kidger.site/jaxtyping/api/array/#array"
                else:
                    uri = "https://docs.kidger.site/jaxtyping/api/array/#dtype"
            elif node["reftype"] == "obj" and target == "jaxtyping.shape":
                uri = "https://docs.kidger.site/jaxtyping/api/array/#shape"
            elif node["reftype"] == "mod":
                uri = "https://docs.kidger.site/jaxtyping/"
            else:
                return None

            newnode = nodes.reference(
                "", "", internal=False, refuri=uri, reftitle="(in jaxtyping)"
            )
            newnode.append(contnode)

            return newnode
        if node["reftarget"] == "plotly.graph_objs._figure.Figure":
            node["reftarget"] = "plotly.graph_objects.Figure"
            return missing_reference(app, env, node, contnode)

    return None


# -- GitHub roles


_EXPLICIT_TITLE_RE = re.compile(r"^(.+?)\s*<(.+)>$", re.DOTALL)

_GITHUB_BASE_URL = "https://github.com"
_REPO = "jeertmans/DiffeRT"


def _make_gh_role(
    url_template: str,
    title_template: str,
) -> Any:
    """Create a GitHub Sphinx role with given URL and title templates."""

    def role(
        _name: str,
        rawtext: str,
        text: str,
        _lineno: int,
        _inliner: Inliner,
        _options: dict[str, Any] | None = None,
        _content: list[str] | None = None,
    ) -> tuple[list[nodes.Node], list[nodes.system_message]]:
        m = _EXPLICIT_TITLE_RE.match(text)
        if m:
            explicit_title, target = m.group(1), m.group(2)
        else:
            explicit_title, target = None, text

        if "#" in target:
            path, fragment = target.split("#", 1)
        else:
            path, fragment = target, ""

        url = url_template.format(
            base=_GITHUB_BASE_URL, repo=_REPO, path=path, fragment=fragment
        )
        fragment_placeholder = "{fragment}"  # ruff:ignore[missing-f-string-syntax]
        if fragment and fragment_placeholder not in url_template:
            url = f"{url}#{fragment}"
        title = explicit_title or title_template.format(path=path, fragment=fragment)

        node = nodes.reference(rawtext, title, refuri=url, classes=["github"])
        return [node], []

    return role


def setup(app: Sphinx) -> None:
    typing.GENERATING_DOCS = True  # type: ignore[ty:unresolved-attribute]

    # Patch to avoid expanding the ArrayLike union type, which takes a lot
    # of space and is less readable.

    class ArrayLike(jaxtyping.Array):
        pass

    jaxtyping.ArrayLike = ArrayLike  # type: ignore[ty:invalid-assignment]

    from typing import TypeVar  # ruff:ignore[import-outside-top-level]

    from differt.geometry import (  # ruff:ignore[import-outside-top-level]
        download_sionna_scenes,
    )

    class ArrayType(jaxtyping.Array):
        def __repr__(self) -> str:
            return "ArrayType"

    import differt.plugins._deepmimo_types  # ruff:ignore[import-outside-top-level]

    differt.plugins._deepmimo_types.ArrayType = TypeVar("ArrayType", bound=ArrayType)  # type: ignore[ty:invalid-assignment,ty:invalid-legacy-type-variable]  # ruff:ignore[private-member-access]

    download_sionna_scenes()  # Put this here so that download does not occur during notebooks execution

    app.connect("autodoc-before-process-signature", fix_sionna_folder)
    app.connect("missing-reference", fix_reference)

    app.add_role(
        "gh-pr",
        _make_gh_role(
            url_template="{base}/{repo}/pull/{path}",
            title_template="PR #{path}",
        ),
    )
    app.add_role(
        "gh-issue",
        _make_gh_role(
            url_template="{base}/{repo}/issues/{path}",
            title_template="Issue #{path}",
        ),
    )
    app.add_role(
        "gh-user",
        _make_gh_role(
            url_template="{base}/{path}",
            title_template="@{path}",
        ),
    )
    app.add_role(
        "ext-gh-issue",
        _make_gh_role(
            url_template="{base}/{path}/issues/{fragment}",
            title_template="Issue {path}#{fragment}",
        ),
    )
