#!/usr/bin/env python3
"""
参考图素材批量下载工具

从 markdown 文件中提取 URL，按域名自动分类，调用合适的工具下载图片。
支持 gallery-dl / yt-dlp / BeautifulSoup4 / Playwright 四种下载引擎。

用法:
    python download.py 参考图素材链接清单.md
    python download.py 参考图素材链接清单.md --dry-run
    python download.py 参考图素材链接清单.md --section 我的暑假
    python download.py 参考图素材链接清单.md --quick
"""

import argparse
import json
import os
import random
import re
import subprocess
import sys
import time
from pathlib import Path
from urllib.parse import urljoin, urlparse

# Windows 终端 UTF-8 输出
if sys.platform == "win32":
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

import requests
from bs4 import BeautifulSoup

# ---------------------------------------------------------------------------
# 常量
# ---------------------------------------------------------------------------

SCRIPT_DIR = Path(__file__).resolve().parent
SITE_RULES_PATH = SCRIPT_DIR / "site_rules.json"

# 章节标题关键词 → 输出子目录
SECTION_MAP = [
    ("我的暑假", "游戏截图/我的暑假"),
    ("放学后少年", "游戏截图/放学后少年"),
    ("城郊", "真实场景"),
    ("真实场景", "真实场景"),
    ("Lowpoly", "美术风格"),
    ("低多边形", "美术风格"),
    ("怀旧美术", "美术风格"),
    ("美术风格", "美术风格"),
]

# 默认 fallback
DEFAULT_SECTION_DIR = "其他"

# 请求头
HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
        "AppleWebKit/537.36 (KHTML, like Gecko) "
        "Chrome/125.0.0.0 Safari/537.36"
    )
}

# 下载延迟范围（秒），防止被封
DELAY_RANGE = (1.5, 4.0)

# 默认最小图片体积 (KB)，过滤掉图标/像素点
DEFAULT_MIN_SIZE_KB = 10


# ---------------------------------------------------------------------------
# Markdown 解析器
# ---------------------------------------------------------------------------

def parse_markdown(md_path: str) -> list[dict]:
    """解析 markdown 文件，返回 [{url, section, subsection, line}]"""
    content = Path(md_path).read_text(encoding="utf-8")
    lines = content.splitlines()

    results = []
    current_section = ""
    current_subsection = ""

    for i, line in enumerate(lines, 1):
        stripped = line.strip()

        # 追踪章节标题
        if stripped.startswith("## "):
            current_section = stripped.lstrip("#").strip()
            current_subsection = ""
        elif stripped.startswith("### "):
            current_subsection = stripped.lstrip("#").strip()

        # 跳过非 URL 行（表格分隔行等）
        if "---" in stripped and not "http" in stripped:
            continue

        # 提取 URL
        for m in re.finditer(r'https?://[^\s\)|\]]+', stripped):
            url = m.group(0).rstrip(".,;:)")  # 清理尾部标点
            results.append({
                "url": url,
                "section": current_section,
                "subsection": current_subsection,
                "line": i,
            })

    return results


def section_to_dir(section: str, subsection: str) -> str:
    """将章节标题映射为输出目录名"""
    # 先检查 section
    for keyword, dir_name in SECTION_MAP:
        if keyword in section:
            return dir_name
    # 再检查 subsection
    for keyword, dir_name in SECTION_MAP:
        if keyword in subsection:
            return dir_name
    return DEFAULT_SECTION_DIR


# ---------------------------------------------------------------------------
# 域名路由器
# ---------------------------------------------------------------------------

def load_site_rules() -> dict:
    """加载 site_rules.json"""
    if SITE_RULES_PATH.exists():
        with open(SITE_RULES_PATH, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}


