r"""
    Autodoc extension dealing with local type references and function signatures.
"""

from __future__ import annotations

from collections.abc import Callable, Mapping
from dataclasses import dataclass
import inspect
import re
from types import FunctionType, ModuleType
from typing import Any, Literal, Optional, Union
from sphinx.application import Sphinx

@dataclass(frozen=True)
class ParsedType:
    r""" Dataclass for a parsed type. """
    name: str
    args: Union[None, str, tuple[ParsedType, ...]] = None
    variadic: bool = False

    def crossref(self, globalns: Optional[Mapping[str, Any]] = None) -> str:
        r""" Generates Sphinx cross-reference link for the given type, using local names. """
        # pylint: disable = eval-used
        if globalns is None:
            globalns = {}
        name, args, variadic = self.name, self.args, self.variadic
        role = "obj"
        if name in globalns:
            obj = globalns[name]
            if isinstance(obj, ModuleType):
                role = "mod"
            elif isinstance(obj, property):
                role = "attr"
            elif isinstance(obj, type):
                role = "class"
            elif isinstance(obj, FunctionType):
                role = "func"
        name_crossref = f":{role}:`{name}`"
        if args is None:
            return name_crossref
        if isinstance(args, str):
            _args = eval(f"({args}, )")
            arg_crossrefs = ", ".join(f"``{repr(arg)}``" for arg in _args)
        else:
            arg_crossrefs = ", ".join((arg.crossref(globalns) for arg in args))
        if variadic:
            arg_crossrefs += ", ..."
        return fr"{name_crossref}\ [{arg_crossrefs}]"

def _find_closing_idx(s: str, c_open: str, c_close: str, idx_open: int = 0) -> int:
    r""" Finds the index where a bracketed/quoted range ends. """
    assert len(c_open) == 1, c_open
    assert len(c_close) == 1, c_close
    assert idx_open < len(s), (idx_open, len(s))
    assert s[idx_open] == c_open, (idx_open, s[idx_open])
    lvl = 1
    idx_close: Optional[int] = None
    for idx in range(idx_open+1, len(s)):
        if s[idx] == c_close:
            lvl -= 1
        elif s[idx] == c_open:
            lvl += 1
        if lvl == 0:
            idx_close = idx
            break
    if idx_close is None:
        error_msg = f"Unbalanced opening symbol found while searching for first outermost {c_open}...{c_close} in {repr(s)}."
        raise ValueError(error_msg)
    return idx_close

def _parse_type(annotation: str) -> tuple[Union[ParsedType, Literal["..."]], str]:
    r""" Parses an annotation into the first type appearing in it, together with the unparsed remainder of the annotation. """
    quote_close_idx = None
    if annotation.startswith("'"):
        quote_close_idx = _find_closing_idx(annotation, "'", "'", 0)
    elif annotation.startswith('"'):
        quote_close_idx = _find_closing_idx(annotation, '"', '"', 0)
    if quote_close_idx is not None:
        if quote_close_idx == 1:
            raise ValueError(f"Cannot parse empty forward reference at start of annotation: {annotation = }")
        annotation = annotation[1:quote_close_idx]+annotation[quote_close_idx+1:]
    try:
        b_open: Optional[int] = annotation.index("[")
    except ValueError:
        b_open = None
    try:
        c_idx = annotation.index(",")
    except ValueError:
        c_idx = None
    if b_open is not None and (c_idx is None or b_open < c_idx):
        b_close = _find_closing_idx(annotation, "[", "]", b_open)
        name = annotation[:b_open]
        args_str = annotation[b_open+1:b_close]
        res = annotation[b_close+1:]
        assert name, (name, args_str, res)
        if name.split(".")[-1] == "Literal":
            return ParsedType(name, args_str), res
        args: list[ParsedType] = []
        variadic = False
        while args_str:
            arg, _args_str = _parse_type(args_str)
            if isinstance(arg, ParsedType):
                args.append(arg)
            else:
                assert arg == "...", arg
                variadic = True
                if _args_str:
                    raise ValueError(f"Ellipsis argument encountered in parametric type, but not at the end of args lis: {annotation = }, {args_str = }")
            if not _args_str:
                break
            if not _args_str.startswith(", "):
                raise ValueError(f"Multiple type parameters must be separated by ', ': {annotation = }, {args_str = }, {arg = } {_args_str = }")
            args_str = _args_str[2:]
        return ParsedType(name, tuple(args), variadic), res
    if c_idx is not None:
        name = annotation[:c_idx]
        res = annotation[c_idx:]
        if name == "...":
            raise ValueError(f"Found ellipsis followed by comma: {annotation = }")
        return ParsedType(name), res
    if "]" in annotation:
        raise ValueError(f"Encountered closing bracket ']' without any opening bracket: {annotation = } ")
    name = annotation
    if name == "...":
        return "...", ""
    return ParsedType(name), ""

