#!/usr/bin/env python3
"""crawl.py — 把紫鸟 / HubStudio 的公开文档归档到本地（raw + markdown + manifest）。

设计要点
- 每个来源一个 group：raw/<group>/... 存原始响应，md/<group>/... 存提取后的 Markdown。
- 幂等/可续抓：某个 URL 的 md 已存在且非空就跳过（中断后重跑即可继续）。
- 礼貌抓取：默认 0.3s 间隔、2 次重试；不并发。
- manifest.tsv 记录 group / url / http / bytes / sha256 / title，便于之后核对完整性。

用法
  ./crawl.py --list               # 只看各 group 的任务数
  ./crawl.py --group zn-help      # 只抓某组
  ./crawl.py                      # 抓全部（groups 见 GROUPS）
"""

from __future__ import annotations

import argparse
import hashlib
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from html.parser import HTMLParser
from pathlib import Path

ROOT = Path(__file__).resolve().parent
UA = ('Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) '
      'Chrome/126.0.0.0 Safari/537.36')
DELAY = 0.3
TIMEOUT = 30


class Extract(HTMLParser):
    """极简 HTML→Markdown：只保留标题/段落/列表/表格/代码块的骨架。"""

    SKIP = {'script', 'style', 'noscript', 'svg', 'head', 'nav', 'footer', 'header', 'form', 'iframe'}
    BLOCK = {'p', 'div', 'section', 'article', 'li', 'tr', 'blockquote', 'pre', 'table'}
    HEAD = {'h1': 1, 'h2': 2, 'h3': 3, 'h4': 4, 'h5': 5, 'h6': 6}

    def __init__(self) -> None:
        super().__init__(convert_charrefs=True)
        self.parts: list[str] = []
        self.skip = 0
        self.in_title = False
        self.title = ''
        self.in_pre = False

    def handle_starttag(self, tag, attrs):
        if tag in self.SKIP:
            self.skip += 1
        elif tag == 'title':
            self.in_title = True
        elif tag == 'pre':
            self.in_pre = True
            self.parts.append('\n```\n')
        elif tag in self.HEAD:
            self.parts.append('\n\n' + '#' * self.HEAD[tag] + ' ')
        elif tag == 'li':
            self.parts.append('\n- ')
        elif tag in self.BLOCK:
            self.parts.append('\n\n')
        elif tag == 'br':
            self.parts.append('\n')
        elif tag == 'td' or tag == 'th':
            self.parts.append(' | ')

    def handle_endtag(self, tag):
        if tag in self.SKIP and self.skip:
            self.skip -= 1
        elif tag == 'title':
            self.in_title = False
        elif tag == 'pre':
            self.in_pre = False
            self.parts.append('\n```\n')
        elif tag in self.HEAD or tag in self.BLOCK:
            self.parts.append('\n')

    def handle_data(self, data):
        if self.in_title and not self.title:
            self.title = data.strip()
        if self.skip:
            return
        text = data if self.in_pre else re.sub(r'\s+', ' ', data)
        if text.strip():
            self.parts.append(text if self.in_pre else text.strip() + ' ')

    def markdown(self) -> str:
        t = ''.join(self.parts)
        t = re.sub(r'[ \t]+', ' ', t)
        t = re.sub(r' *\n *', '\n', t)
        t = re.sub(r'\n{3,}', '\n\n', t)
        return t.strip()


def fetch(url: str) -> tuple[int, bytes, str]:
    # 中文路径/查询必须百分号编码，否则 urllib 直接失败（实测 blog 的 /tag/中文 会返回 0）
    parts = urllib.parse.urlsplit(url)
    url = urllib.parse.urlunsplit((
        parts.scheme, parts.netloc,
        urllib.parse.quote(parts.path, safe='/%'),
        urllib.parse.quote(parts.query, safe='=&%'),
        parts.fragment))
    req = urllib.request.Request(url, headers={
        'User-Agent': UA, 'Accept': '*/*', 'Accept-Language': 'zh-CN,zh;q=0.9'})
    for attempt in range(3):
        try:
            with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
                return r.status, r.read(), r.headers.get('Content-Type', '')
        except urllib.error.HTTPError as e:
            if e.code in (403, 404, 410):
                return e.code, b'', ''
            time.sleep(1.5 * (attempt + 1))
        except Exception:  # noqa: BLE001
            time.sleep(1.5 * (attempt + 1))
    return 0, b'', ''


