#!/usr/bin/env python3
"""fetch_source.py — 按 sources/*.yaml 自动抓取 L2 原始数据到本地 staging。

只处理**可自动化**的 fetch_method：
  - curl       : 逐 item 下载单文件（item.url，或 source.url_template 渲染 {item_id}）
  - git-clone  : 浅克隆一次仓库，再按 item.src_path 打包成 .tar.gz

manual-download / manual-export / generated 会被跳过（需人工介入）。

用法：
  # 看有哪些可自动抓取（不下载）
  python tools/fetch_source.py --list

  # 试跑：只打印将要下载的 URL / 目标路径
  python tools/fetch_source.py --source phet --dry-run

  # 实际下载到 .staging/，并回填 sha256/size/fetch_date/status=fetched
  python tools/fetch_source.py --source phet --apply
  python tools/fetch_source.py --source chinese-poetry --apply
  python tools/fetch_source.py --subject physics --apply     # 按学科批量
  python tools/fetch_source.py --all-auto --apply            # 所有可自动化来源

下载产物落在 <repo>/.staging/{source_id}/...（已 gitignore），随后用
  python tools/oss_sync.py --raw-root .staging --apply
上传到 OSS 并对齐 oss_key。

注意：回填 YAML 走 yaml_patch（行级替换），**保留文件内注释**。
"""
import os, sys, glob, hashlib, datetime, argparse, subprocess, shutil, tarfile, urllib.request

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from yaml_patch import patch_item_fields, patch_source_fields

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
STAGING = os.path.join(REPO, ".staging")
CACHE = os.path.join(REPO, ".cache", "repos")
AUTO_METHODS = {"curl", "git-clone"}
UA = "edu-content-fetcher/1.0 (+https://git.aibeaver.cn/fengjunchen/edu-content)"


def sha256_file(p):
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def load_sources():
    import yaml
    out = []
    for path in sorted(glob.glob(os.path.join(REPO, "sources", "**", "*.yaml"), recursive=True)):
        try:
            doc = yaml.safe_load(open(path, encoding="utf-8"))
        except Exception as e:
            print(f"WARN: {path} 解析失败: {e}")
            continue
        for s in (doc or {}).get("sources", []) or []:
            out.append((path, s))
    return out


def item_url(source, item):
    if item.get("url"):
        return item["url"]
    tpl = source.get("url_template")
    if tpl:
        return tpl.replace("{item_id}", item["id"])
    return None


def fetch_curl(source, item, dest_dir, dry_run):
    url = item_url(source, item)
    if not url:
        return None, "无 url / url_template，跳过"
    fname = os.path.basename(item.get("oss_key") or "") or f"{item['id']}"
    dest = os.path.join(dest_dir, fname)
    if dry_run:
        return dest, f"[dry-run] {url} -> {dest}"
    os.makedirs(dest_dir, exist_ok=True)
    req = urllib.request.Request(url, headers={"User-Agent": UA})
    try:
        with urllib.request.urlopen(req, timeout=120) as r, open(dest, "wb") as f:
            shutil.copyfileobj(r, f)
    except Exception as e:
        return None, f"下载失败: {type(e).__name__}: {e}"
    return dest, f"ok {os.path.getsize(dest)}B"


def ensure_clone(source, dry_run):
    """浅克隆到 .cache/repos/{id}，返回本地路径。"""
    url = source.get("url")
    if not url:
        return None, "source.url 为空，无法 clone"
    local = os.path.join(CACHE, source["id"])
    if os.path.isdir(os.path.join(local, ".git")):
        return local, "复用已有克隆"
    if dry_run:
        return local, f"[dry-run] git clone --depth 1 {url} {local}"
    os.makedirs(os.path.dirname(local), exist_ok=True)
    r = subprocess.run(["git", "clone", "--depth", "1", url, local],
                       capture_output=True, text=True, timeout=1800)
    if r.returncode != 0:
        return None, f"clone 失败: {(r.stderr or '').strip()[:200]}"
    return local, "克隆完成"


