"""
FastAPI Docstring-to-OpenAPI Enricher
=====================================

自动解析 Google Style docstring，将 Args / Returns / Raises 注入到 OpenAPI Schema 的对应位置，
让 vitepress-openapi 等工具能够正确渲染参数表格和返回值说明。

功能特性
--------
1. 参数描述注入：路径参数、查询参数、Header、Cookie、请求体字段
2. 返回值描述注入：自动写入 2xx 响应的 description 及 schema.description
3. 异常描述注入：以 Markdown 格式附加到 operation description
4. 嵌套 Schema 递归处理：支持 Pydantic v1/v2 的嵌套模型、List、Optional 等
5. \f 截断兼容：与 FastAPI 原生行为保持一致
6. 全局 Schema 补全：为 components/schemas 中缺少描述的字段补充 docstring 说明
7. 零侵入：只需替换 ``app.openapi``，业务代码无需任何修改

依赖安装
--------
    pip install docstring-parser

使用方式
--------
    from fastapi import FastAPI
    from fastapi_docstring_openapi import enrich_openapi_from_docstrings

    app = FastAPI()
    enrich_openapi_from_docstrings(app)

    # 你的路由注册...

完整示例见文件底部 ``if __name__ == "__main__":`` 块。
"""

from __future__ import annotations

import inspect
import logging
from typing import Any, Dict, Optional, Set

from docstring_parser import Docstring, DocstringStyle, ParseError, parse
from fastapi import FastAPI
from fastapi.openapi.utils import get_openapi

logger = logging.getLogger(__name__)


def _parse_docstring(func: Any, style: DocstringStyle = DocstringStyle.AUTO) -> Optional[Docstring]:  # noqa ANN401 有意为之
    """
    解析函数的 Google Style docstring，失败时返回 None。

    Returns:
        成功解析时返回 docstring_parser.Docstring。
        失败时返回 None。
    """
    doc = inspect.getdoc(func)
    if not doc:
        return None
    if "\f" in doc:
        doc = doc.split("\f")[0].strip()
    try:
        return parse(doc, style=style)
    except ParseError:
        return None
    except Exception:
        logger.warning(
            "Unexpected error parsing docstring for %s",
            getattr(func, "__name__", func),
            exc_info=True,
        )
        return None


def _build_param_lookup(parsed_doc: Docstring) -> Dict[str, str]:
    """
    从解析后的 docstring 构建 {参数名: 描述} 映射。

    Returns:
        {参数名: 描述}
    """
    lookup: Dict[str, str] = {}
    for param in getattr(parsed_doc, "params", []) or []:
        name = getattr(param, "arg_name", None)
        desc = getattr(param, "description", None) or ""
        if name:
            # 合并多行描述并清理缩进
            lookup[name] = _dedent_description(desc)
    return lookup


def _build_returns_description(parsed_doc: Docstring) -> Optional[str]:
    """
    从解析后的 docstring 构建返回值描述。

    Returns:
        可能的字符串为返回值描述。
    """
    ret = getattr(parsed_doc, "returns", None)
    if not ret:
        return None
    type_name = getattr(ret, "type_name", None) or ""
    desc = getattr(ret, "description", None) or ""
    if type_name and desc:
        return f"**{type_name}**: {desc}"
    return desc or None


def _build_raises_description(parsed_doc: Docstring) -> Optional[str]:
    """
    从解析后的 docstring 构建异常描述（Markdown 列表）。

    Returns:
        可能的字符串为异常描述。
    """
    raises = getattr(parsed_doc, "raises", None)
    if not raises:
        return None
    parts: list[str] = []
    for exc in raises:
        type_name = getattr(exc, "type_name", None) or ""
        desc = getattr(exc, "description", None) or ""
        if type_name and desc:
            parts.append(f"- **{type_name}**: {desc}")
        elif type_name:
            parts.append(f"- **{type_name}**")
        elif desc:
            parts.append(f"- {desc}")
    return "\n".join(parts) if parts else None


def _dedent_description(text: str) -> str:
    """
    清理 docstring 描述中的多余缩进，保留段落结构。

    Returns:
        docstring 字符串
    """
    if not text:
        return ""
    lines = text.splitlines()
    # 找到最小公共缩进（排除空行）
    min_indent = min(
        (len(line) - len(line.lstrip()) for line in lines if line.strip()),
        default=0,
    )
    cleaned = [line[min_indent:] if line.strip() else "" for line in lines]
    return "\n".join(cleaned).strip()


def _get_clean_operation_description(parsed_doc: Docstring) -> str:
    """
    生成干净的 operation description。

    Returns:
        策略：保留 summary + long_description，去掉 Args/Returns/Raises 等机器块。
    """
    parts: list[str] = []
    short = getattr(parsed_doc, "short_description", None)
    long_ = getattr(parsed_doc, "long_description", None)

    if short:
        parts.append(short)
    if long_:
        parts.append(long_)

    return "\n\n".join(parts).strip()


