ForcePilot/backend/package/yuxi/external_systems/integrations/notion/operations.py

"""Notion discover / preview_tools / create_tools handler 实现。

handler 签名遵循 ``OperationHandler`` 协议：
``(source_type: str, system_config: dict[str, Any]) -> Awaitable[Any]``。

handler 内部职责：
1. 从 ``system_config`` 提取 ``base_url`` / ``auth_config`` / ``_resolved_token``
2. 构建 ``auth_headers``（token 由 use_cases 层注入到 ``system_config["_resolved_token"]``）
3. 调用 Notion API 聚合元数据（search / databases / users）
4. 调用 ``NotionConfigGenerator`` 生成工具配置草稿

429 退避分层（见设计文档 §6.2）：
- discover/preview/create 阶段：本模块 ``_request_with_retry`` 实现（解析 Retry-After）
- 运行时工具执行：通过工具 ``retry_policy.retry_status_codes: [429, 502, 503, 504]`` 配置
"""

from __future__ import annotations

import asyncio
import time
from typing import TYPE_CHECKING, Any

from yuxi.external_systems.exceptions import (
    AuthError,
    DomainValidationError,
    ExecutionError,
    RateLimitExceededError,
)
from yuxi.external_systems.integrations.notion.constants import (
    _NOTION_HEADERS,
    _RETRY_BASE_DELAY,
    _RETRY_MAX_ATTEMPTS,
)
from yuxi.external_systems.integrations.notion.error_extractor import (
    map_notion_error,
)
from yuxi.external_systems.integrations.notion.generators import (
    NotionConfigGenerator,
)
from yuxi.external_systems.integrations.schemas import GeneratedToolsDraft

if TYPE_CHECKING:
    pass


# ── system_config 解析辅助 ────────────────────────────────────────────────


def _build_auth_headers(system_config: dict[str, Any]) -> dict[str, str]:
    """从 system_config 提取 token，构建 Notion 请求头。

    system_config 中 token 由 use_cases 层通过 ``token_manager.get_token`` 预处理后
    注入到 ``system_config["_resolved_token"]["access_token"]``。handler 不自行换取
    token——token 获取是 use_cases 层的职责。
    """
    token_info = system_config.get("_resolved_token")
    if not isinstance(token_info, dict) or not token_info.get("access_token"):
        raise AuthError("notion: system_config 缺少 _resolved_token，use_cases 层未注入 token")
    access_token = token_info["access_token"]
    headers = dict(_NOTION_HEADERS)
    headers["Authorization"] = f"Bearer {access_token}"
    return headers


def _resolve_base_url(system_config: dict[str, Any]) -> str:
    """从 system_config 提取 base_url。"""
    connection_config = system_config.get("connection_config") or {}
    base_url = connection_config.get("base_url")
    if not base_url:
        raise DomainValidationError("notion: system_config.connection_config.base_url 不能为空")
    return str(base_url).rstrip("/")


def _resolve_selected_tools(system_config: dict[str, Any]) -> list[str] | None:
    """从 system_config 提取用户选择的工具 slug 列表。"""
    discovery_options = system_config.get("discovery_options") or {}
    selected = discovery_options.get("selected_tools")
    return list(selected) if selected else None


# ── 429 退避 HTTP 请求 ────────────────────────────────────────────────────