def route_url(url: str, site_rules: dict) -> dict:
    """根据 URL 域名和 site_rules 确定下载引擎"""
    parsed = urlparse(url)
    domain = parsed.netloc.lower()

    # 去掉 www. 前缀尝试匹配
    domain_no_www = re.sub(r"^www\.", "", domain)

    # 精确匹配
    if domain in site_rules:
        return {"tool": site_rules[domain].get("engine", "bs4"), "rules": site_rules[domain]}
    if domain_no_www in site_rules:
        return {"tool": site_rules[domain_no_www].get("engine", "bs4"), "rules": site_rules[domain_no_www]}

    # 后缀匹配（处理子域名）
    for rule_domain, rules in site_rules.items():
        if domain.endswith(rule_domain) or rule_domain.endswith(domain):
            return {"tool": rules.get("engine", "bs4"), "rules": rules}

    # 默认使用 BS4
    return {"tool": "bs4", "rules": {}}


# ---------------------------------------------------------------------------
# 下载引擎：gallery-dl
# ---------------------------------------------------------------------------

def download_gallery_dl(url: str, output_dir: str, rules: dict) -> bool:
    """调用 gallery-dl 下载图片"""
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    cmd = ["gallery-dl", "-D", str(output_path)]

    # Pinterest / Instagram 需要 cookies
    if rules.get("cookies_required"):
        cmd.extend(["--cookies-from-browser", "chrome"])

    cmd.append(url)

    print(f"  [gallery-dl] {url}")
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.returncode == 0:
            print(f"    ✓ 完成")
            return True
        else:
            print(f"    ✗ 失败: {result.stderr[:200]}")
            return False
    except subprocess.TimeoutExpired:
        print(f"    ✗ 超时")
        return False
    except FileNotFoundError:
        print(f"    ✗ gallery-dl 未安装 (pip install gallery-dl)")
        return False


# ---------------------------------------------------------------------------
# 下载引擎：yt-dlp
# ---------------------------------------------------------------------------

def download_yt_dlp(url: str, output_dir: str, rules: dict) -> bool:
    """调用 yt-dlp 下载 YouTube 视频缩略图"""
    try:
        import yt_dlp
    except ImportError:
        print(f"    ✗ yt-dlp 未安装 (pip install yt-dlp)")
        return False

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    print(f"  [yt-dlp] {url}")

    opts = {
        "writethumbnail": True,
        "skip_download": True,
        "convertthumbnails": {"default": "png"},
        "outtmpl": str(output_path / "%(title)s.%(ext)s"),
        "quiet": True,
        "no_warnings": True,
    }

    try:
        with yt_dlp.YoutubeDL(opts) as ydl:
            ydl.download([url])
        print(f"    ✓ 缩略图已下载")
        return True
    except Exception as e:
        print(f"    ✗ 失败: {e}")
        return False


# ---------------------------------------------------------------------------
# 下载引擎：BeautifulSoup4 静态抓取
# ---------------------------------------------------------------------------

