crawl_docs.py

# -*- coding: utf-8 -*-
"""
Pico XR 开发文档爬虫
爬取 https://developer-cn.picoxr.com/document/ 下所有开发文档
并保存为结构化 JSON 数据库
"""

import sys
import io
# 强制 stdout/stderr 使用 UTF-8（Windows GBK 环境下需要）
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')

import asyncio
import json
import re
import time
from pathlib import Path
from urllib.parse import urljoin, urlparse

from crawl4ai import AsyncWebCrawler, BrowserConfig, CrawlerRunConfig, CacheMode

BASE_URL = "https://developer-cn.picoxr.com"
START_URL = "https://developer-cn.picoxr.com/document/#develop"
DATA_DIR = Path(__file__).parent / "data"
DATA_DIR.mkdir(exist_ok=True)

# 文档页面的 URL 前缀白名单（只爬这些路径）
ALLOWED_PREFIXES = [
    "/document/unity/",
    "/document/unreal/",
    "/document/native/",
    "/document/openxr/",
    "/document/discover/",
    "/document/design/",
    "/document/distribute/",
    "/document/safe-info/",
]

# 已知要爬的分类入口页
CATEGORY_ENTRIES = [
    # Unity SDK
    "https://developer-cn.picoxr.com/document/unity/",
    # Unreal SDK
    "https://developer-cn.picoxr.com/document/unreal/",
    # Native SDK
    "https://developer-cn.picoxr.com/document/native/",
    # OpenXR
    "https://developer-cn.picoxr.com/document/openxr/",
    # 设备介绍
    "https://developer-cn.picoxr.com/document/discover/",
    # 设计指南
    "https://developer-cn.picoxr.com/document/design/",
    # 分发上架
    "https://developer-cn.picoxr.com/document/distribute/",
    # 安全信息
    "https://developer-cn.picoxr.com/document/safe-info/",
]


def is_allowed_url(url: str) -> bool:
    """判断 URL 是否在允许爬取的范围内"""
    parsed = urlparse(url)
    if parsed.netloc and parsed.netloc != "developer-cn.picoxr.com":
        return False
    path = parsed.path
    return any(path.startswith(prefix) for prefix in ALLOWED_PREFIXES)


def normalize_url(url: str, base: str = BASE_URL) -> str:
    """规范化 URL"""
    if url.startswith("http"):
        return url.split("#")[0].rstrip("/")
    full = urljoin(base, url)
    return full.split("#")[0].rstrip("/")


async def crawl_page(crawler: AsyncWebCrawler, url: str, config: CrawlerRunConfig) -> dict | None:
    """爬取单个页面，返回文档数据"""
    try:
        result = await crawler.arun(url=url, config=config)
        if not result.success:
            print(f"  [ERR] 失败: {url} - {result.error_message}")
            return None

        # 提取内部链接
        internal_links = []
        for link in result.links.get("internal", []):
            href = link.get("href", "")
            if href:
                normalized = normalize_url(href)
                if is_allowed_url(normalized):
                    internal_links.append(normalized)

        # 提取页面标题
        title = result.metadata.get("title", "") or ""
        title = title.replace(" - PICO 开发者平台", "").replace(" | PICO Developer", "").strip()

        # 获取 markdown 内容
        markdown = ""
        if hasattr(result.markdown, "fit_markdown") and result.markdown.fit_markdown:
            markdown = result.markdown.fit_markdown
        elif hasattr(result.markdown, "raw_markdown"):
            markdown = result.markdown.raw_markdown
        else:
            markdown = str(result.markdown)

        # 清理 markdown
        markdown = clean_markdown(markdown)

        return {
            "url": url,
            "title": title,
            "markdown": markdown,
            "links": list(set(internal_links)),
            "crawled_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
        }
    except Exception as e:
        print(f"  [ERR] 异常: {url} - {e}")
        return None


def clean_markdown(text: str) -> str:
    """清理 markdown 内容"""
    if not text:
        return ""
    # 移除过多的空行
    text = re.sub(r"\n{4,}", "\n\n\n", text)
    # 移除纯导航链接行
    lines = text.split("\n")
    cleaned = []
    for line in lines:
        # 跳过只有导航链接的行
        stripped = line.strip()
        if stripped and len(stripped) < 3:
            continue
        cleaned.append(line)
    return "\n".join(cleaned).strip()


async def discover_all_urls(crawler: AsyncWebCrawler) -> set[str]:
    """发现所有文档 URL"""
    discovered = set()
    to_visit = set(CATEGORY_ENTRIES)
    visited = set()

    config = CrawlerRunConfig(
        cache_mode=CacheMode.ENABLED,
        page_timeout=30000,
        exclude_external_links=True,
        css_selector="body",  # 获取整个页面以发现链接
    )

    print(f"🔍 开始发现文档链接，入口页: {len(to_visit)} 个")

    # BFS 发现所有链接（最多 3 层深度）
    for depth in range(3):
        if not to_visit:
            break
        print(f"\n📡 第 {depth + 1} 层，待访问: {len(to_visit)} 个 URL")

        # 批量爬取
        results = await crawler.arun_many(
            urls=list(to_visit),
            config=config,
        )

        next_to_visit = set()
        for result in results:
            if not result.success:
                continue
            url = result.url.split("#")[0].rstrip("/")
            visited.add(url)
            discovered.add(url)

            # 收集内部链接
            for link in result.links.get("internal", []):
                href = link.get("href", "")
                if not href:
                    continue
                normalized = normalize_url(href)
                if is_allowed_url(normalized) and normalized not in visited:
                    next_to_visit.add(normalized)

        to_visit = next_to_visit - visited
        print(f"  [OK] 已发现: {len(discovered)} 个，新增待访问: {len(to_visit)} 个")

    print(f"\n🎯 共发现 {len(discovered)} 个文档 URL")
    return discovered