async def _request_with_retry(
    base_url: str,
    endpoint: str,
    headers: dict[str, str],
    *,
    method_verb: str = "GET",
    params: dict[str, Any] | None = None,
    json_body: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """带 429 退避的 HTTP 请求（discover/preview/create 阶段使用）。

    解析 ``Retry-After`` 头部（秒数），按指数退避重试，最多 ``_RETRY_MAX_ATTEMPTS`` 次。
    其他 4xx/5xx 错误通过 ``map_notion_error`` 转换为 ``ExternalSystemError`` 子类抛出。
    网络层异常（超时 / 连接错误）转换为 ``ExecutionError``。

    Returns:
        解析后的 JSON 响应 dict。
    """
    import httpx

    url = f"{base_url}/{endpoint}"
    last_error: Exception | None = None

    for attempt in range(_RETRY_MAX_ATTEMPTS):
        try:
            async with httpx.AsyncClient() as client:
                if method_verb == "GET":
                    resp = await client.get(url, headers=headers, params=params)
                else:
                    resp = await client.request(method_verb, url, headers=headers, json=json_body)

            if resp.status_code == 429:
                if attempt == _RETRY_MAX_ATTEMPTS - 1:
                    retry_after_header = resp.headers.get("Retry-After", "60")
                    try:
                        retry_after = int(retry_after_header)
                    except ValueError:
                        retry_after = 60
                    raise RateLimitExceededError(
                        f"notion: 限流，已重试 {_RETRY_MAX_ATTEMPTS} 次仍失败: {url}",
                        retry_after=retry_after,
                        limit_type="qps",
                    )
                retry_after_header = resp.headers.get("Retry-After")
                if retry_after_header:
                    try:
                        delay = float(retry_after_header)
                    except ValueError:
                        delay = _RETRY_BASE_DELAY * (2**attempt)
                else:
                    delay = _RETRY_BASE_DELAY * (2**attempt)
                await asyncio.sleep(delay)
                continue

            if resp.status_code >= 400:
                try:
                    error_body = resp.json()
                except Exception:
                    error_body = {}
                raise map_notion_error(resp.status_code, error_body)

            try:
                return resp.json()
            except Exception as exc:
                raise ExecutionError(f"notion: 响应 JSON 解析失败: {url}: {exc}") from exc

        except httpx.HTTPError as exc:
            last_error = exc
            if attempt == _RETRY_MAX_ATTEMPTS - 1:
                break
            await asyncio.sleep(_RETRY_BASE_DELAY * (2**attempt))
            continue

    raise ExecutionError(f"notion: 请求 {url} 失败，已重试 {_RETRY_MAX_ATTEMPTS} 次: {last_error}")


# ── cursor 分页自动迭代 ──────────────────────────────────────────────────


async def _paginate(
    base_url: str,
    endpoint: str,
    headers: dict[str, str],
    params: dict[str, Any],
) -> list[dict[str, Any]]:
    """cursor 分页自动迭代，聚合所有页结果（GET 请求）。

    Args:
        endpoint: API 端点（如 ``users``）
        params: 初始查询参数（不含 ``start_cursor``）

    Returns:
        聚合后的资源列表
    """
    aggregated: list[dict[str, Any]] = []
    cursor: str | None = None
    while True:
        query = {**params}
        if cursor:
            query["start_cursor"] = cursor
        resp = await _request_with_retry(
            base_url,
            endpoint,
            headers,
            method_verb="GET",
            params=query,
        )
        results = resp.get("results", [])
        if isinstance(results, list):
            aggregated.extend(results)
        if not resp.get("has_more"):
            break
        cursor = resp.get("next_cursor")
        if not cursor:
            break
    return aggregated


async def _paginate_search(
    base_url: str,
    headers: dict[str, str],
    payload: dict[str, Any],
) -> list[dict[str, Any]]:
    """search 接口专用分页（POST + JSON body）。"""
    aggregated: list[dict[str, Any]] = []
    cursor: str | None = None
    while True:
        body = {**payload}
        if cursor:
            body["start_cursor"] = cursor
        resp = await _request_with_retry(
            base_url,
            "search",
            headers,
            method_verb="POST",
            json_body=body,
        )
        results = resp.get("results", [])
        if isinstance(results, list):
            aggregated.extend(results)
        if not resp.get("has_more"):
            break
        cursor = resp.get("next_cursor")
        if not cursor:
            break
    return aggregated


# ── Notion 资源标题提取 ──────────────────────────────────────────────────


def _extract_database_title(db: dict[str, Any]) -> str:
    """从 Notion database 对象提取标题（title 字段为 rich_text 数组）。"""
    title_list = db.get("title", [])
    if isinstance(title_list, list):
        return "".join(t.get("plain_text", "") for t in title_list if isinstance(t, dict))
    return ""


def _extract_page_title(page: dict[str, Any]) -> str:
    """从 Notion page 对象提取标题（properties 中 type=title 的属性）。"""
    properties = page.get("properties", {})
    if not isinstance(properties, dict):
        return ""
    for prop in properties.values():
        if not isinstance(prop, dict):
            continue
        if prop.get("type") == "title":
            title_list = prop.get("title", [])
            if isinstance(title_list, list):
                return "".join(t.get("plain_text", "") for t in title_list if isinstance(t, dict))
    return ""


# ── 元数据发现辅助 ────────────────────────────────────────────────────────


async def _discover_databases_with_properties(
    base_url: str,
    headers: dict[str, str],
) -> list[dict[str, Any]]:
    """发现数据库并获取每个数据库的完整结构（含 properties）。

    用于 preview/create 阶段：先通过 search 获取数据库列表，
    再对每个数据库调用 ``GET /v1/databases/{id}`` 获取完整 properties。
    """
    databases = await _paginate_search(
        base_url,
        headers,
        payload={"page_size": 100, "filter": {"value": "database", "property": "object"}},
    )
    result: list[dict[str, Any]] = []
    for db in databases:
        if not isinstance(db, dict):
            continue
        db_id = db.get("id")
        if not db_id:
            continue
        resp = await _request_with_retry(
            base_url,
            f"databases/{db_id}",
            headers,
            method_verb="GET",
        )
        result.append(
            {
                "id": db_id,
                "title": _extract_database_title(db),
                "properties": resp.get("properties", {}) if isinstance(resp, dict) else {},
            }
        )
    return result


# ── 三个 OperationHandler 实现 ────────────────────────────────────────────


async def discover_notion(
    source_type: str,
    system_config: dict[str, Any],
) -> list[dict[str, Any]]:
    """discover handler：发现 Notion 工作区可用资源（数据库 + 页面 + 用户）。

    Returns:
        资源列表，每项形如：
        - ``{"type": "database", "id": "...", "title": "...", "url": "...", "parent_type": "..."}``
        - ``{"type": "page", "id": "...", "title": "...", "url": "...", "parent_type": "..."}``
        - ``{"type": "user", "id": "...", "name": "...", "type_detail": "...", "email": "..."}``
    """
    base_url = _resolve_base_url(system_config)
    headers = _build_auth_headers(system_config)

    # 通过 search 接口聚合所有可访问的 database 与 page
    databases = await _paginate_search(
        base_url,
        headers,
        payload={"page_size": 100, "filter": {"value": "database", "property": "object"}},
    )
    pages = await _paginate_search(
        base_url,
        headers,
        payload={"page_size": 100, "filter": {"value": "page", "property": "object"}},
    )
    users = await _paginate(base_url, "users", headers, params={"page_size": 100})

    return [
        *[
            {
                "type": "database",
                "id": r.get("id", ""),
                "title": _extract_database_title(r),
                "url": r.get("url"),
                "parent_type": (r.get("parent") or {}).get("type"),
            }
            for r in databases
            if isinstance(r, dict)
        ],
        *[
            {
                "type": "page",
                "id": r.get("id", ""),
                "title": _extract_page_title(r),
                "url": r.get("url"),
                "parent_type": (r.get("parent") or {}).get("type"),
            }
            for r in pages
            if isinstance(r, dict)
        ],
        *[
            {
                "type": "user",
                "id": u.get("id", ""),
                "name": u.get("name"),
                "type_detail": u.get("type"),
                "email": (u.get("person") or {}).get("email") if isinstance(u.get("person"), dict) else None,
            }
            for u in users
            if isinstance(u, dict)
        ],
    ]


async def preview_notion_tools(
    source_type: str,
    system_config: dict[str, Any],
) -> list[dict[str, Any]]:
    """preview_tools handler：预览将生成的工具列表。

    Returns:
        工具预览列表，每项为 ExternalToolCreateInput 兼容的 dict。
    """
    payload = await _build_generator_payload(system_config)
    generator = NotionConfigGenerator()
    return await generator.generate(payload)


async def create_notion_tools(
    source_type: str,
    system_config: dict[str, Any],
) -> GeneratedToolsDraft:
    """create_tools handler：生成工具草稿（不直接持久化）。

    Returns:
        GeneratedToolsDraft，由 use_cases 层统一持久化。
    """
    payload = await _build_generator_payload(system_config)
    generator = NotionConfigGenerator()
    tool_configs = await generator.generate(payload)
    return GeneratedToolsDraft(tool_configs=tool_configs, override_existing=False)


# ── 生成器 payload 构造 ───────────────────────────────────────────────────


async def _build_generator_payload(system_config: dict[str, Any]) -> dict[str, Any]:
    """从 system_config 构造生成器 payload。

    流程：
    1. 解析 base_url / auth_headers / auth_type / auth_config
    2. 解析用户选择的工具 slug 列表（可选）
    3. 发现数据库并获取完整结构（含 properties）
    4. 汇总为生成器 payload
    """
    base_url = _resolve_base_url(system_config)
    auth_headers = _build_auth_headers(system_config)
    auth_config = system_config.get("auth_config") or {}
    auth_type = system_config.get("auth_type", "bearer")
    selected_tools = _resolve_selected_tools(system_config)

    discovered_databases = await _discover_databases_with_properties(base_url, auth_headers)

    payload: dict[str, Any] = {
        "base_url": base_url,
        "auth_headers": auth_headers,
        "auth_config": auth_config,
        "auth_type": auth_type,
        "discovered_databases": discovered_databases,
    }
    if selected_tools is not None:
        payload["selected_tools"] = selected_tools
    return payload


# ── 健康检查 ─────────────────────────────────────────────────────────────


async def health_check_notion(system_config: dict[str, Any]) -> dict[str, Any]:
    """Notion 健康检查：调用 ``GET /v1/users/me`` 验证 Token 有效性。

    供 use_cases 层直接调用，不注册到 ``IntegrationOperationRegistry``。

    Returns::

        {
            "status": "healthy" | "unhealthy",
            "bot_id": "...",
            "workspace_name": "Acme",
            "duration_ms": 123,
        }
    """
    start = time.monotonic()
    try:
        base_url = _resolve_base_url(system_config)
        headers = _build_auth_headers(system_config)
        resp = await _request_with_retry(
            base_url,
            "users/me",
            headers,
            method_verb="GET",
        )
        bot_id: str | None = None
        if resp.get("type") == "bot":
            bot = resp.get("bot") or {}
            owner = bot.get("owner") or {}
            user = owner.get("user") or {}
            bot_id = user.get("id") if isinstance(user, dict) else None
        return {
            "status": "healthy",
            "bot_id": bot_id,
            "workspace_name": resp.get("workspace_name"),
            "duration_ms": int((time.monotonic() - start) * 1000),
        }
    except Exception as exc:
        return {
            "status": "unhealthy",
            "detail": str(exc),
            "duration_ms": int((time.monotonic() - start) * 1000),
        }