Modules/_docstrings.py

from __future__ import annotations

import inspect
import re
from typing import Any, Annotated, get_args, get_origin, get_type_hints


def _typename(tp: Any) -> str:
    """Return a readable type name from a type or annotation."""
    try:
        if hasattr(tp, "__name__"):
            return tp.__name__  # e.g. int, str
        if getattr(tp, "__module__", None) and getattr(tp, "__qualname__", None):
            return f"{tp.__module__}.{tp.__qualname__}"
        return str(tp).replace("typing.", "")
    except Exception:
        return str(tp)


def _parse_string_annotation(annot_str: str) -> tuple[str | None, str | None]:
    """
    Parse a string annotation like "Annotated[Optional[str], 'description']"
    and extract the base type name and the description metadata.
    
    Returns (base_type_name, description) or (None, None) if parsing fails.
    """
    if not isinstance(annot_str, str):
        return None, None
    
    # Match Annotated[..., 'description'] or Annotated[..., "description"]
    # Pattern: Annotated[<base_type>, '<description>'] or with double quotes
    match = re.match(
        r"^Annotated\[(.+?),\s*['\"](.+?)['\"]\s*\]$",
        annot_str.strip(),
        re.DOTALL,
    )
    if match:
        base_type_str = match.group(1).strip()
        description = match.group(2)
        # Simplify Optional[X] -> just the base type for display
        opt_match = re.match(r"^Optional\[(.+)\]$", base_type_str)
        if opt_match:
            base_type_str = opt_match.group(1).strip()
        return base_type_str, description
    
    return None, None


def _extract_base_and_meta(annotation: Any) -> tuple[Any, str | None]:
    """Given an annotation, return (base_type, first string metadata) if Annotated, else (annotation, None)."""
    try:
        # Handle string annotations from PEP 563 (__future__.annotations)
        if isinstance(annotation, str):
            base_str, meta = _parse_string_annotation(annotation)
            if meta:
                return base_str or annotation, meta
            return annotation, None
        
        if get_origin(annotation) is Annotated:
            args = get_args(annotation)
            base = args[0] if args else annotation
            # Grab the first string metadata if present
            for meta in args[1:]:
                if isinstance(meta, str):
                    return base, meta
            return base, None
        return annotation, None
    except Exception:
        return annotation, None


def autodoc(summary: str | None = None, returns: str | None = None, *, force: bool = False):
    """
    Decorator that auto-generates a concise Google-style docstring from a function's
    type hints and Annotated metadata. Useful for Gradio MCP where docstrings are
    used for tool descriptions and parameter docs.

    Args:
        summary: Optional one-line summary for the function. If not provided,
            will generate a simple sentence from the function name.
        returns: Optional return value description. If not provided, only the
            return type will be listed (if available).
        force: When True, overwrite an existing docstring. Default False.

    Returns:
        The original function with its __doc__ populated (unless skipped).
    """

    def decorator(func):
        # Skip if docstring already present and not forcing
        if not force and func.__doc__ and func.__doc__.strip():
            return func

        try:
            # include_extras=True to retain Annotated metadata
            hints = get_type_hints(func, include_extras=True, globalns=getattr(func, "__globals__", None))
        except Exception:
            hints = {}

        sig = inspect.signature(func)

        lines: list[str] = []
        # Summary line
        if summary and summary.strip():
            lines.append(summary.strip())
        else:
            pretty = func.__name__.replace("_", " ").strip().capitalize()
            if not pretty.endswith("."):
                pretty += "."
            lines.append(pretty)

        # Args section
        if sig.parameters:
            lines.append("")
            lines.append("Args:")
            for name, param in sig.parameters.items():
                if name == "self":
                    continue
                annot = hints.get(name, param.annotation)
                base, meta = _extract_base_and_meta(annot)
                tname = _typename(base) if base is not inspect._empty else None
                desc = meta or ""
                if tname and tname != str(inspect._empty):
                    lines.append(f"    {name} ({tname}): {desc}".rstrip())
                else:
                    lines.append(f"    {name}: {desc}".rstrip())

        # Returns section
        ret_hint = hints.get("return", sig.return_annotation)
        if returns or (ret_hint and ret_hint is not inspect.Signature.empty):
            lines.append("")
            lines.append("Returns:")
            if returns:
                lines.append(f"    {returns}")
            else:
                base, meta = _extract_base_and_meta(ret_hint)
                rtype = _typename(base)
                if meta:
                    lines.append(f"    {rtype}: {meta}")
                else:
                    lines.append(f"    {rtype}")

        func.__doc__ = "\n".join(lines).strip() + "\n"
        return func

    return decorator


__all__ = ["autodoc"]