def parse_type(annotation: str) -> ParsedType:
    r""" Parses an annotation into a type. """
    parsed_type, residual_string = _parse_type(annotation)
    if residual_string:
        raise ValueError(f"Annotation was not entirely consumed by parsing: {annotation = }, {parsed_type = }, {residual_string = }")
    if not isinstance(parsed_type, ParsedType):
        raise ValueError(f"Cannot parse ellipsis on its own: {annotation = }")
    return parsed_type

def sigdoc(fun: FunctionType, lines: list[str]) -> None:
    r"""
        Returns doclines documenting the parameter and return type of the given function
    """
    # pylint: disable = too-many-branches
    doc = "\n".join(lines)
    lines.append("")
    # FIXME: if an :rtype: line already exists, remove it here and re-append it after all param type lines.
    globalns = fun.__globals__
    sig = inspect.signature(fun)
    for p in sig.parameters.values():
        annotation = p.annotation
        if annotation == p.empty:
            continue
        if not isinstance(annotation, str):
            print(f"WARNING! Found non-string annotation: {repr(annotation)}. Did you forget to import annotation from __future__?.")
            annotation = str(annotation)
        t = parse_type(annotation)
        tx = t.crossref(globalns)
        default = p.default if p.default != p.empty else None
        is_args = p.kind == p.VAR_POSITIONAL
        is_kwargs = p.kind == p.VAR_KEYWORD
        if is_args:
            extra_info = "variadic positional"
        elif is_kwargs:
            extra_info = "variadic keyword"
        elif default is not None:
            extra_info = f"default = ``{default}``"
        else:
            extra_info = None
        if extra_info is None:
            line = f":type {p.name}: {tx}"
        else:
            line = f":type {p.name}: {tx}; {extra_info}"
        if f":param {p.name}:" not in doc:
            lines.append(f":param {p.name}:")
        if f":type {p.name}:" not in doc:
            lines.append(line)
    if sig.return_annotation == sig.empty:
        return
    t = parse_type(sig.return_annotation)
    tx = t.crossref(globalns)
    line = f":rtype: {tx}"
    if ":rtype:" not in doc:
        lines.append(line)

def sigdoc_handler(app: Sphinx, what: str, fullname: str, obj: Any, options: Any, lines: list[str]) -> None:
    r"""
        Handler for Sphinx Autodoc's event
        `autodoc-process-docstring<https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html#event-autodoc-process-docstring>`_
        which replaces cross-references specified in terms of module globals with their fully qualified version.
    """
    # pylint: disable = too-many-arguments
    if what not in ("function", "method", "property"):
        return
    if what == "property":
        fun: FunctionType = obj.fget
    else:
        fun = obj
    sigdoc(fun, lines)

def simple_crossref_pattern(name: str) -> re.Pattern[str]:
    r"""
        Pattern for simple imports:

        .. code-block :: python

            f":{role}:`{name}`"        # e.g. ":class:`MyClass`"
            f":{role}:`~{name}`"       # e.g. ":class:`~MyClass`"
            f":{role}:`{name}{tail}`"  # e.g. ":attr:`MyClass.my_property.my_subproperty`"
            f":{role}:`~{name}{tail}`" # e.g. ":attr:`~MyClass.my_property.my_subproperty`"

    """
    return re.compile(rf":([a-z]+):`(~)?{name}(\.[\.a-zA-Z0-9_]+)?`")

def simple_crossref_repl(name: str, fullname: str) -> Callable[[re.Match[str]], str]:
    r"""
        Replacement function for the pattern generated by :func:`simple_crossref_pattern`:

        .. code-block :: python

            f":{role}:`~{fullname}`"                    # e.g. ":class:`~mymod.mysubmod.MyClass`"
            f":{role}:`~{fullname}`"                    # e.g. ":class:`~mymod.mysubmod.MyClass`"
            f":{role}:`{name}{tail}<{fullname}{tail}>`" # e.g. ":attr:`MyClass.my_property.my_subproperty<mymod.mysubmod.MyClass.my_property.my_subproperty>`"
            f":{role}:`~{fullname}{tail}`"              # e.g. ":attr:`~mymod.mysubmod.MyClass.my_property.my_subproperty`"

    """
    def repl(match: re.Match[str]) -> str:
        role = match[1]
        short = match[2] is not None
        tail = match[3]
        if tail is None:
            return f":{role}:`~{fullname}`"
        if short:
            return f":{role}:`~{fullname}{tail}`"
        return f":{role}:`{name}{tail}<{fullname}{tail}>`"
    return repl

def labelled_crossref_pattern(name: str) -> re.Pattern[str]:
    r"""
        Pattern for labelled imports:

        .. code-block :: python

            f":{role}:`{label}<{name}>`"       # e.g. ":class:`my class<MyClass>`"
            f":{role}:`{label}<{name}{tail}>`" # e.g. ":attr:`my_property<MyClass.my_property>`"

    """
    return re.compile(rf":([a-z]+):`([\.a-zA-Z0-9_]+)<{name}(\.[\.a-zA-Z0-9_]+)?>`")