def safe_path(url: str) -> Path:
    p = urllib.parse.urlparse(url)
    path = p.path.strip('/') or 'index'
    if path.endswith('/'):
        path += 'index'
    path = re.sub(r'[^A-Za-z0-9._/-]', '_', path)
    return Path(p.netloc.replace(':', '_')) / (path if path.endswith(('.md', '.html')) else path + '.html')


def write(group: str, url: str, code: int, body: bytes, ctype: str) -> str:
    raw_dir, md_dir = ROOT / 'raw' / group, ROOT / 'md' / group
    rel = safe_path(url)
    raw_path = raw_dir / rel
    md_path = md_dir / rel.with_suffix('.md')
    raw_path.parent.mkdir(parents=True, exist_ok=True)
    md_path.parent.mkdir(parents=True, exist_ok=True)
    raw_path.write_bytes(body)

    if 'markdown' in ctype or url.endswith('.md'):
        md_text = body.decode('utf-8', 'replace')
        title = next((ln.lstrip('# ').strip() for ln in md_text.splitlines() if ln.startswith('#')), '')
    else:
        ex = Extract()
        try:
            ex.feed(body.decode('utf-8', 'replace'))
        except Exception:  # noqa: BLE001
            pass
        md_text = ex.markdown()
        title = ex.title

    header = f'<!-- source: {url} | fetched: {time.strftime("%Y-%m-%d %H:%M:%S")} | http: {code} -->\n\n'
    md_path.write_text(header + md_text, encoding='utf-8')
    sha = hashlib.sha256(body).hexdigest()
    with (ROOT / 'manifest.tsv').open('a', encoding='utf-8') as f:
        f.write(f'{group}\t{url}\t{code}\t{len(body)}\t{sha[:16]}\t{title[:80]}\n')
    return title[:60]


def load_urls(group: str) -> list[str]:
    f = ROOT / f'urls-{group}.txt'
    if not f.exists():
        return []
    return [ln.strip() for ln in f.read_text(encoding='utf-8').splitlines() if ln.strip()]


def discover_hs_support(limit: int = 600) -> list[str]:
    """HubStudio support 站没有 sitemap：从首页 BFS 发现 /xxxx/xxxx 形式的文章链接。"""
    seed = 'https://support-orig.hubstudio.cn/'
    seen: set[str] = set()
    queue = [seed]
    out: list[str] = []
    pat = re.compile(r'href="(/[0-9a-f]{4}/[0-9a-f]{4,6})"')
    while queue and len(out) < limit:
        u = queue.pop(0)
        if u in seen:
            continue
        seen.add(u)
        code, body, _ = fetch(u)
        if code != 200 or not body:
            continue
        out.append(u)
        html = body.decode('utf-8', 'replace')
        for m in pat.findall(html):
            full = urllib.parse.urljoin(seed, m)
            if full not in seen and full not in queue:
                queue.append(full)
        time.sleep(DELAY)
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument('--group', action='append', default=[])
    ap.add_argument('--list', action='store_true')
    args = ap.parse_args()

    groups = {
        'zn-help': lambda: load_urls('zn-help'),
        'zn-blog': lambda: load_urls('zn-blog'),
        'zn-news': lambda: load_urls('zn-news'),
        'hs-api': lambda: load_urls('hs-api'),
        'hs-support': discover_hs_support,
    }
    if args.list:
        for g, fn in groups.items():
            print(f'  {g:12s} {len(fn())} 条')
        return 0

    want = args.group or list(groups)
    total = 0
    for g in want:
        urls = groups[g]()
        (ROOT / f'urls-{g}.txt').write_text('\n'.join(urls) + '\n', encoding='utf-8')
        ok = skip = fail = 0
        print(f'== {g}: {len(urls)} 条 ==', flush=True)
        for i, u in enumerate(urls, 1):
            md = ROOT / 'md' / g / safe_path(u).with_suffix('.md')
            if md.exists() and md.stat().st_size > 200:
                skip += 1
                continue
            code, body, ctype = fetch(u)
            if code == 200 and body:
                title = write(g, u, code, body, ctype)
                ok += 1
                if i % 20 == 0:
                    print(f'   [{i}/{len(urls)}] {title[:48]}', flush=True)
            else:
                fail += 1
                with (ROOT / 'manifest.tsv').open('a', encoding='utf-8') as f:
                    f.write(f'{g}\t{u}\t{code}\t0\t-\t(fail)\n')
            time.sleep(DELAY)
        print(f'   -> ok={ok} skip={skip} fail={fail}', flush=True)
        total += ok
    print(f'完成，新抓 {total} 页', flush=True)
    return 0


if __name__ == '__main__':
    sys.exit(main())