

from __future__ import annotations

import argparse
import importlib
import inspect
import pkgutil
import sys
from pathlib import Path
from types import ModuleType
from typing import Iterable

_ALLOW_LIST = [
    "verl.third_party.vllm.LLM",
    "verl.third_party.vllm.parallel_state",
    "verl.utils.profiler.WorkerProfiler",
    "verl.utils.profiler.WorkerProfilerExtension",
    "verl.utils.profiler.log_gpu_memory_usage",
    "verl.utils.profiler.log_print",
    "verl.utils.profiler.mark_annotate",
    "verl.utils.profiler.mark_end_range",
    "verl.utils.profiler.mark_start_range",
    "verl.models.mcore.qwen2_5_vl.get_vision_model_config",
    "verl.models.mcore.qwen2_5_vl.get_vision_projection_config",
]

def iter_submodules(root: ModuleType) -> Iterable[ModuleType]:
    yield root
    if getattr(root, "__path__", None):
        for mod_info in pkgutil.walk_packages(root.__path__, prefix=f"{root.__name__}."):
            try:
                yield importlib.import_module(mod_info.name)
            except Exception as exc:
                print(f"[warn] Skipping {mod_info.name!r}: {exc}", file=sys.stderr)

def names_missing_doc(mod: ModuleType) -> list[str]:
    missing: list[str] = []
    public = getattr(mod, "__all__", [])
    for name in public:
        obj = getattr(mod, name, None)
        if f"{mod.__name__}.{name}" in _ALLOW_LIST:
            continue
        if obj is None:

            missing.append(f"{mod.__name__}.{name}  (not found)")
            continue

        if inspect.isfunction(obj) or inspect.isclass(obj):
            doc = inspect.getdoc(obj)
            if not doc or not doc.strip():
                missing.append(f"{mod.__name__}.{name}")
    return missing

def check_module(qualname: str) -> list[str]:
    try:
        module = importlib.import_module(qualname)
    except ModuleNotFoundError as exc:
        print(f"[error] Cannot import '{qualname}': {exc}", file=sys.stderr)
        return [qualname]

    missing: list[str] = []
    for submod in iter_submodules(module):
        missing.extend(names_missing_doc(submod))
    return missing

def autodiscover_packages() -> list[str]:
    pkgs: list[str] = []
    for p in Path.cwd().iterdir():
        if p.is_dir() and (p / "__init__.py").exists():
            pkgs.append(p.name)
    return pkgs

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "modules",
        nargs="*",
        help="Fully-qualified module or package names (defaults to every top-level package found in CWD).",
    )
    args = parser.parse_args()

    targets = args.modules or autodiscover_packages()
    if not targets:
        raise ValueError("[error] No modules specified and none detected automatically.")

    all_missing: list[str] = []
    for modname in targets:
        all_missing.extend(check_module(modname))

    if all_missing:
        print("\nMissing docstrings:")
        for name in sorted(all_missing):
            print(f"  - {name}")
        raise ValueError("Missing docstrings detected. Please enhance them with docs accordingly.")

    print("✅ All exported functions/classes have docstrings.")

if __name__ == "__main__":
    main()
