# Copyright (c) Meta Platforms, Inc. and affiliates.
# SPDX-License-Identifier: LGPL-2.1-or-later

import ast
from typing import Any, List, Optional, Pattern, Sequence, Tuple, cast

from drgndoc.namespace import BoundNode, Namespace, ResolvedNode
from drgndoc.parse import (
    Class,
    DocumentedNode,
    Function,
    FunctionSignature,
    Module,
    Variable,
)
from drgndoc.visitor import NodeVisitor


def _is_name_constant(node: ast.Constant) -> bool:
    return node.value is None or node.value is True or node.value is False


class _FormatVisitor(NodeVisitor):
    def __init__(
        self,
        namespace: Namespace,
        substitutions: Sequence[Tuple[Pattern[str], Any]],
        modules: Sequence[BoundNode[Module]],
        classes: Sequence[BoundNode[Class]],
        context_module: Optional[str],
        context_class: Optional[str],
    ) -> None:
        self._namespace = namespace
        self._substitutions = substitutions
        self._modules = modules
        self._classes = classes
        self._context_module = context_module
        self._context_class = context_class
        self._parts: List[str] = []

    def visit(  # type: ignore[override]  # This is intentionally incompatible with the supertype.
        self,
        node: ast.AST,
        *,
        rst: bool,
        qualify_typing: bool,
        qualify_typeshed: bool,
    ) -> str:
        self._rst = rst
        self._qualify_typing = qualify_typing
        self._qualify_typeshed = qualify_typeshed
        super().visit(node)
        ret = "".join(self._parts)
        self._parts.clear()
        return ret

    def generic_visit(self, node: ast.AST) -> None:
        raise NotImplementedError(
            f"{node.__class__.__name__} formatting is not implemented"
        )

    @staticmethod
    def _check_ctx_is_load(node: Any) -> None:
        if not isinstance(node.ctx, ast.Load):
            raise NotImplementedError(
                f"{node.ctx.__class__.__name__} formatting is not implemented"
            )

    def visit_Constant(
        self, node: ast.Constant, parent: Optional[ast.AST], sibling: Optional[ast.AST]
    ) -> None:
        if node.value is ...:
            self._parts.append("...")
        else:
            obj = self._rst and _is_name_constant(node)
            quote = self._rst and not isinstance(node.value, (int, float))
            if obj:
                self._parts.append(":py:obj:`")
            elif quote:
                self._parts.append("``")
            self._parts.append(repr(node.value))
            if obj:
                self._parts.append("`")
            elif quote:
                self._parts.append("``")

    def _append_resolved_name(self, name: str) -> None:
        if self._rst:
            self._parts.append(":py:obj:`")

        resolved = self._namespace.resolve_name_in_scope(
            self._modules, self._classes, name
        )
        if isinstance(resolved, ResolvedNode):
            target = resolved.qualified_name()
        else:
            target = resolved
        for pattern, repl in self._substitutions:
            target, num_subs = pattern.subn(repl, target)
            if num_subs:
                break

        title = target
        if not self._qualify_typing and title.startswith("typing."):
            title = title[len("typing.") :]
        elif not self._qualify_typeshed and title.startswith("_typeshed."):
            title = title[len("_typeshed.") :]
        elif self._context_module and title.startswith(self._context_module + "."):
            title = title[len(self._context_module) + 1 :]
            if self._context_class and title.startswith(self._context_class + "."):
                title = title[len(self._context_class) + 1 :]
        self._parts.append(title)

        if self._rst:
            if title != target:
                self._parts.append(" <")
                self._parts.append(target)
                self._parts.append(">")
            self._parts.append("`")

    def visit_Name(
        self, node: ast.Name, parent: Optional[ast.AST], sibling: Optional[ast.AST]
    ) -> None:
        self._check_ctx_is_load(node)
        self._append_resolved_name(node.id)

    def visit_Attribute(
        self, node: ast.Attribute, parent: Optional[ast.AST], sibling: Optional[ast.AST]
    ) -> None:
        self._check_ctx_is_load(node)
        name_stack = [node.attr]
        while True:
            value = node.value
            if isinstance(value, ast.Attribute):
                name_stack.append(value.attr)
                node = value
                continue
            elif isinstance(value, ast.Name):
                name_stack.append(value.id)
                name_stack.reverse()
                self._append_resolved_name(".".join(name_stack))
            elif isinstance(value, ast.Constant) and _is_name_constant(value):
                name_stack.append(repr(value.value))
                name_stack.reverse()
                self._append_resolved_name(".".join(name_stack))
            elif isinstance(value, ast.Constant) and not isinstance(
                value.value, (type(...), int, float)
            ):
                name_stack.append(repr(value.value))
                name_stack.reverse()
                if self._rst:
                    self._parts.append("``")
                self._parts.append(".".join(name_stack))
                if self._rst:
                    self._parts.append("``")
            else:
                self._visit(value, node, None)
                name_stack.append("")
                name_stack.reverse()
                if isinstance(value, ast.Constant) and isinstance(value.value, int):
                    # "1.foo()" is a syntax error without parentheses or an
                    # extra space.
                    self._parts.append(" ")
                elif self._rst:
                    # Make sure the "``" doesn't get squashed into a previous
                    # special character.
                    self._parts.append("\\ ")
                if self._rst:
                    self._parts.append("``")
                self._parts.append(".".join(name_stack))
                if self._rst:
                    self._parts.append("``")
            break

    def visit_Subscript(
        self, node: ast.Subscript, parent: Optional[ast.AST], sibling: Optional[ast.AST]
    ) -> None:
        self._check_ctx_is_load(node)
        self._visit(node.value, node, None)
        if self._rst:
            self._parts.append("\\")
        self._parts.append("[")
        self._visit(node.slice, node, None)
        if self._rst:
            self._parts.append("\\")
        self._parts.append("]")

    def visit_Tuple(
        self, node: ast.Tuple, parent: Optional[ast.AST], sibling: Optional[ast.AST]
    ) -> None:
        self._check_ctx_is_load(node)
        parens = (
            len(node.elts) == 0
            or not isinstance(parent, ast.Subscript)
            or node is not parent.slice
        )
        if parens:
            self._parts.append("(")
        for i, elt in enumerate(node.elts):
            if i > 0:
                self._parts.append(", ")
            self._visit(elt, node, node.elts[i + 1] if i < len(node.elts) - 1 else None)
        if len(node.elts) == 1:
            self._parts.append(",")
        if parens:
            self._parts.append(")")

    def visit_List(
        self, node: ast.List, parent: Optional[ast.AST], sibling: Optional[ast.AST]
    ) -> None:
        self._check_ctx_is_load(node)
        if self._rst:
            self._parts.append("\\")
        self._parts.append("[")
        for i, elt in enumerate(node.elts):
            if i > 0:
                self._parts.append(", ")
            self._visit(elt, node, node.elts[i + 1] if i < len(node.elts) - 1 else None)
        if self._rst:
            self._parts.append("\\")
        self._parts.append("]")

    def visit_UnaryOp(
        self, node: ast.UnaryOp, parent: Optional[ast.AST], sibling: Optional[ast.AST]
    ) -> None:
        if isinstance(node.op, ast.UAdd):
            self._parts.append("+")
        elif isinstance(node.op, ast.USub):
            self._parts.append("-")
        elif isinstance(node.op, ast.Not):
            self._parts.append("not ")
        elif isinstance(node.op, ast.Invert):
            self._parts.append("~")
        else:
            raise NotImplementedError(
                f"{node.op.__class__.__name__} formatting is not implemented"
            )
        parens = not isinstance(node.operand, (ast.Constant, ast.Name))
        if parens:
            self._parts.append("(")
        self._visit(node.operand, node, None)
        if parens:
            self._parts.append(")")