def download_bs4(url: str, output_dir: str, rules: dict) -> bool:
    """用 requests + BeautifulSoup 抓取静态页面中的图片"""
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    min_size_kb = rules.get("min_size_kb", DEFAULT_MIN_SIZE_KB)
    img_selector = rules.get("img_selector")

    print(f"  [bs4] {url}")

    try:
        resp = requests.get(url, headers=HEADERS, timeout=30)
        resp.raise_for_status()
    except requests.RequestException as e:
        print(f"    ✗ 请求失败: {e}")
        return False

    soup = BeautifulSoup(resp.text, "lxml")

    # 提取图片 URL
    img_urls = []

    if img_selector:
        # 先尝试自定义 CSS 选择器
        for img in soup.select(img_selector):
            src = img.get("src") or img.get("data-src") or img.get("data-original")
            if src:
                img_urls.append(urljoin(url, src))

    # 选择器无结果时 fallback 到通用提取
    if not img_urls:
        for img in soup.find_all("img"):
            for attr in ("src", "data-src", "data-original", "data-lazy-src"):
                src = img.get(attr)
                if src:
                    img_urls.append(urljoin(url, src))
                    break

    # 额外策略：扫描 <a> 标签中直接链接的图片
    for a in soup.find_all("a", href=True):
        href = a["href"]
        if any(href.lower().endswith(ext) for ext in (".jpg", ".jpeg", ".png", ".gif", ".webp")):
            img_urls.append(urljoin(url, href))

    if not img_urls:
        print(f"    ⚠ 未找到图片")
        return False

    # 去重
    img_urls = list(dict.fromkeys(img_urls))

    # 下载图片
    downloaded = 0
    skipped = 0
    for img_url in img_urls:
        # 跳过 base64 内嵌图片
        if img_url.startswith("data:"):
            continue

        # 跳过 SVG 图标、追踪像素等
        parsed_path = urlparse(img_url).path.lower()
        if parsed_path.endswith((".svg",)):
            continue

        try:
            # 先 HEAD 请求检查大小
            head_resp = requests.head(img_url, headers=HEADERS, timeout=10, allow_redirects=True)
            content_length = int(head_resp.headers.get("content-length", 0))
            if content_length < min_size_kb * 1024:
                skipped += 1
                continue

            # 下载图片
            img_resp = requests.get(img_url, headers=HEADERS, timeout=30, stream=True)
            img_resp.raise_for_status()

            # 生成文件名
            filename = Path(urlparse(img_url).path).name
            if not filename or "." not in filename:
                ext = ".jpg"
                content_type = img_resp.headers.get("content-type", "")
                if "png" in content_type:
                    ext = ".png"
                elif "webp" in content_type:
                    ext = ".webp"
                filename = f"img_{downloaded + 1}{ext}"

            # 避免文件名冲突
            filepath = output_path / filename
            counter = 1
            while filepath.exists():
                stem = filepath.stem
                suffix = filepath.suffix
                filepath = output_path / f"{stem}_{counter}{suffix}"
                counter += 1

            with open(filepath, "wb") as f:
                for chunk in img_resp.iter_content(8192):
                    f.write(chunk)

            downloaded += 1

            # 延迟
            time.sleep(random.uniform(*DELAY_RANGE))

        except Exception as e:
            print(f"    ⚠ 跳过 {img_url[:80]}: {e}")

    print(f"    ✓ 下载 {downloaded} 张，跳过 {skipped} 张小图")
    return downloaded > 0


# ---------------------------------------------------------------------------
# 下载引擎：Playwright 动态抓取
# ---------------------------------------------------------------------------

