#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# SPDX-License-Identifier: LGPL-2.1-or-later

import argparse
import functools
import sys
from typing import cast

from drgndoc.format import Formatter
from drgndoc.namespace import Namespace, ResolvedNode
from drgndoc.parse import Class, DocumentedNode, Node, parse_paths
from drgndoc.util import dot_join

escapes = []
for c in range(256):
    if c == 0:
        e = r"\0"
    elif c == 7:
        e = r"\a"
    elif c == 8:
        e = r"\b"
    elif c == 9:
        e = r"\t"
    elif c == 10:
        e = r"\n"
    elif c == 11:
        e = r"\v"
    elif c == 12:
        e = r"\f"
    elif c == 13:
        e = r"\r"
    elif c == 34:
        e = r"\""
    elif c == 92:
        e = r"\\"
    elif 32 <= c <= 126:
        e = chr(c)
    else:
        e = f"\\x{c:02x}"
    escapes.append(e)


def escape_string(s: str) -> str:
    return "".join([escapes[c] for c in s.encode("utf-8")])


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="generate docstring definitions for a C extension from Python source code/stub files"
    )
    parser.add_argument(
        "--header", "-H", action="store_true", help="generate header file"
    )
    parser.add_argument(
        "-m",
        "--module",
        dest="modules",
        metavar="MODULE[:NAME]",
        action="append",
        help="generate docstrings for the given module instead of all modules "
        "(may be given multiple times); "
        "an alternate name to use for the generated variables may also be given",
    )
    parser.add_argument(
        "paths", metavar="PATH", nargs="+", help="module or package path"
    )
    args = parser.parse_args()

    modules = parse_paths(args.paths, functools.partial(print, file=sys.stderr))
    namespace = Namespace(modules)
    formatter = Formatter(namespace)
    output_file = sys.stdout

    if args.header:
        output_file.write(
            """\
/*
 * Generated by drgndoc.docstrings -H.
 *
 * Before Python 3.7, various docstring fields were defined as char * (see
 * https://bugs.python.org/issue28761). We still want the strings to be
 * read-only, so just cast away the const.
 */

"""
        )
    else:
        output_file.write("/* Generated by drgndoc.docstrings. */\n\n")

    def aux(resolved: ResolvedNode[Node], name: str) -> None:
        node = resolved.node
        if node.has_docstring():
            var_name = name.replace(".", "_") + "_DOC"
            if args.header:
                output_file.write("extern ")
            output_file.write(f"const char {var_name}[]")
            if not args.header:
                output_file.write(" =")
                lines = formatter.format(
                    cast(ResolvedNode[DocumentedNode], resolved),
                    name.rpartition(".")[2],
                    rst=False,
                )
                if lines:
                    for i, line in enumerate(lines):
                        output_file.write(f'\n\t"{escape_string(line)}')
                        if i != len(lines) - 1:
                            output_file.write("\\n")
                        output_file.write('"')
                else:
                    output_file.write(' ""')
            output_file.write(";\n")
            if args.header:
                output_file.write(f"#define {var_name} (char *){var_name}\n")
        for attr in resolved.attrs():
            if isinstance(node, Class) and attr.name == "__init__":
                continue
            aux(attr, dot_join(name, attr.name))

    for module in args.modules or namespace.modules.keys():
        module, _, name = module.partition(":")
        resolved = namespace.resolve_global_name(module)
        if isinstance(resolved, ResolvedNode):
            aux(resolved, name or module)
        else:
            sys.exit(f"name {module} not found")