def _inject_schema_properties(
    schema: Dict[str, Any],
    param_lookup: Dict[str, str],
    visited_refs: Optional[Set[str]] = None,
    openapi_schema: Optional[Dict[str, Any]] = None,
) -> None:
    """
    递归注入 schema properties 的描述。

    支持对象、数组、allOf/anyOf/oneOf、$ref 引用（全局 components/schemas 补全）。
    """
    if visited_refs is None:
        visited_refs = set()
    if not isinstance(schema, dict):
        return

    # 处理 $ref：在 components/schemas 中查找并补全
    ref = schema.get("$ref")
    if ref and openapi_schema:
        if ref not in visited_refs:
            visited_refs.add(ref)
            # 提取 schema 名，例如 "#/components/schemas/User" -> "User"
            schema_name = ref.split("/")[-1]
            components = openapi_schema.get("components", {}).get("schemas", {})
            target = components.get(schema_name)
            if target:
                _inject_schema_properties(target, param_lookup, visited_refs, openapi_schema)
        return

    # 处理 properties（对象字段）
    properties = schema.get("properties")
    if isinstance(properties, dict):
        for prop_name, prop_schema in properties.items():
            if prop_name in param_lookup:  # noqa SIM102
                # 只有当 docstring 描述非空时才写入
                if param_lookup[prop_name]:
                    prop_schema["description"] = param_lookup[prop_name]
            # 递归处理子 schema
            _inject_schema_properties(prop_schema, param_lookup, visited_refs, openapi_schema)

    # 处理 items（数组类型）
    items = schema.get("items")
    if items:
        _inject_schema_properties(items, param_lookup, visited_refs, openapi_schema)

    # 处理 allOf / anyOf / oneOf
    for key in ("allOf", "anyOf", "oneOf"):
        for sub in schema.get(key, []):
            _inject_schema_properties(sub, param_lookup, visited_refs, openapi_schema)

    # 处理 additionalProperties
    additional = schema.get("additionalProperties")
    if isinstance(additional, dict):
        _inject_schema_properties(additional, param_lookup, visited_refs, openapi_schema)


def enrich_openapi_from_docstrings(
    app: FastAPI,
    *,
    prefer_docstring_over_field: bool = True,
    append_raises_to_description: bool = True,
    docstring_style: DocstringStyle = DocstringStyle.AUTO,
) -> None:
    """
    为 FastAPI 应用启用 docstring 驱动的 OpenAPI 增强。

    参数
    ----
    app : FastAPI
        目标应用实例。
    prefer_docstring_over_field : bool, 默认 True
        当 docstring 中的参数描述与 Pydantic Field(description=...) 冲突时，
        是否优先使用 docstring 的描述。
    append_raises_to_description : bool, 默认 True
        是否将 Raises 段落以 Markdown 格式追加到 operation description。
    docstring_style : DocstringStyle, 默认 AUTO
        解析风格。团队统一用 Google 时可显式传入 ``DocstringStyle.GOOGLE``。
    """

    # 可选择保存原始的 openapi 函数（如果有）
    # original_openapi = app.openapi

    def custom_openapi() -> Dict[str, Any]:
        # 如果已经有缓存，直接返回
        if app.openapi_schema:
            return app.openapi_schema

        # 生成基础 schema
        openapi_schema = get_openapi(
            title=app.title,
            version=app.version,
            openapi_version=app.openapi_version,
            description=app.description,
            routes=app.routes,
        )

        for route in app.routes:
            if not hasattr(route, "endpoint") or not route.endpoint:
                continue

            path = getattr(route, "path", None)
            methods = [m.lower() for m in route.methods] if hasattr(route, "methods") else []
            if not path or not methods:
                continue

            # 解析 docstring
            parsed = _parse_docstring(route.endpoint, style=docstring_style)
            if not parsed:
                continue

            param_lookup = _build_param_lookup(parsed)
            returns_desc = _build_returns_description(parsed)
            raises_desc = _build_raises_description(parsed)
            clean_desc = _get_clean_operation_description(parsed)

            for method in methods:
                if method == "head":
                    # FastAPI 通常不把 HEAD 单独写入 OpenAPI，或者与 GET 共享
                    continue
                if method not in openapi_schema.get("paths", {}).get(path, {}):
                    continue

                op = openapi_schema["paths"][path][method]

                # ---------- 1. 更新 operation description ----------
                if clean_desc:
                    op["description"] = clean_desc

                # ---------- 2. 注入参数描述（路径/查询/Header/Cookie）----------
                for param in op.get("parameters", []):
                    param_name = param.get("name")
                    if param_name and param_name in param_lookup:
                        existing = param.get("description", "")
                        new_desc = param_lookup[param_name]
                        if new_desc and (prefer_docstring_over_field or not existing):
                            param["description"] = new_desc

                # ---------- 3. 注入请求体 schema 字段描述 ----------
                if "requestBody" in op:
                    content = op["requestBody"].get("content", {})
                    for media_obj in content.values():
                        schema = media_obj.get("schema", {})
                        if schema:
                            _inject_schema_properties(schema, param_lookup, openapi_schema=openapi_schema)

                # ---------- 4. 注入返回值描述 ----------
                if returns_desc:
                    for code, resp in op.get("responses", {}).items():
                        if code.startswith("2"):  # 2xx 成功响应
                            # 更新响应级 description
                            resp["description"] = returns_desc
                            # 同时注入到响应 schema 的 description（如果存在）
                            for media_obj in resp.get("content", {}).values():
                                schema = media_obj.get("schema", {})
                                if schema and not schema.get("description"):
                                    schema["description"] = returns_desc
                                    # 递归注入 schema 内部字段
                                    _inject_schema_properties(schema, param_lookup, openapi_schema=openapi_schema)

                # ---------- 5. 异常描述追加 ----------
                if append_raises_to_description and raises_desc:
                    current_desc = op.get("description", "")
                    raises_md = f"## 异常\n\n{raises_desc}"
                    # 避免重复追加
                    if raises_md not in current_desc:
                        op["description"] = f"{current_desc}\n\n{raises_md}".strip()

        app.openapi_schema = openapi_schema
        return app.openapi_schema

    app.openapi = custom_openapi  # type: ignore[method-assign]