class Formatter:
    def __init__(
        self,
        namespace: Namespace,
        substitutions: Sequence[Tuple[Pattern[str], Any]] = (),
    ) -> None:
        self._namespace = namespace
        self._substitutions = substitutions

    def _format_function_signature(
        self,
        node: FunctionSignature,
        modules: Sequence[BoundNode[Module]],
        classes: Sequence[BoundNode[Class]],
        context_module: Optional[str],
        context_class: Optional[str],
        rst: bool,
        want_rtype: bool,
    ) -> Tuple[str, List[str]]:
        visitor = _FormatVisitor(
            self._namespace,
            self._substitutions,
            modules,
            classes,
            context_module,
            context_class,
        )
        assert node.docstring is not None
        lines = node.docstring.splitlines()
        if rst:
            lines = ["    " + line for line in lines]

        signature = ["("]
        need_comma = False

        def visit_arg(
            arg: ast.arg, default: Optional[ast.expr] = None, name: Optional[str] = None
        ) -> None:
            nonlocal need_comma
            if need_comma:
                signature.append(", ")
            signature.append(arg.arg if name is None else name)

            default_sep = "="
            if arg.annotation:
                signature.append(": ")
                signature.append(
                    visitor.visit(
                        arg.annotation,
                        rst=False,
                        qualify_typing=rst,
                        qualify_typeshed=False,
                    )
                )
                default_sep = " = "

            if default:
                signature.append(default_sep)
                signature.append(
                    visitor.visit(
                        default, rst=False, qualify_typing=True, qualify_typeshed=True
                    )
                )
            need_comma = True

        try:
            posargs = node.args.posonlyargs + node.args.args
            num_posonlyargs = len(node.args.posonlyargs)
        except AttributeError:
            posargs = node.args.args
            num_posonlyargs = 0

        # Type checkers treat parameters with names that begin but don't end
        # with __ as positional-only:
        # https://typing.readthedocs.io/en/latest/spec/historical.html#positional-only-parameters
        # We translate those to the PEP 570 syntax.
        def _is_posonly(arg: ast.arg) -> bool:
            return arg.arg.startswith("__") and not arg.arg.endswith("__")

        num_pep_570_posonlyargs = num_posonlyargs
        if (
            num_posonlyargs == 0
            and classes
            and not node.has_decorator("staticmethod")
            and len(posargs) > 1
            and _is_posonly(posargs[1])
        ):
            num_posonlyargs = 2
        while num_posonlyargs < len(posargs) and _is_posonly(posargs[num_posonlyargs]):
            num_posonlyargs += 1

        for i, arg in enumerate(posargs):
            default: Optional[ast.expr]
            if i >= len(posargs) - len(node.args.defaults):
                default = node.args.defaults[
                    i - (len(posargs) - len(node.args.defaults))
                ]
            else:
                default = None
            if i == 0 and classes and not node.has_decorator("staticmethod"):
                # Skip self for methods and cls for class methods.
                continue
            visit_arg(
                arg,
                default,
                name=(
                    arg.arg[2:]
                    if num_pep_570_posonlyargs <= i < num_posonlyargs
                    else arg.arg
                ),
            )
            if i == num_posonlyargs - 1:
                signature.append(", /")

        if node.args.vararg:
            visit_arg(node.args.vararg, name="*" + node.args.vararg.arg)

        if node.args.kwonlyargs:
            if not node.args.vararg:
                if need_comma:
                    signature.append(", ")
                signature.append("*")
                need_comma = True
            for i, arg in enumerate(node.args.kwonlyargs):
                visit_arg(arg, node.args.kw_defaults[i])

        if node.args.kwarg:
            visit_arg(node.args.kwarg, name="**" + node.args.kwarg.arg)

        signature.append(")")

        if want_rtype and node.returns:
            signature.append(" -> ")
            signature.append(
                visitor.visit(
                    node.returns, rst=False, qualify_typing=rst, qualify_typeshed=False
                )
            )

        return "".join(signature), lines

    def _format_class(
        self,
        resolved: ResolvedNode[Class],
        name: str,
        context_module: Optional[str] = None,
        context_class: Optional[str] = None,
        rst: bool = True,
    ) -> List[str]:
        node = resolved.node

        init_signatures: List[FunctionSignature] = []
        try:
            init = resolved.attr("__init__")
        except KeyError:
            pass
        else:
            if isinstance(init.node, Function):
                init_signatures = [
                    signature
                    for signature in init.node.signatures
                    if signature.docstring is not None
                ]

                init_context_class = resolved.name
                if context_class:
                    init_context_class = context_class + "." + init_context_class

        lines = []

        if rst and len(init_signatures) == 1 and node.docstring is None:
            class_signature, class_docstring_lines = self._format_function_signature(
                init_signatures[0],
                init.modules,
                init.classes,
                context_module,
                init_context_class,
                rst,
                False,
            )
            del init_signatures[0]
        else:
            class_signature = ""
            class_docstring_lines = (
                node.docstring.splitlines() if node.docstring else []
            )

        if rst:
            lines.append(f".. py:class:: {name}{class_signature}")

        if node.bases:
            visitor = _FormatVisitor(
                self._namespace,
                self._substitutions,
                resolved.modules,
                resolved.classes,
                context_module,
                context_class,
            )
            bases = [
                visitor.visit(
                    base, rst=rst, qualify_typing=False, qualify_typeshed=False
                )
                for base in node.bases
            ]
            if lines:
                lines.append("")
            lines.append(("    " if rst else "") + "Bases: " + ", ".join(bases))

        if class_docstring_lines:
            if lines:
                lines.append("")
            if rst:
                for line in class_docstring_lines:
                    lines.append("    " + line)
            else:
                lines.extend(class_docstring_lines)

        for i, signature_node in enumerate(init_signatures):
            if lines:
                lines.append("")

            signature, signature_lines = self._format_function_signature(
                signature_node,
                init.modules,
                init.classes,
                context_module,
                init_context_class,
                rst,
                False,
            )

            if rst:
                lines.append(f"    .. py:method:: {name}{signature}")
                lines.append("        :noindex:")
            elif signature:
                lines.append(f"{name}{signature}")
            lines.append("")
            if rst:
                for line in signature_lines:
                    lines.append("    " + line)
            else:
                lines.extend(signature_lines)
        return lines

    def _format_function(
        self,
        resolved: ResolvedNode[Function],
        name: str,
        context_module: Optional[str] = None,
        context_class: Optional[str] = None,
        rst: bool = True,
    ) -> List[str]:
        node = resolved.node

        lines = []
        for i, signature_node in enumerate(
            signature
            for signature in node.signatures
            if signature.docstring is not None
        ):
            if i > 0:
                lines.append("")
            signature, signature_lines = self._format_function_signature(
                signature_node,
                resolved.modules,
                resolved.classes,
                context_module,
                context_class,
                rst,
                True,
            )

            if rst:
                directive = "py:method" if resolved.classes else "py:function"
                lines.append(f".. {directive}:: {name}{signature}")
                if i > 0:
                    lines.append("    :noindex:")
                if node.async_:
                    lines.append("    :async:")
                if signature_node.has_decorator("classmethod") or name in (
                    "__init_subclass__",
                    "__class_getitem__",
                ):
                    lines.append("    :classmethod:")
                if signature_node.has_decorator("staticmethod"):
                    lines.append("    :staticmethod:")
            else:
                lines.append(f"{name}{signature}")
            if signature_lines:
                lines.append("")
                lines.extend(signature_lines)
        return lines

    def _format_variable(
        self,
        resolved: ResolvedNode[Variable],
        name: str,
        context_module: Optional[str],
        context_class: Optional[str],
        rst: bool,
    ) -> List[str]:
        node = resolved.node
        assert node.docstring is not None
        docstring_lines = node.docstring.splitlines()

        visitor = _FormatVisitor(
            self._namespace,
            self._substitutions,
            resolved.modules,
            resolved.classes,
            context_module,
            context_class,
        )
        if rst:
            directive = "py:attribute" if resolved.classes else "py:data"
            lines = [f".. {directive}:: {name}"]
            if node.annotation:
                lines.append(
                    "    :type: "
                    + visitor.visit(
                        node.annotation,
                        rst=False,
                        qualify_typing=True,
                        qualify_typeshed=False,
                    )
                )
            if docstring_lines:
                lines.append("")
            for line in docstring_lines:
                lines.append("    " + line)
            return lines
        else:
            if node.annotation:
                if docstring_lines:
                    docstring_lines.insert(0, "")
                docstring_lines.insert(
                    0,
                    visitor.visit(
                        node.annotation,
                        rst=False,
                        qualify_typing=False,
                        qualify_typeshed=False,
                    ),
                )
            return docstring_lines

    def format(
        self,
        resolved: ResolvedNode[DocumentedNode],
        name: Optional[str] = None,
        context_module: Optional[str] = None,
        context_class: Optional[str] = None,
        rst: bool = True,
    ) -> List[str]:
        node = resolved.node
        if not node.has_docstring():
            return []

        if name is None:
            name = resolved.name
        if context_module is None and resolved.modules:
            context_module = ".".join([module.name for module in resolved.modules])
        if context_class is None and resolved.classes:
            context_module = ".".join([class_.name for class_ in resolved.classes])

        if isinstance(node, Class):
            return self._format_class(
                cast(ResolvedNode[Class], resolved),
                name,
                context_module,
                context_class,
                rst,
            )
        elif isinstance(node, Function):
            return self._format_function(
                cast(ResolvedNode[Function], resolved),
                name,
                context_module,
                context_class,
                rst,
            )
        elif isinstance(node, Variable):
            return self._format_variable(
                cast(ResolvedNode[Variable], resolved),
                name,
                context_module,
                context_class,
                rst,
            )
        else:
            assert isinstance(node, Module)
            assert node.docstring is not None
            return node.docstring.splitlines()