#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从书稿 Markdown 中抓取远程图片（![alt](url) 与 <img src="...">），下载到本地 images/，
并把正文中的 URL 替换为相对路径 images/文件名。

用法：
  python3 scripts/download_chapter_images_from_md.py /path/to/第140场.md
  python3 scripts/download_chapter_images_from_md.py --images-dir /path/to/images a.md b.md
  python3 scripts/download_chapter_images_from_md.py --dry-run a.md

依赖：requests（与 content_download.py 一致）
"""
from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path
from urllib.parse import unquote, urlparse

try:
    import requests
except ImportError:
    print("请安装: pip install requests", file=sys.stderr)
    sys.exit(1)

MD_IMG = re.compile(r"!\[[^\]]*\]\((https?://[^)\s]+)\)")
HTML_IMG = re.compile(r'<img[^>]+src=["\'](https?://[^"\']+)["\']', re.I)


def collect_urls(text: str) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    for m in MD_IMG.finditer(text):
        u = m.group(1).strip()
        if u not in seen:
            seen.add(u)
            out.append(u)
    for m in HTML_IMG.finditer(text):
        u = m.group(1).strip()
        if u not in seen:
            seen.add(u)
            out.append(u)
    return out


def safe_filename(url: str, index: int) -> str:
    path = unquote(urlparse(url).path)
    base = Path(path).name
    if not base or base == "/" or ".." in base:
        base = f"image_{index:02d}.bin"
    base = re.sub(r"[^\w.\-一-龥]", "_", base)
    return base


def download_one(url: str, dest: Path, session: requests.Session) -> None:
    r = session.get(url, timeout=60, headers={"User-Agent": "Mozilla/5.0"})
    r.raise_for_status()
    dest.write_bytes(r.content)


def process_file(
    md_path: Path,
    images_dir: Path,
    dry_run: bool,
    session: requests.Session,
) -> tuple[int, int]:
    text = md_path.read_text(encoding="utf-8")
    urls = collect_urls(text)
    if not urls:
        print(f"{md_path}: 无远程图片")
        return 0, 0

    used_names: dict[str, str] = {}  # url -> local filename
    basename_count: dict[str, int] = {}
    ok = fail = 0

    for i, url in enumerate(urls, 1):
        name = safe_filename(url, i)
        stem = Path(name).stem
        ext = Path(name).suffix
        if name in basename_count:
            basename_count[name] += 1
            name = f"{stem}_{basename_count[name]}{ext}"
        else:
            basename_count[name] = 0

        local_rel = f"images/{name}"
        used_names[url] = name

        dest = images_dir / name
        print(f"  GET {url} -> {local_rel}")
        if dry_run:
            ok += 1
            continue
        images_dir.mkdir(parents=True, exist_ok=True)
        try:
            download_one(url, dest, session)
            ok += 1
        except Exception as e:
            print(f"  FAIL {url}: {e}", file=sys.stderr)
            fail += 1

    if dry_run or fail:
        return ok, fail

    new_text = text
    for url, name in used_names.items():
        local_rel = f"images/{name}"
        new_text = new_text.replace(f"]({url})", f"]({local_rel})")
        new_text = new_text.replace(f'src="{url}"', f'src="{local_rel}"')
        new_text = new_text.replace(f"src='{url}'", f"src='{local_rel}'")

    md_path.write_text(new_text, encoding="utf-8")
    print(f"  已写回: {md_path}")
    return ok, fail


def main() -> None:
    ap = argparse.ArgumentParser(description="从 md 下载远程插图并改为相对路径")
    ap.add_argument("md", nargs="+", type=Path, help="Markdown 文件路径")
    ap.add_argument(
        "--images-dir",
        type=Path,
        default=None,
        help="图片目录，默认与第一个 md 同级的 images/",
    )
    ap.add_argument("--dry-run", action="store_true", help="只打印不下载、不改文件")
    args = ap.parse_args()

    first = args.md[0].resolve()
    images_dir = (
        args.images_dir.resolve()
        if args.images_dir
        else first.parent / "images"
    )

    session = requests.Session()
    total_ok = total_fail = 0
    for p in args.md:
        p = p.resolve()
        if not p.is_file():
            print(f"跳过（非文件）: {p}", file=sys.stderr)
            continue
        print(f"=== {p.name}")
        ok, fail = process_file(p, images_dir, args.dry_run, session)
        total_ok += ok
        total_fail += fail

    print(f"完成: 成功 {total_ok}, 失败 {total_fail}")
    if total_fail:
        sys.exit(1)


if __name__ == "__main__":
    main()