def labelled_crossref_repl(name: str, fullname: str) -> Callable[[re.Match[str]], str]:
    r"""
        Replacement function for the pattern generated by :func:`labelled_crossref_pattern`:

        .. code-block :: python

            f":{role}:`{label}<{fullname}>`"       # e.g. ":class:`my class<mymod.mysubmod.MyClass>`"
            f":{role}:`{label}<{fullname}{tail}>`" # e.g. ":attr:`my_property<mymod.mysubmod.MyClass.my_property>`"

    """
    def repl(match: re.Match[str]) -> str:
        role = match[1]
        label = match[2]
        tail = match[3]
        if tail is None:
            return f":{role}:`{label}<{fullname}>`"
        return f":{role}:`{label}<{fullname}{tail}>`"
    return repl

_crossref_subs: list[tuple[Callable[[str], re.Pattern[str]],
                           Callable[[str, str], Callable[[re.Match[str]], str]]]] = [
    (simple_crossref_pattern, simple_crossref_repl),
    (labelled_crossref_pattern, labelled_crossref_repl),
]
r"""
    Substitution patterns and replacement functions for various kinds of cross-reference scenarios.
"""

def _get_module_by_name(modname: str) -> ModuleType:
    r"""
        Gathers a module object by name.
    """
    # pylint: disable = exec-used, eval-used
    exec(f"import {modname.split('.')[0]}")
    mod: ModuleType = eval(modname)
    if not isinstance(mod, ModuleType):
        return None
    return mod

def _get_obj_mod(app: Sphinx, what: str, fullname: str, obj: Any) -> Optional[ModuleType]:
    r"""
        Gathers the containing module for the given ``obj``.
    """
    autodoc_type_aliases = app.config.__dict__.get("autodoc_type_aliases")
    name = fullname.split(".")[-1]
    obj_mod: Optional[ModuleType]
    if autodoc_type_aliases is not None:
        if name in autodoc_type_aliases and fullname == autodoc_type_aliases[name]:
            modname = ".".join(fullname.split(".")[:-1])
            obj_mod = _get_module_by_name(modname)
            return obj_mod
    if what == "module":
        obj_mod = obj
    elif what in ("function", "class", "method"):
        obj_mod = inspect.getmodule(obj)
    elif what == "property":
        obj_mod = inspect.getmodule(obj.fget)
    elif what == "data":
        modname = ".".join(fullname.split(".")[:-1])
        obj_mod = _get_module_by_name(modname)
    elif what == "attribute":
        modname = ".".join(fullname.split(".")[:-2])
        obj_mod = _get_module_by_name(modname)
    else:
        print(f"WARNING! Encountered unexpected value for {what = } at {fullname = }")
        obj_mod = None
    return obj_mod

def _build_fullname_dict(app: Sphinx, fullname: str, obj_mod: Optional[ModuleType], ) -> dict[str, str]:
    r"""
        Builds a dictionary of substitutions from module global names to their fully qualified names,
        based on :func:`inspect.getmodule` and `autodoc_type_aliases` (if specified in the Sphinx app config).
    """
    autodoc_type_aliases = app.config.__dict__.get("autodoc_type_aliases")
    fullname_dict: dict[str, str] = {}
    if obj_mod is not None:
        globalns = obj_mod.__dict__
        for g_name, g_obj in globalns.items():
            if isinstance(g_obj, (FunctionType, type)):
                g_mod = inspect.getmodule(g_obj)
            elif isinstance(g_obj, ModuleType):
                g_mod = g_obj
            else:
                g_mod = inspect.getmodule(g_obj)
            if g_mod is None or g_mod == obj_mod:
                continue
            if g_name not in g_mod.__dict__:
                continue
            g_modname = g_mod.__name__
            fullname_dict[g_name] = f"{g_modname}.{g_name}"
    if autodoc_type_aliases is not None:
        for a_name, a_fullname in autodoc_type_aliases.items():
            if a_name not in fullname_dict:
                fullname_dict[a_name] = a_fullname
    return fullname_dict

def local_crossref_handler(app: Sphinx, what: str, fullname: str, obj: Any, options: Any, lines: list[str]) -> None:
    r"""
        Handler for Sphinx Autodoc's event
        `autodoc-process-docstring<https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html#event-autodoc-process-docstring>`_
        which replaces cross-references specified in terms of module globals with their fully qualified version.
    """
    # pylint: disable = too-many-arguments, too-many-locals
    obj_mod = _get_obj_mod(app, what, fullname, obj)
    fullname_dict = _build_fullname_dict(app, fullname, obj_mod)
    for sub_name, sub_fullname in fullname_dict.items():
        for idx, line in enumerate(lines):
            for pattern_fun, repl_fun in _crossref_subs:
                pattern = pattern_fun(sub_name)
                repl = repl_fun(sub_name, sub_fullname)
                line = re.sub(pattern, repl, line)
            lines[idx] = line

def setup(app: Sphinx) -> None:
    r"""
        Registers handlers for Sphinx Autodoc's event
        `autodoc-process-docstring<https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html#event-autodoc-process-docstring>`_
    """
    app.connect("autodoc-process-docstring", sigdoc_handler)
    app.connect("autodoc-process-docstring", local_crossref_handler)