def fetch_git_item(source, item, repo_dir, dest_dir, dry_run):
    """按 item.src_path 从克隆仓库取内容。

    目录 → 打包 .tar.gz；单文件 → 直接复制（除非 oss_key 显式要 .tar.gz）。
    否则会产出「名叫 KET.json 其实是 tar.gz」的坑数据。
    """
    rel = item.get("src_path")
    if not rel:
        return None, "缺 src_path，跳过"
    src = os.path.join(repo_dir, rel)
    fname = os.path.basename(item.get("oss_key") or f"{item['id']}.tar.gz")
    dest = os.path.join(dest_dir, fname)
    want_tar = fname.endswith((".tar.gz", ".tgz"))
    if dry_run:
        action = "tar" if want_tar else "copy"
        return dest, f"[dry-run] {action} {rel} -> {dest}"
    if not os.path.exists(src):
        return None, f"仓库内路径不存在: {rel}"
    if os.path.isdir(src) and not want_tar:
        return None, f"src_path 是目录但 oss_key 不是 .tar.gz（{fname}），请修正 oss_key"
    os.makedirs(dest_dir, exist_ok=True)
    if want_tar:
        with tarfile.open(dest, "w:gz") as tf:
            tf.add(src, arcname=rel)
    else:
        shutil.copy2(src, dest)
    return dest, f"ok {os.path.getsize(dest)}B"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--source", help="按 source id 抓取")
    ap.add_argument("--subject", help="按学科代码抓取（chinese/physics/science/...）")
    ap.add_argument("--all-auto", action="store_true", help="所有 fetch_method 可自动化的来源")
    ap.add_argument("--list", action="store_true", help="只列出可自动抓取的来源")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--apply", action="store_true")
    ap.add_argument("--limit", type=int, default=0, help="每个 source 最多抓 N 个条目（调试用）")
    args = ap.parse_args()

    all_sources = load_sources()

    if args.list:
        print(f"{'source':22} {'method':12} {'items':>5}  subjects")
        print("-" * 70)
        for path, s in all_sources:
            mark = "" if s.get("fetch_method") in AUTO_METHODS else "  (需人工)"
            print(f"{s['id']:22} {str(s.get('fetch_method')):12} "
                  f"{len(s.get('items') or []):5}  {','.join(s.get('subjects') or [])}{mark}")
        return

    # 选择目标
    targets = []
    for path, s in all_sources:
        if s.get("fetch_method") not in AUTO_METHODS:
            continue
        if args.source and s["id"] != args.source:
            continue
        if args.subject and args.subject not in (s.get("subjects") or []):
            continue
        if not (args.source or args.subject or args.all_auto):
            continue
        targets.append((path, s))

    if not targets:
        print("没有匹配的可自动抓取来源。用 --list 查看，或检查 --source/--subject 拼写。")
        sys.exit(0 if (args.source or args.subject or args.all_auto) else 2)

    dry = args.dry_run or not args.apply
    today = datetime.date.today().isoformat()
    total_ok = total_fail = 0

    for path, s in targets:
        sid = s["id"]
        method = s["fetch_method"]
        items = s.get("items") or []
        if args.limit:
            items = items[:args.limit]
        print(f"\n=== {sid} ({method}) — {len(items)} 个条目 ===")
        if not items:
            print("  items 为空，无可抓取内容（该来源尚未登记具体条目）")
            continue

        dest_dir = os.path.join(STAGING, sid)
        repo_dir = None
        if method == "git-clone":
            repo_dir, msg = ensure_clone(s, dry)
            print(f"  clone: {msg}")
            if repo_dir is None:
                total_fail += len(items)
                continue

        ok_here = 0
        for it in items:
            if method == "curl":
                dest, msg = fetch_curl(s, it, dest_dir, dry)
            else:
                dest, msg = fetch_git_item(s, it, repo_dir, dest_dir, dry)
            print(f"  {it['id']:34} {msg}")
            if dest is None:
                total_fail += 1
                continue
            total_ok += 1
            ok_here += 1
            if not dry and os.path.isfile(dest):
                patch_item_fields(path, sid, it["id"], {
                    "sha256": sha256_file(dest),
                    "size": os.path.getsize(dest),
                    "fetch_date": today,
                    "status": "fetched",
                })
        if not dry and ok_here:
            patch_source_fields(path, sid, {"status": "fetched", "fetched_at": today})

    print(f"\n合计：成功 {total_ok} | 失败/跳过 {total_fail}")
    if dry:
        print("[dry-run] 未实际下载。加 --apply 执行。")
    else:
        print(f"产物在 {STAGING}/，已回填 sha256/size 到 sources YAML（注释保留）。")
        print("下一步上传 OSS： python tools/oss_sync.py --raw-root .staging --apply")


if __name__ == "__main__":
    main()