def download_playwright(url: str, output_dir: str, rules: dict) -> bool:
    """用 Playwright headless 浏览器抓取动态页面中的图片"""
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    min_size_kb = rules.get("min_size_kb", DEFAULT_MIN_SIZE_KB)
    wait_seconds = rules.get("wait_seconds", 3)
    scroll_to_load = rules.get("scroll_to_load", False)

    print(f"  [playwright] {url}")

    try:
        from playwright.sync_api import sync_playwright
    except ImportError:
        print(f"    ✗ playwright 未安装 (pip install playwright && playwright install chromium)")
        return False

    try:
        with sync_playwright() as p:
            browser = p.chromium.launch(headless=True)
            page = browser.new_page()
            page.goto(url, wait_until="networkidle", timeout=30000)

            # 等待额外渲染时间
            page.wait_for_timeout(wait_seconds * 1000)

            # 滚动页面触发懒加载
            if scroll_to_load:
                _scroll_page(page)

            # 提取图片 URL
            img_elements = page.query_selector_all("img")
            img_urls = []
            for img in img_elements:
                src = img.get_attribute("src") or img.get_attribute("data-src") or img.get_attribute("data-original")
                if src and not src.startswith("data:"):
                    img_urls.append(urljoin(url, src))

            browser.close()

    except Exception as e:
        print(f"    ✗ Playwright 失败: {e}")
        return False

    if not img_urls:
        print(f"    ⚠ 未找到图片")
        return False

    # 去重
    img_urls = list(dict.fromkeys(img_urls))

    # 下载图片
    downloaded = 0
    skipped = 0
    for img_url in img_urls:
        parsed_path = urlparse(img_url).path.lower()
        if parsed_path.endswith((".svg",)):
            continue

        try:
            head_resp = requests.head(img_url, headers=HEADERS, timeout=10, allow_redirects=True)
            content_length = int(head_resp.headers.get("content-length", 0))
            if content_length < min_size_kb * 1024:
                skipped += 1
                continue

            img_resp = requests.get(img_url, headers=HEADERS, timeout=30, stream=True)
            img_resp.raise_for_status()

            filename = Path(urlparse(img_url).path).name
            if not filename or "." not in filename:
                ext = ".jpg"
                content_type = img_resp.headers.get("content-type", "")
                if "png" in content_type:
                    ext = ".png"
                elif "webp" in content_type:
                    ext = ".webp"
                filename = f"img_{downloaded + 1}{ext}"

            filepath = output_path / filename
            counter = 1
            while filepath.exists():
                stem = filepath.stem
                suffix = filepath.suffix
                filepath = output_path / f"{stem}_{counter}{suffix}"
                counter += 1

            with open(filepath, "wb") as f:
                for chunk in img_resp.iter_content(8192):
                    f.write(chunk)

            downloaded += 1
            time.sleep(random.uniform(*DELAY_RANGE))

        except Exception as e:
            print(f"    ⚠ 跳过 {img_url[:80]}: {e}")

    print(f"    ✓ 下载 {downloaded} 张，跳过 {skipped} 张小图")
    return downloaded > 0


def _scroll_page(page):
    """滚动页面到底部，触发懒加载"""
    try:
        last_height = page.evaluate("document.body.scrollHeight")
        for _ in range(10):  # 最多滚动 10 次
            page.evaluate("window.scrollTo(0, document.body.scrollHeight)")
            page.wait_for_timeout(1000)
            new_height = page.evaluate("document.body.scrollHeight")
            if new_height == last_height:
                break
            last_height = new_height
        # 滚回顶部
        page.evaluate("window.scrollTo(0, 0)")
    except Exception:
        pass


# ---------------------------------------------------------------------------
# 调度器
# ---------------------------------------------------------------------------

TOOL_DISPATCHER = {
    "gallery-dl": download_gallery_dl,
    "yt-dlp": download_yt_dlp,
    "bs4": download_bs4,
    "playwright": download_playwright,
}


def process_urls(
    entries: list[dict],
    output_root: str,
    quick: bool = False,
    section_filter: str = "",
) -> None:
    """处理所有 URL 条目"""
    site_rules = load_site_rules()

    # 按 section 过滤
    if section_filter:
        entries = [e for e in entries if section_filter in e["section"] or section_filter in e["subsection"]]
        if not entries:
            print(f"⚠ 没有匹配章节 '{section_filter}' 的 URL")
            return

    # 统计
    stats = {"total": 0, "success": 0, "fail": 0, "skip": 0}
    by_tool = {}

    print(f"\n{'='*60}")
    print(f"  参考图素材批量下载")
    print(f"  共 {len(entries)} 个 URL")
    if quick:
        print(f"  模式：快速（仅 gallery-dl + yt-dlp）")
    print(f"{'='*60}\n")

    for i, entry in enumerate(entries, 1):
        url = entry["url"]
        section = entry["section"]
        dir_name = section_to_dir(section, entry["subsection"])
        output_dir = os.path.join(output_root, dir_name)

        route = route_url(url, site_rules)
        tool = route["tool"]
        rules = route["rules"]

        # quick 模式跳过 BS4/Playwright
        if quick and tool in ("bs4", "playwright"):
            print(f"  [{i}/{len(entries)}] 跳过 (quick模式): {url[:80]}")
            stats["skip"] += 1
            continue

        # 水印警告
        if rules.get("watermark_warning"):
            print(f"  ⚠ 注意：该站点图片带水印，仅供参考")

        print(f"\n[{i}/{len(entries)}] {section} → {dir_name}")

        stats["total"] += 1
        by_tool[tool] = by_tool.get(tool, 0) + 1

        downloader = TOOL_DISPATCHER.get(tool, download_bs4)
        success = downloader(url, output_dir, rules)

        if success:
            stats["success"] += 1
        else:
            stats["fail"] += 1

        # 请求间延迟
        if i < len(entries):
            delay = random.uniform(*DELAY_RANGE)
            time.sleep(delay)

    # 打印统计
    print(f"\n{'='*60}")
    print(f"  下载完成！")
    print(f"  总计: {stats['total']} 个 URL")
    print(f"  成功: {stats['success']}")
    print(f"  失败: {stats['fail']}")
    print(f"  跳过: {stats['skip']}")
    print(f"  引擎分布: {by_tool}")
    print(f"{'='*60}\n")