async def crawl_all_docs():
    """爬取所有文档"""
    browser_config = BrowserConfig(
        headless=True,
        viewport_width=1920,
        viewport_height=1080,
    )

    crawl_config = CrawlerRunConfig(
        cache_mode=CacheMode.ENABLED,
        page_timeout=30000,
        remove_overlay_elements=True,
        exclude_external_links=True,
        excluded_tags=["nav", "footer", "header", "aside", ".sidebar", "#sidebar"],
        css_selector=".doc-content, main, article, .content-wrapper, #content, body",
    )

    all_docs = []
    failed_urls = []

    async with AsyncWebCrawler(config=browser_config) as crawler:
        # 第一步：发现所有 URL
        print("=" * 60)
        print("🚀 Pico XR 开发文档爬虫启动")
        print("=" * 60)

        all_urls = await discover_all_urls(crawler)

        # 保存 URL 列表
        urls_file = DATA_DIR / "discovered_urls.json"
        with open(urls_file, "w", encoding="utf-8") as f:
            json.dump(sorted(all_urls), f, ensure_ascii=False, indent=2)
        print(f"\n💾 URL 列表已保存到: {urls_file}")

        # 第二步：爬取所有页面内容
        print(f"\n📖 开始爬取 {len(all_urls)} 个页面内容...")
        print("=" * 60)

        urls_list = sorted(all_urls)
        batch_size = 5  # 每批并发 5 个

        for i in range(0, len(urls_list), batch_size):
            batch = urls_list[i : i + batch_size]
            batch_num = i // batch_size + 1
            total_batches = (len(urls_list) + batch_size - 1) // batch_size
            print(f"\n📦 批次 {batch_num}/{total_batches}: {len(batch)} 个页面")

            results = await crawler.arun_many(
                urls=batch,
                config=crawl_config,
            )

            for result in results:
                if not result.success:
                    failed_urls.append(result.url)
                    print(f"  [ERR] {result.url}")
                    continue

                url = result.url.split("#")[0].rstrip("/")
                title = result.metadata.get("title", "") or ""
                title = title.replace(" - PICO 开发者平台", "").replace(" | PICO Developer", "").strip()

                markdown = ""
                if hasattr(result.markdown, "fit_markdown") and result.markdown.fit_markdown:
                    markdown = result.markdown.fit_markdown
                elif hasattr(result.markdown, "raw_markdown"):
                    markdown = result.markdown.raw_markdown
                else:
                    markdown = str(result.markdown)

                markdown = clean_markdown(markdown)

                if not markdown or len(markdown) < 50:
                    print(f"  [WARN]  内容太少: {url} ({len(markdown)} chars)")
                    continue

                # 生成文档 ID（从 URL 路径派生）
                path_parts = urlparse(url).path.strip("/").split("/")
                doc_id = "_".join(path_parts[-3:]) if len(path_parts) >= 3 else "_".join(path_parts)

                doc = {
                    "id": doc_id,
                    "url": url,
                    "title": title,
                    "category": get_category(url),
                    "markdown": markdown,
                    "crawled_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
                }
                all_docs.append(doc)
                print(f"  [OK] [{len(all_docs)}] {title[:50]} ({len(markdown)} chars)")

            # 保存中间结果
            if i % (batch_size * 10) == 0 and all_docs:
                save_docs(all_docs)

            await asyncio.sleep(1)  # 礼貌延迟

    # 最终保存
    save_docs(all_docs)

    # 保存失败列表
    if failed_urls:
        with open(DATA_DIR / "failed_urls.json", "w", encoding="utf-8") as f:
            json.dump(failed_urls, f, ensure_ascii=False, indent=2)
        print(f"\n[WARN]  {len(failed_urls)} 个页面爬取失败，已保存到 failed_urls.json")

    print("\n" + "=" * 60)
    print(f"🎉 爬取完成！共获取 {len(all_docs)} 篇文档")
    print(f"💾 数据已保存到: {DATA_DIR}")
    print("=" * 60)

    return all_docs


def get_category(url: str) -> str:
    """从 URL 推断文档分类"""
    path = urlparse(url).path
    if "/unity/" in path:
        return "Unity SDK"
    elif "/unreal/" in path:
        return "Unreal SDK"
    elif "/native/" in path:
        return "Native SDK"
    elif "/openxr/" in path:
        return "OpenXR"
    elif "/discover/" in path:
        return "设备介绍"
    elif "/design/" in path:
        return "设计指南"
    elif "/distribute/" in path:
        return "分发上架"
    elif "/safe-info/" in path:
        return "安全信息"
    else:
        return "其他"


def save_docs(docs: list[dict]):
    """保存文档到 JSON 文件"""
    # 全量数据（含 markdown）
    full_path = DATA_DIR / "pico_docs_full.json"
    with open(full_path, "w", encoding="utf-8") as f:
        json.dump(docs, f, ensure_ascii=False, indent=2)

    # 索引数据（不含 markdown，用于快速检索）
    index = [
        {
            "id": doc["id"],
            "url": doc["url"],
            "title": doc["title"],
            "category": doc["category"],
            "summary": doc["markdown"][:300] + "..." if len(doc["markdown"]) > 300 else doc["markdown"],
        }
        for doc in docs
    ]
    index_path = DATA_DIR / "pico_docs_index.json"
    with open(index_path, "w", encoding="utf-8") as f:
        json.dump(index, f, ensure_ascii=False, indent=2)

    print(f"\n💾 已保存 {len(docs)} 篇文档 -> {full_path}")


if __name__ == "__main__":
    asyncio.run(crawl_all_docs())