

from __future__ import annotations

import argparse
import ast
import importlib
import inspect
import pathlib
import sys

def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Verify that imported functions/classes have docstrings.")
    p.add_argument(
        "--target-file",
        default="verl/trainer/ppo/ray_trainer.py",
        help="Path to the Python source file to analyse (e.g. verl/trainer/ppo/ray_trainer.py)",
    )
    p.add_argument(
        "--allow-list",
        default=["omegaconf.open_dict"],
        help="a list of third_party dependencies that do not have proper docs :(",
    )
    p.add_argument(
        "--project-root",
        default=".",
        help="Directory to prepend to PYTHONPATH so local packages resolve (default: .)",
    )
    p.add_argument(
        "--quiet",
        action="store_true",
        help="Suppress success message (still prints errors).",
    )
    return p.parse_args()

def _import_attr(module_name: str, attr_name: str):
    module = importlib.import_module(module_name)
    return getattr(module, attr_name)

def _check_file(py_file: pathlib.Path, project_root: pathlib.Path, allow_list: list[str]) -> list[str]:

    sys.path.insert(0, str(project_root.resolve()))

    tree = ast.parse(py_file.read_text(), filename=str(py_file))
    problems: list[str] = []

    for node in ast.walk(tree):
        if not isinstance(node, ast.ImportFrom):
            continue

        module_name = "." * node.level + (node.module or "")
        for alias in node.names:
            if alias.name == "*":
                problems.append(
                    f"{py_file}:{node.lineno} - wildcard import `from {module_name} import *` cannot be verified."
                )
                continue

            imported_name = alias.name

            try:
                obj = _import_attr(module_name, imported_name)
            except Exception:
                pass

                continue

            if f"{module_name}.{imported_name}" in allow_list:
                continue
            if inspect.isfunction(obj) or inspect.isclass(obj):
                doc = inspect.getdoc(obj)
                if not (doc and doc.strip()):
                    kind = "class" if inspect.isclass(obj) else "function"
                    problems.append(
                        f"{py_file}:{node.lineno} - {kind} `{module_name}.{imported_name}` is missing a docstring."
                    )

    return problems

def main() -> None:
    args = _parse_args()
    target_path = pathlib.Path(args.target_file).resolve()
    project_root = pathlib.Path(args.project_root).resolve()

    if not target_path.is_file():
        raise Exception(f"❌ Target file not found: {target_path}")

    errors = _check_file(target_path, project_root, args.allow_list)

    if errors:
        print("Docstring verification failed:\n")
        print("\n".join(f" • {e}" for e in errors))
        raise Exception("❌ Docstring verification failed.")

    if not args.quiet:
        print(f"✅ All explicitly imported functions/classes in {target_path} have docstrings.")

if __name__ == "__main__":
    main()