def dry_run(entries: list[dict]) -> None:
    """仅打印 URL 分类结果，不下载"""
    site_rules = load_site_rules()

    print(f"\n{'='*70}")
    print(f"  Dry Run — URL 分类预览 ({len(entries)} 个 URL)")
    print(f"{'='*70}\n")

    current_section = ""
    by_tool = {}
    by_dir = {}

    for entry in entries:
        url = entry["url"]
        section = entry["section"]
        dir_name = section_to_dir(section, entry["subsection"])
        route = route_url(url, site_rules)
        tool = route["tool"]

        by_tool[tool] = by_tool.get(tool, 0) + 1
        by_dir[dir_name] = by_dir.get(dir_name, 0) + 1

        if section != current_section:
            current_section = section
            print(f"\n  >> {section}")
            print(f"  {'~'*60}")

        watermark = " [水印]" if route.get("rules", {}).get("watermark_warning") else ""
        cookies = " [需Cookie]" if route.get("rules", {}).get("cookies_required") else ""
        print(f"    [{tool:>11}] {dir_name:<20} {url[:65]}{watermark}{cookies}")

    print(f"\n{'~'*70}")
    print(f"  引擎分布: {by_tool}")
    print(f"  目录分布: {by_dir}")
    print(f"{'='*70}\n")


# ---------------------------------------------------------------------------
# CLI 入口
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="参考图素材批量下载工具 — 从 markdown 提取 URL 并自动下载图片",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python download.py 参考图素材链接清单.md
  python download.py 参考图素材链接清单.md --dry-run
  python download.py 参考图素材链接清单.md --section 我的暑假
  python download.py 参考图素材链接清单.md --quick
  python download.py 参考图素材链接清单.md -o ./output
        """,
    )
    parser.add_argument(
        "markdown",
        help="包含 URL 的 markdown 文件路径",
    )
    parser.add_argument(
        "-o", "--output",
        default=None,
        help="输出根目录（默认：markdown 文件所在目录）",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="仅打印 URL 分类结果，不下载",
    )
    parser.add_argument(
        "--section",
        default="",
        help="仅处理包含指定关键词的章节",
    )
    parser.add_argument(
        "--quick",
        action="store_true",
        help="快速模式：仅使用 gallery-dl 和 yt-dlp（跳过自定义抓取）",
    )

    args = parser.parse_args()

    # 检查 markdown 文件
    md_path = Path(args.markdown)
    if not md_path.exists():
        print(f"✗ 文件不存在: {md_path}")
        sys.exit(1)

    # 解析 URL
    entries = parse_markdown(str(md_path))
    if not entries:
        print("⚠ 未在文件中找到任何 URL")
        sys.exit(0)

    # 确定输出目录
    if args.output:
        output_root = args.output
    else:
        output_root = str(md_path.parent)

    # 执行
    if args.dry_run:
        dry_run(entries)
    else:
        process_urls(
            entries,
            output_root=output_root,
            quick=args.quick,
            section_filter=args.section,
        )


if __name__ == "__main__":
    main()