# -*- coding: utf-8 -*-

import argparse
import ast
import inspect
import itertools
import operator
import os
import pathlib
import re
import sys
import typing

# Docstring sections: https://www.sphinx-doc.org/en/master/usage/domains/python.html#info-field-lists
ptype_key = "type"
rtype_key = "rtype"
param_set = {"param", "parameter", "arg", "argument", "keyword", "key"}
return_set = {"return", "returns"}
raises_set = {"raises", "raise", "except", "exception"}
ignore_set = {"var", "ivar", "cvar", "vartype", "meta"}  # At this moment, it's ignored

# Summary (everything before the first section or the end of the string)
summary_regex = re.compile(r'^(.*?)(?=^:|\Z)', flags=re.DOTALL | re.MULTILINE)
# Section (start with ':' and end before next ':' at start of line or end of string)
section_regex = re.compile(r'(^:.*?)(?=^:|\Z)', flags=re.DOTALL | re.MULTILINE)
# Trailing empty lines
trailing_regex = re.compile("(?:^\\s*$){2,}\\Z", flags=re.MULTILINE)
# Consecutive empty lines (not at end)
empty_lines_regex = re.compile("(?:^[ \t]*\r?\n){2,}(?=[^\r\n])", re.MULTILINE)


class ParsedDocsParam(typing.NamedTuple):
    section_key: str
    sep: bool  # True if ':' was present
    param_name: str | None
    param_type: str | None


class ParsedDocsReturn(typing.NamedTuple):
    section_key: str
    sep: bool  # True if ':' was present
    return_type: str | None
    description: str | None  # Only for ´:return:´


class ParsedDocsRaise(typing.NamedTuple):
    section_key: str
    sep: bool  # True if ':' was present
    error_types: list[str]


class ParsedDocs(typing.NamedTuple):
    summary: str | None  # The summary section (before the first section)
    params: list[ParsedDocsParam]
    returns: list[ParsedDocsReturn]
    raises: list[ParsedDocsRaise]
    invalid: list[str]  # invalid sections

    rawdocs: str | None  # Raw docstring
    docs: str | None  # The full docstring
    docs_ini_lineno: int | None  # First line number of the docstring, None if no docstring
    docs_end_lineno: int | None  # End line number of the docstring, None if no docstring
    code_ini_lineno: int | None  # First line number of the code block after the docstring, None if no code


