#!/usr/bin/env python3
"""按目录递归下载 HF 仓库文件到本地（走 hf-mirror）。

Piper 中文 medium 档没有 lexicon.txt，改用 espeak-ng-data 做音素化，
而那是一堆小文件，用 shell 逐个 curl 太啰嗦，所以写这个。

用法：
    python3 hf_fetch.py <repo> <本地目录> [远端前缀]
    python3 hf_fetch.py csukuangfj/vits-piper-zh_CN-huayan-medium \
        models/tts_piper_medium espeak-ng-data
"""

import json
import os
import sys
import urllib.parse
import urllib.request

MIRROR = os.environ.get("HF_MIRROR", "https://hf-mirror.com")
UA = {"User-Agent": "curl/8.5.0"}  # hf-mirror 对 urllib 默认 UA 返回 403


def _open(url, timeout):
    return urllib.request.urlopen(urllib.request.Request(url, headers=UA), timeout=timeout)


def api_tree(repo):
    url = f"{MIRROR}/api/models/{repo}/tree/main?recursive=true"
    with _open(url, 60) as r:
        return json.load(r)


def fetch(url, dest):
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    tmp = dest + ".part"
    with _open(url, 300) as r, open(tmp, "wb") as f:
        while True:
            b = r.read(1 << 16)
            if not b:
                break
            f.write(b)
    os.replace(tmp, dest)


def main():
    repo, outdir = sys.argv[1], sys.argv[2]
    prefix = sys.argv[3] if len(sys.argv) > 3 else ""
    files = [f for f in api_tree(repo) if f["type"] == "file"]
    if prefix:
        files = [f for f in files if f["path"].startswith(prefix)]
    print(f"{repo}: 取 {len(files)} 个文件 → {outdir}")
    for f in files:
        rel = f["path"]
        dest = os.path.join(outdir, rel)
        if os.path.exists(dest) and os.path.getsize(dest) == f.get("size", -1):
            continue
        # espeak 的 voices/!v/ 下有带空格的文件名（如 "Mr serious"），必须转义
        url = f"{MIRROR}/{repo}/resolve/main/{urllib.parse.quote(rel)}"
        try:
            fetch(url, dest)
        except Exception as e:
            print(f"  跳过 {rel}: {type(e).__name__}: {e}")
            continue
        print(f"  {rel}  ({f.get('size',0)/1024:.0f} KB)")
    print("完成")


if __name__ == "__main__":
    main()