joke/crawler/site_crawler.py

"""
单站点深度采集工具。
从 JSON 文件按编号加载站点，或直接指定 URL，深度采集该站点所有笑话。

用法:
  python crawler/site_crawler.py --id 1                    # 采集 site_finder 发现的 #1 站点
  python crawler/site_crawler.py --id 1,2,3               # 批量采集多个站点
  python crawler/site_crawler.py --url https://...         # 直接采集指定 URL
  python crawler/site_crawler.py --id 1 --no-headless      # 显示浏览器窗口
"""
import argparse
import asyncio
import json
import os
import sys

if sys.platform == "win32":
    import shutil
    os.environ["PYTHONIOENCODING"] = "utf-8"
    os.environ["TERM"] = "dumb"
    try:
        shutil.get_terminal_size()
    except Exception:
        pass

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from crawler.processor import Processor


def parse_args():
    parser = argparse.ArgumentParser(description="单站点深度采集工具")
    parser.add_argument(
        "--id",
        type=str,
        default=None,
        help="站点编号（从 site_finder 生成的 JSON 读取），多个用逗号分隔",
    )
    parser.add_argument(
        "--url",
        type=str,
        default=None,
        help="直接指定站点 URL（与 --id 二选一）",
    )
    parser.add_argument(
        "--sites-file",
        type=str,
        default="joke_sites.json",
        help="站点列表 JSON 文件路径（默认: joke_sites.json）",
    )
    parser.add_argument(
        "--api-base",
        type=str,
        default=os.getenv("API_BASE", "http://localhost:8001"),
        help="API 服务地址",
    )
    parser.add_argument(
        "--username",
        type=str,
        default=os.getenv("CRAWL_USERNAME", "admin"),
        help="管理员用户名",
    )
    parser.add_argument(
        "--password",
        type=str,
        default=os.getenv("CRAWL_PASSWORD", "admin123"),
        help="管理员密码",
    )
    parser.add_argument(
        "--no-headless",
        action="store_true",
        help="显示浏览器窗口",
    )
    return parser.parse_args()


def load_sites(path: str) -> list[dict]:
    """从 JSON 文件加载站点列表"""
    if not os.path.exists(path):
        print(f"错误: 站点文件 {path} 不存在，请先运行 site_finder.py")
        sys.exit(1)
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        if not data:
            print(f"错误: 站点文件 {path} 为空")
            sys.exit(1)
        return data
    except Exception as e:
        print(f"错误: 读取站点文件失败: {e}")
        sys.exit(1)


def resolve_sites(args) -> list[str]:
    """解析 --id 或 --url 参数，返回待采集的 URL 列表"""
    if args.url:
        return [args.url]

    if not args.id:
        print("错误: 请指定 --id 或 --url")
        print("  例如: python crawler/site_crawler.py --id 1")
        print("  例如: python crawler/site_crawler.py --url https://xiaohua.com")
        sys.exit(1)

    # 解析编号列表 "1,2,3" → [1, 2, 3]
    try:
        ids = [int(x.strip()) for x in args.id.split(",") if x.strip()]
    except ValueError:
        print("错误: --id 参数必须是数字，多个用逗号分隔")
        sys.exit(1)

    sites = load_sites(args.sites_file)
    found = []
    for sid in ids:
        match = [s for s in sites if s["id"] == sid]
        if match:
            found.append(match[0])
            print(f"  [*] 站点 #{sid}: {match[0]['domain']} — {match[0]['title'][:40]}")
        else:
            print(f"  [!] 站点 #{sid} 未找到（可用编号: {[s['id'] for s in sites[:10]]}...)")

    if not found:
        print("错误: 没有找到有效的站点编号")
        sys.exit(1)

    return [s["url"] for s in found]


def main():
    args = parse_args()
    headless = not args.no_headless
    urls = resolve_sites(args)

    print(f"=" * 50)
    print(f"单站点深度采集工具")
    print(f"目标站点: {len(urls)} 个")
    for u in urls:
        print(f"  - {u}")
    print(f"API 地址: {args.api_base}")
    print(f"浏览器: {'显示窗口' if not headless else '无头模式'}")
    print(f"=" * 50)

    try:
        import crawl4ai
        print(f"crawl4ai 版本: {crawl4ai.__version__}")
    except ImportError:
        print("错误: crawl4ai 未安装，请先运行: pip install crawl4ai")
        return

    processor = Processor(
        api_base=args.api_base,
        username=args.username,
        password=args.password,
        headless=headless,
    )

    try:
        for url in urls:
            asyncio.run(processor.crawl_site(url))
    except KeyboardInterrupt:
        print("\n用户中断，退出")


if __name__ == "__main__":
    main()