class Violations:
    # DOC0xx: Docstring section issues
    DOC001 = ("DOC001", "Unknown docstring section ({!r})")
    DOC002 = ("DOC002", "Malformed section ({!r})")
    DOC003 = ("DOC003", "Missing blank line after docstring")
    DOC004 = ("DOC004", "Missing blank line between summary and sections")
    DOC005 = ("DOC005", "Too many consecutive empty lines")
    DOC006 = ("DOC006", "Trailing empty lines")

    # DOC1xx: Argument issues
    DOC101 = ("DOC101", "Parameter documented but not in signature ({!r})")
    DOC102 = ("DOC102", "Invalid parameter type syntax ({!r})")
    DOC103 = ("DOC103", "Parameter type already in signature ({!r})")
    DOC104 = ("DOC104", "Parameter type mismatch with hint ({!r} != {!r})")
    DOC105 = ("DOC105", "Duplicated parameter ({!r})")
    # DOC2xx: Return issues
    DOC201 = ("DOC201", "Return documented but function has no return")
    DOC202 = ("DOC202", "Invalid return type syntax ({!r})")
    DOC203 = ("DOC203", "Return type already in signature ({!r})")
    DOC204 = ("DOC204", "Return type mismatch with annotation ({!r} != {!r})")
    DOC205 = ("DOC205", "Duplicated return section ({!r})")
    # DOC3xx: Raise issues
    DOC302 = ("DOC302", "Invalid exception type syntax ({!r})")
    DOC305 = ("DOC305", "Duplicated exception type ({!r})")

    @staticmethod
    def is_valid_syntax(value, /):
        try:
            ast.parse(value, mode="eval")
            return True
        except SyntaxError:
            return False

    @classmethod
    def is_valid_type_hint(cls, hint, /):
        if hint:
            try:
                ast.parse(f"x: {hint}")
                return True
            except SyntaxError:
                return False
        else:
            return False

    @classmethod
    def validate_empty_lines(cls, parsed, /):
        text = parsed.rawdocs
        if text:
            hits = empty_lines_regex.finditer(text)
            present = next(filter(None, hits), None)
            if present:
                yield Violations.DOC005, ()

            hits = trailing_regex.finditer(text)
            present = next(filter(None, hits), None)
            if present:
                yield Violations.DOC006, ()

    @classmethod
    def validate_summary(cls, parsed, /):
        if parsed.summary:
            summary_lines = parsed.summary.splitlines()
            br_tail_count = sum(1 for _ in itertools.takewhile(operator.not_, reversed(summary_lines)))

            if br_tail_count == 0 and (parsed.params + parsed.returns + parsed.raises):
                # Only applies if there are sections after the summary
                yield Violations.DOC004, ()  # Missing blank line between summary and sections

    @classmethod
    def validate_params(cls, parsed, parameters, /):
        bag = set()
        for param in parsed.params:
            section_key, sep, name, kind = param
            type_hint = parameters.get(name)
            # Malformed param: missing ':' or missing name/type when required
            if not (sep and name) or (section_key == ptype_key and not kind):
                yield Violations.DOC002, (section_key,)
            if name and name not in parameters:  # Documented but not in signature
                yield Violations.DOC101, (name,)
            if kind and not Violations.is_valid_type_hint(kind):  # Invalid type syntax
                yield Violations.DOC102, (kind,)
            if kind and type_hint and kind == type_hint:  # Redundant type
                yield Violations.DOC103, (kind,)
            if kind and type_hint and kind != type_hint:  # Mismatched type
                yield Violations.DOC104, (kind, type_hint)
            if name and name in bag:  # Duplicated parameter
                yield Violations.DOC105, (name,)
            if name:
                bag.add(name)

    @classmethod
    def validate_return(cls, parsed, type_hint, has_returns, is_implemented, /):
        if parsed.returns and (not has_returns and is_implemented):
            yield Violations.DOC201, ()  # Return documented but implementation has no return

        bag = set()
        for doc_return in parsed.returns:
            section_key, sep, return_type, description = doc_return
            # Malformed return: missing ':' or missing type/description when required
            if not sep:
                # Missing ':' separator
                yield Violations.DOC002, (section_key,)
            elif section_key in return_set and not description:
                # :return:´ without description
                yield Violations.DOC002, (section_key,)
            elif section_key == rtype_key and not return_type:  # ´:rtype:´ without type
                yield Violations.DOC002, (section_key,)
            # Invalid or redundant return type checks
            if return_type and not Violations.is_valid_type_hint(return_type):  # Invalid return type
                yield Violations.DOC202, (return_type,)
            if return_type and type_hint and return_type == type_hint:  # Redundant return type
                yield Violations.DOC203, (return_type,)
            if return_type and type_hint and return_type != type_hint:  # Mismatched return type
                yield Violations.DOC204, (return_type, type_hint)
            # Duplicate return section
            if section_key in bag:
                yield Violations.DOC205, (section_key,)  # Duplicate return section
            bag.add(section_key)

    @classmethod
    def validate_raises(cls, parsed, /):
        bag = set()
        for section_key, sep, error_types in parsed.raises:
            if not (sep and error_types):  # Missing ':' or missing error types
                yield Violations.DOC002, (section_key,)

            is_invalid = any(not Violations.is_valid_syntax(e) for e in error_types)
            if is_invalid:  # Invalid exception type syntax
                yield Violations.DOC302, (is_invalid,)

            for error_type in error_types:
                if error_type in bag:  # Duplicate exception type
                    yield Violations.DOC305, (error_type,)
                bag.add(error_type)

    @classmethod
    def discover(cls, parsed, parameters, has_returns, is_implemented, /):
        if parsed.code_ini_lineno and parsed.docs_end_lineno and parsed.code_ini_lineno - parsed.docs_end_lineno == 1:
            yield Violations.DOC003, ()
        yield from ((Violations.DOC001, (section_key,)) for section_key in parsed.invalid)

        yield from cls.validate_empty_lines(parsed)
        yield from cls.validate_summary(parsed)
        yield from cls.validate_params(parsed, parameters)
        yield from cls.validate_return(parsed, parameters.get("return"), has_returns, is_implemented)
        yield from cls.validate_raises(parsed)


def parse_section_param(section_key, sep, parts_a, /):
    if len(parts_a) == 1:  # ´:param:´
        param_name = None
        param_type = None
    elif len(parts_a) == 2:  # ´:param [ParamName]:´
        param_name = parts_a[1]
        param_type = None
    else:  # ´:param [ParamType] [ParamName]:´
        param_type = " ".join(parts_a[1:-1])  # unsplit type
        param_name = parts_a[-1]
    return ParsedDocsParam(section_key, sep, param_name, param_type)


def parse_section_type(section_key, sep, parts_a, parts_b, /):
    if len(parts_a) == 1:  # ´:type:´
        param_name = None
    else:  # >=2   ´:type [ParamName]:´
        param_name = " ".join(parts_a[1:])

    if parts_b:  # ´:type [ParamName]: [ParamType]´
        param_type = " ".join(parts_b)
    else:  # ´:type [ParamName]:´ (without type)
        param_type = None
    return ParsedDocsParam(section_key, sep, param_name, param_type)


def parse_section_rtype(section_key, sep, parts_b, /):
    if parts_b:  # ´:rtype: [ReturnType]´
        return_type = " ".join(parts_b)
    else:  # ´:rtype:´ (without type)
        return_type = None
    return ParsedDocsReturn(section_key, sep, return_type, None)


def parse_section_raise(section_key, sep, parts_a, /):
    if len(parts_a) == 1:  # ´:raise [ErrorType]: [ErrorDescription]´
        error_types = []
    else:  # >=2  ´:raise [ErrorType1, ErrorType2]: [ErrorDescription]´
        error_types = "".join(parts_a[1:]).split(",")  # commas separate multiple error types
    return ParsedDocsRaise(section_key, sep, error_types)


def parse_section_return(section_key, sep, parts_a, parts_b, /):
    if parts_b:  # ´:return: [ReturnDescription]´
        description = " ".join(parts_b)
    else:  # ´:return:´ (without description)
        description = None
    return ParsedDocsReturn(section_key, sep, None, description)


def itersections(docstring, /):
    if docstring:
        yield from (chunk.strip() for match in section_regex.finditer(docstring) if (chunk := match.group(0)))


def get_summary(docstring, /):
    if docstring and (match := summary_regex.match(docstring)):
        return match.group(0)
    else:
        return None


def parse_docs(node, /):
    params = list()
    raises = list()
    returns = list()
    invalid = list()
    rawdocs = ast.get_docstring(node, clean=False)
    docstring = rawdocs if rawdocs is None else inspect.cleandoc(rawdocs)

    for section in itersections(docstring):
        a, sep, b = section.lstrip(":").partition(":")
        sep = bool(sep)  # True if ':' was present
        b = b.splitlines()[0].strip() if (b := b.strip()) else ""  # Only first line of description is relevant.
        parts_a = a.split()
        parts_b = b.split()
        section_key = parts_a[0].lower()
        if section_key in param_set:
            params.append(parse_section_param(section_key, sep, parts_a))
        elif section_key == ptype_key:
            params.append(parse_section_type(section_key, sep, parts_a, parts_b))
        elif section_key in rtype_key:
            returns.append(parse_section_rtype(section_key, sep, parts_b))
        elif section_key in raises_set:
            raises.append(parse_section_raise(section_key, sep, parts_a))
        elif section_key in return_set:
            returns.append(parse_section_return(section_key, sep, parts_a, parts_b))
        elif section_key not in ignore_set:
            invalid.append(section_key)

    if docstring:
        summary = get_summary(docstring)
        first = node.body[0]
        docs_ini_lineno, docs_end_lineno = (first.lineno, first.end_lineno)
        code_ini_lineno = node.body[1].lineno if len(node.body) > 1 else None
    else:
        summary = None
        docs_ini_lineno = docs_end_lineno = None
        code_ini_lineno = node.body[0].lineno if node.body else None

    return ParsedDocs(
        summary=summary,
        rawdocs=rawdocs,
        docs=docstring,
        docs_ini_lineno=docs_ini_lineno,
        docs_end_lineno=docs_end_lineno,
        code_ini_lineno=code_ini_lineno,
        params=params,
        returns=returns,
        raises=raises,
        invalid=invalid,
    )


def has_return_or_yield(node, /):
    """
    Returns True if the given AST node (ast.FunctionDef or ast.AsyncFunctionDef)
    contains a 'return', 'yield', or 'yield from' statement in its main body
    (excluding nested functions).

    :param ast.FunctionDef | ast.AsyncFunctionDef node: Root node to explore.
    :return: Whether the function has either "return" or "yield".
    :rtype: bool
    """

    class ReturnYieldVisitor(ast.NodeVisitor):
        def __init__(self):
            self.found = False
            self.depth = 0

        def visit_nested(self, node):
            if self.depth == 0:
                self.depth += 1
                self.generic_visit(node)
                self.depth -= 1

        def visit_FunctionDef(self, node):
            self.visit_nested(node)

        def visit_AsyncFunctionDef(self, node):
            self.visit_nested(node)

        def visit_Return(self, node):
            self.found = True

        def visit_Yield(self, node):
            self.found = True

        def visit_YieldFrom(self, node):
            self.found = True

    visitor = ReturnYieldVisitor()
    visitor.visit(node)
    return visitor.found


def is_not_implemented(node, /, rawdocs=None):
    """
    Returns True if the given AST node's body contains exactly one statement
    and that statement is one of the following:
      - `pass`
      - `raise NotImplementedError` or `raise NotImplementedError()`
      - `...` (Ellipsis literal)
    Otherwise, returns False.

    :param ast.FunctionDef | ast.AsyncFunctionDef node: Root node to explore.
    :param str | None rawdocs: Raw docstring.
    """

    if rawdocs and len(node.body) > 2 or (not rawdocs and len(node.body) > 1):
        return False  # More than one statement
    elif rawdocs and len(node.body) == 1:
        return True  # Only docstring, no code
    else:  # One statement (without docstring or after docstring)
        stmt = node.body[-1]

        if isinstance(stmt, ast.Pass):  # Case: "pass"
            return True

        elif isinstance(stmt, ast.Raise):  # Case: "raise NotImplementedError()"
            exc = stmt.exc
            error_name = NotImplementedError.__name__
            if (
                    isinstance(exc, ast.Call)
                    and isinstance(exc.func, ast.Name)
                    and exc.func.id == error_name
            ):
                return True

            else:  # Case: "raise NotImplementedError"
                return isinstance(exc, ast.Name) and exc.id == error_name

        elif isinstance(stmt, ast.Expr):  # Case: "..."
            return isinstance(stmt.value, ast.Constant) and (stmt.value.value is Ellipsis)
        else:
            return False


def checker(node, /):
    """
    Returns True if the given AST node (ast.FunctionDef or ast.AsyncFunctionDef)
    contains a 'return', 'yield', or 'yield from' statement in its main body
    (excluding nested functions).

    :param ast.FunctionDef | ast.AsyncFunctionDef node: Root node to explore.
    :return: Generator with data of each violation.
    :rtype: typing.Iterator[tuple[int, str, dict]]
    """

    params = dict(get_params(node))
    parsed = parse_docs(node)
    lineno = parsed.docs_ini_lineno
    has_returns = has_return_or_yield(node)
    is_implemented = not is_not_implemented(node, rawdocs=parsed.rawdocs)

    for (code, msg), ctx in Violations.discover(parsed, params, has_returns, is_implemented):
        yield (lineno, code, msg, ctx)


def get_args(node, /):
    yield from node.args.posonlyargs  # Positional-only args
    yield from node.args.args  # Regular args
    yield from node.args.kwonlyargs  # Keyword-only args
    yield from filter(None, (node.args.vararg, node.args.kwarg))  # *args, **kwargs


def get_params(node, /):
    for arg in get_args(node):
        yield (arg.arg, ast.unparse(ann) if (ann := arg.annotation) else None)
    if node.returns:
        yield ("return", ast.unparse(node.returns))


def check_node(filename, node, /):
    fmt = "{}:{}: [{}] {}".format
    result = True
    for lineno, code, msg, ctx in checker(node):
        print(fmt(filename, lineno, code, msg.format(*ctx)))
        result = False

    return result


def walk_module(data, filename, /):
    """
    Yields each function node from the parsed "data" tree.

    :param bytes data: Content to parse.
    :param str filename: File name to use when printing messages.
    :return: Each root node of every function/method.
    :rtype: Iterator[ast.FunctionDef | ast.AsyncFunctionDef]
    """

    try:
        tree = ast.parse(data, filename=filename)
    except SyntaxError:
        pass
    else:
        klasses = (ast.FunctionDef, ast.AsyncFunctionDef)
        for node in ast.walk(tree):
            if isinstance(node, klasses):
                yield node


def walk(paths, ignore_dirs, /):
    ignore_dirs = frozenset(ignore_dirs)
    for path in paths:
        if os.path.isfile(path) and path.endswith(".py"):
            yield pathlib.Path(path)
        elif os.path.isdir(path):
            for rpath in pathlib.Path(path).rglob("*.py"):
                if not any(part in ignore_dirs for part in rpath.parts):  # Skip ignored dirs
                    yield rpath


def main():
    ignore = [".venv", ".env", ".git", ".pytest_cache", ".ruff_cache", "__pycache__", "site-packages"]
    parser = argparse.ArgumentParser(description="Sphinx docstring checker")
    parser.add_argument("files", nargs="*", help="files or dirs to check", default=[os.getcwd()])
    parser.add_argument("--ignore", nargs="*", help="directories to ignore", default=ignore)
    args = parser.parse_args()
    result = False
    for path in walk(args.files, args.ignore):
        filename = str(path)
        for node in walk_module(path.read_bytes(), filename):
            if not check_node(filename, node):
                result = True

    return 1 if result else 0


if __name__ == "__main__":
    sys.exit(main())
