zlib-dl/zlib_dl.py

#!/usr/bin/env python3
"""
ZLibrary Book Downloader - AI-friendly CLI tool

用法:
  zlib_dl.py search  --title "书名" --author "作者"
  zlib_dl.py download --title "书名" --author "作者" [--proxy-file proxies.txt]

支持搜索参数: title, author, isbn, doi, publisher, 或自由查询 query
输出格式: 默认人类友好, 加 --json 输出结构化 JSON
"""

import argparse
import gzip
import hashlib
import json
import os
import platform
import re
import shutil
import signal
import subprocess
import sys
import time
import unicodedata
import urllib.request
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Optional, List, Dict, Any, Tuple

try:
    from curl_cffi import requests as cffi_requests
    from bs4 import BeautifulSoup
except ImportError:
    print("错误: 缺少依赖。请运行: pip install curl_cffi beautifulsoup4 lxml", file=sys.stderr)
    sys.exit(1)


# ============================================================
# 数据模型
# ============================================================

@dataclass
class BookResult:
    """图书搜索结果"""
    id: str
    title: str
    author: str
    format: str
    filesize: str
    language: str
    year: str
    publisher: str
    isbn: str
    quality: str
    rating: str
    download_path: str
    page_url: str
    detail_path: str = ""
    deleted: str = "0"

    def to_dict(self) -> dict:
        return asdict(self)

    def display_line(self, index: int) -> str:
        """人类可读的单行显示"""
        parts = [
            f"{index:3d}.",
            f"[{self.format:5s}]",
            f"{self.filesize:>10s}",
            f"| {self.title[:55]:<55s}",
            f"| {self.author[:30]:<30s}",
            f"| {self.language[:8]:<8s}",
            f"| {self.year}",
        ]
        if self.quality and self.quality != "0":
            parts.append(f"| Q:{self.quality}")
        if self.isbn:
            parts.append(f"| ISBN:{self.isbn}")
        return " ".join(parts)


@dataclass
class SearchResponse:
    """搜索响应（AI友好）"""
    status: str  # "success" | "error" | "no_results"
    query: Dict[str, str]
    total_results: int
    results: List[BookResult]
    recommended: Dict[str, Optional[BookResult]]  # 按格式推荐
    message: str

    def to_dict(self) -> dict:
        d = {
            "status": self.status,
            "query": self.query,
            "total_results": self.total_results,
            "results": [r.to_dict() for r in self.results],
            "recommended": {
                fmt: (r.to_dict() if r else None)
                for fmt, r in self.recommended.items()
            },
            "message": self.message,
        }
        return d


@dataclass
class DownloadResult:
    """下载结果"""
    status: str  # "success" | "error" | "skipped"
    title: str
    author: str
    format: str
    filesize: str
    filepath: str
    message: str
    proxy_used: str = ""

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass
class ProxyConfig:
    """代理配置"""
    name: str
    type: str  # 'socks5', 'http', 'ss'
    host: str
    port: int
    # SS 特有
    method: str = ""
    password: str = ""
    obfs: str = ""
    obfs_host: str = ""
    # 运行时
    local_port: int = 0
    downloads_count: int = 0
    _process: Any = field(default=None, repr=False)


# ============================================================
# PoW 挑战解决器
# ============================================================

class ChallengeSolver:
    """解决 zlib.li 的 SHA1 工作量证明挑战"""

    @staticmethod
    def solve(html: str) -> Optional[Tuple[str, str]]:
        """
        解析并解决 PoW 挑战。
        返回 (c_token, c_time) 或 None（如果不是挑战页面）
        """
        if "Checking your browser" not in html:
            return None

        # 提取目标哈希
        match = re.search(r"\['([A-F0-9]{40})'", html)
        if not match:
            return None

        target_hash = match.group(1)
        n1 = int(target_hash[0], 16)

        start_time = time.time()
        i = 0
        while True:
            data = target_hash + str(i)
            digest = hashlib.sha1(data.encode()).digest()
            if digest[n1] == 0xB0 and digest[n1 + 1] == 0x0B:
                elapsed = time.time() - start_time
                c_token = target_hash + str(i)
                c_time = f"{elapsed:.3f}"
                return (c_token, c_time)
            i += 1
            if i > 10_000_000:
                return None


# ============================================================
# ZLib 会话管理
# ============================================================

class ZLibSession:
    """管理与 zlib.li 的会话，含 PoW 挑战处理"""

    BASE_URL = "https://zlib.li"

    def __init__(self, proxy_url: Optional[str] = None, verbose: bool = False):
        self.session = cffi_requests.Session()
        self.proxy_url = proxy_url
        self.verbose = verbose
        self._challenge_solved = False

    def _get_proxies(self) -> Optional[dict]:
        if self.proxy_url:
            return {"https": self.proxy_url, "http": self.proxy_url}
        return None

    def _debug_log_exit_ip(self, stage: str):
        """在 verbose 模式下打印当前会话出口 IP，便于核对代理是否一致"""
        if not self.verbose:
            return

        proxies = self._get_proxies()
        probe_targets = [
            ("http://httpbin.org/ip", lambda data: data.get("origin", "unknown")),
            ("http://ip-api.com/json", lambda data: data.get("query", "unknown")),
            ("https://api.ipify.org?format=json", lambda data: data.get("ip", "unknown")),
        ]

        proxy_desc = self.proxy_url if self.proxy_url else "直连"
        last_error = None
        for probe_url, parser in probe_targets:
            try:
                resp = self.session.get(
                    probe_url,
                    impersonate="chrome",
                    proxies=proxies,
                    timeout=15,
                )
                ip = parser(resp.json())
                print(
                    f"  🌐 [{stage}] 当前出口 IP: {ip} (proxy={proxy_desc}, probe={probe_url})",
                    file=sys.stderr,
                )
                return
            except Exception as e:
                last_error = e

        print(f"  🌐 [{stage}] 当前出口 IP 获取失败: {last_error}", file=sys.stderr)

    def _request(self, url: str, **kwargs) -> cffi_requests.Response:
        """发起请求，自动处理 PoW 挑战"""
        proxies = self._get_proxies()
        resp = self.session.get(
            url, impersonate="chrome", proxies=proxies, timeout=30, **kwargs
        )

        # 检查是否需要解决挑战
        if resp.status_code == 503 and "Checking your browser" in resp.text:
            if self.verbose:
                print("  ⏳ 正在解决 PoW 挑战...", file=sys.stderr)

            result = ChallengeSolver.solve(resp.text)
            if result is None:
                raise RuntimeError("无法解决 PoW 挑战")

            c_token, c_time = result
            self.session.cookies.set("c_token", c_token, domain="zlib.li", path="/")
            self.session.cookies.set("c_time", c_time, domain="zlib.li", path="/")
            self._challenge_solved = True

            if self.verbose:
                print(f"  ✅ 挑战已解决 (耗时 {c_time}s)", file=sys.stderr)

            # 重新请求
            resp = self.session.get(
                url, impersonate="chrome", proxies=proxies, timeout=30, **kwargs
            )

        return resp

    def search(self, query: str, max_results: int = 50) -> List[BookResult]:
        """搜索图书，返回结果列表"""
        self._debug_log_exit_ip("search")
        url = f"{self.BASE_URL}/s/?q={query}"
        resp = self._request(url)

        if resp.status_code != 200:
            raise RuntimeError(f"搜索失败: HTTP {resp.status_code}")

        soup = BeautifulSoup(resp.text, "lxml")
        bookcards = soup.find_all("z-bookcard")

        results = []
        for card in bookcards[:max_results]:
            title_el = card.find("div", slot="title")
            author_el = card.find("div", slot="author")

            result = BookResult(
                id=card.get("id", "").strip(),
                title=title_el.get_text(strip=True) if title_el else "",
                author=author_el.get_text(strip=True) if author_el else "",
                format=card.get("extension", "").strip(),
                filesize=card.get("filesize", "").strip(),
                language=card.get("language", "").strip(),
                year=card.get("year", "").strip(),
                publisher=card.get("publisher", "").strip(),
                isbn=card.get("isbn", "").strip(),
                quality=card.get("quality", "").strip(),
                rating=card.get("rating", "").strip(),
                download_path=card.get("download", "").strip().replace("\n", "").replace("\r", ""),
                page_url=card.get("href", "").strip().replace("\n", "").replace("\r", ""),
                deleted=card.get("deleted", "0").strip(),
            )
            # 跳过已删除的
            if result.deleted == "1":
                continue
            results.append(result)

        return results

    def _request_download(self, url: str, **kwargs) -> cffi_requests.Response:
        """发起下载请求，自动处理 PoW 挑战（支持 allow_redirects）"""
        proxies = self._get_proxies()
        timeout = kwargs.pop("timeout", (20, 90))
        resp = self.session.get(
            url, impersonate="chrome", proxies=proxies, timeout=timeout,
            allow_redirects=True, **kwargs
        )

        # 检查是否需要解决挑战 (503 或 200 但返回挑战页面)
        needs_challenge = (
            (resp.status_code == 503 and "Checking your browser" in resp.text)
            or (resp.status_code == 200
                and "text/html" in resp.headers.get("Content-Type", "")
                and "Checking your browser" in resp.text)
        )

        if needs_challenge:
            if self.verbose:
                print("  ⏳ 下载遇到 PoW 挑战，正在解决...", file=sys.stderr)

            result = ChallengeSolver.solve(resp.text)
            if result is None:
                raise RuntimeError("无法解决 PoW 挑战")

            c_token, c_time = result
            from urllib.parse import urlparse
            chal_domain = urlparse(resp.url).hostname or "zlib.li"
            self.session.cookies.set("c_token", c_token, domain=chal_domain, path="/")
            self.session.cookies.set("c_time", c_time, domain=chal_domain, path="/")

            if self.verbose:
                print(f"  ✅ 挑战已解决 (耗时 {c_time}s，节点 {chal_domain})", file=sys.stderr)

            # 无论什么页面，PoW 是在哪个域名的哪个具体 URL 上发生的，就全盘用带着 Cookie 的新状态重新向该域名发出最初的请求
            resp = self.session.get(
                resp.url, impersonate="chrome", proxies=proxies, timeout=timeout,
                allow_redirects=True, **kwargs
            )

        return resp

    def download_file(
        self, download_path: str, page_url: str, output_path: str
    ) -> Tuple[bool, str]:
        """
        下载文件。
        """
        if not download_path.startswith("/"):
            download_path = "/" + download_path
        
        url = f"{self.BASE_URL}{download_path}"
        detail_url = f"{self.BASE_URL}{page_url}" if page_url else ""

        # 模拟自然人行为：
        # 根据用户的经验，代理切换后，必须先在同一个全新代理 Session 下访问一遍“书籍详情页”，
        # 且下载链接是被CDN基于Session重新动态分配的。如果强行用旧链接，会被踢走。
        if detail_url:
            if self.verbose:
                print(f"  🔗 正在预热书籍详情页以注册当前代理指纹和获取新 Session...", file=sys.stderr)
            self._debug_log_exit_ip("warmup")
            try:
                dt_resp = self._request_download(detail_url, headers={"Referer": self.BASE_URL + "/"})
                import re
                # 尝试从详情页 HTML 中抓取服务器新动态生成的专属下载路径
                match = re.search(r'href="(/dl/[^"]+)"', dt_resp.text)
                if match:
                    dynamic_dl = match.group(1)
                    url = f"{self.BASE_URL}{dynamic_dl}"
                    if self.verbose:
                        print(f"  ⚡ 成功从详情页解析到动态下载链接: {dynamic_dl}", file=sys.stderr)
                else:
                    try:
                        with open("failed_dl_page.html", "w") as f:
                            f.write(dt_resp.text)
                    except Exception:
                        pass
                    return False, f"未能从书籍详情页解析出最新的真实下载链接，已保存现场到 failed_dl_page.html"
            except Exception as e:
                return False, f"预热书籍详情页或解析出错: {e}"

        try:
            self._debug_log_exit_ip("download")
            resp = self._request_download(
                url,
                headers={"Referer": detail_url} if detail_url else {},
                stream=True,
                timeout=(20, 120),
            )
        except Exception as e:
            return False, f"下载请求失败: {e}"

        if resp.status_code != 200:
            # 检查是否下载限制
            if resp.status_code == 429:
                return False, "下载次数已达上限 (每IP每天5本)"
            content = ""
            try:
                content = resp.text[:500]
            except Exception:
                pass
            if "limit" in content.lower():
                return False, "下载次数已达上限 (每IP每天5本)"
            return False, f"下载失败: HTTP {resp.status_code}"

        content_type = resp.headers.get("Content-Type", "")
        if "text/html" in content_type:
            # 下载接口返回了 HTML，这说明被系统拦截甚至踢回了首页
            html_text = resp.text[:2000].lower()
            if "daily limit" in html_text or "limit reached" in html_text or "已达下载上限" in html_text:
                return False, "下载次数已达真实上限"
            elif "sign in" in html_text or "login" in html_text or "登录" in html_text:
                return False, "系统要求登录或达到匿名限制"
            else:
                return False, "下载失败，服务端依然返回了 HTML 重定向网页！"

        # 写入文件（流式下载：只要持续有数据就继续）
        os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
        tmp_path = output_path + ".part"
        total_bytes = 0

        try:
            content_length = resp.headers.get("Content-Length", "").strip()
            expected_total = int(content_length) if content_length.isdigit() else 0

            if self.verbose:
                if expected_total > 0:
                    print(
                        f"  📦 服务器文件大小: {expected_total / (1024 * 1024):.2f} MB",
                        file=sys.stderr,
                    )
                else:
                    print("  📦 服务器未返回文件大小（Content-Length）", file=sys.stderr)

            started_at = time.time()
            last_log_at = started_at

            with open(tmp_path, "wb") as f:
                for chunk in resp.iter_content(chunk_size=262144):
                    if not chunk:
                        continue
                    f.write(chunk)
                    total_bytes += len(chunk)

                    now = time.time()
                    if self.verbose and now - last_log_at >= 3:
                        elapsed = max(0.001, now - started_at)
                        speed = total_bytes / elapsed
                        if expected_total > 0:
                            pct = (total_bytes / expected_total) * 100
                            print(
                                f"  ⏬ 已下载 {total_bytes / (1024 * 1024):.2f} / {expected_total / (1024 * 1024):.2f} MB ({pct:.1f}%), 速度 {speed / (1024 * 1024):.2f} MB/s",
                                file=sys.stderr,
                            )
                        else:
                            print(
                                f"  ⏬ 已下载 {total_bytes / (1024 * 1024):.2f} MB, 速度 {speed / (1024 * 1024):.2f} MB/s",
                                file=sys.stderr,
                            )
                        last_log_at = now

            os.replace(tmp_path, output_path)
            if expected_total > 0:
                return True, f"下载完成: {total_bytes} bytes / 期望 {expected_total} bytes"
            return True, f"下载完成: {total_bytes} bytes"
        except Exception as e:
            try:
                if os.path.exists(tmp_path):
                    os.remove(tmp_path)
            except Exception:
                pass
            return False, f"下载流中断: {e}"
        finally:
            try:
                resp.close()
            except Exception:
                pass


# ============================================================
# 代理管理器
# ============================================================

class ProxyManager:
    """管理代理列表和轮询"""

    MAX_DOWNLOADS_PER_PROXY = 5

    def __init__(self, proxy_file: Optional[str] = None, verbose: bool = False):
        self.proxies: List[ProxyConfig] = []
        self.current_index = -1  # -1 means direct (no proxy)
        self.verbose = verbose
        self._active_processes: List[subprocess.Popen] = []
        self._mihomo_bin: Optional[str] = None

        if proxy_file:
            self.proxies = self._load_proxies(proxy_file)
            if self.verbose:
                print(f"  📡 已加载 {len(self.proxies)} 个代理", file=sys.stderr)

    def _ensure_mihomo(self) -> Optional[str]:
        """下载并配置 mihomo 代理引擎到当前环境"""
        if self._mihomo_bin and os.path.exists(self._mihomo_bin):
            return self._mihomo_bin

        # 优先选择虚拟环境的 bin 目录，否则取当前目录的 bin
        base_dir = Path(sys.prefix) if sys.prefix != sys.base_prefix else Path.cwd()
        bin_dir = base_dir / "bin"
        bin_dir.mkdir(parents=True, exist_ok=True)

        mihomo_path = bin_dir / ("mihomo.exe" if platform.system().lower() == "windows" else "mihomo")

        if mihomo_path.exists():
            self._mihomo_bin = str(mihomo_path)
            return self._mihomo_bin

        system = platform.system().lower()
        machine = platform.machine().lower()
        is_arm = machine in ("arm64", "aarch64")

        if system == "darwin":
            arch_str = "arm64" if is_arm else "amd64-compatible"
        elif system == "linux":
            arch_str = "arm64" if is_arm else "amd64-compatible"
        else:
            print(f"代理错误: 暂不支持自动下载当前系统 ({system}) 的 mihomo", file=sys.stderr)
            return None

        version = "v1.19.21"
        url = f"https://github.com/MetaCubeX/mihomo/releases/download/{version}/mihomo-{system}-{arch_str}-{version}.gz"

        if self.verbose:
            print(f"  ⬇️ 首次使用 SS 代理，正在自动下载代理引擎...", file=sys.stderr)
            print(f"     URL: {url}", file=sys.stderr)

        try:
            req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
            with urllib.request.urlopen(req, timeout=120) as response:
                with gzip.GzipFile(fileobj=response) as uncompressed:
                    with open(mihomo_path, 'wb') as f:
                        f.write(uncompressed.read())

            os.chmod(mihomo_path, 0o755)
            self._mihomo_bin = str(mihomo_path)

            if self.verbose:
                print(f"  ✅ 代理引擎下载完成: {mihomo_path}", file=sys.stderr)
            return self._mihomo_bin
        except Exception as e:
            print(f"代理错误: 自动下载代理引擎失败: {e}", file=sys.stderr)
            return None

    def _load_proxies(self, filepath: str) -> List[ProxyConfig]:
        """加载代理文件，支持 Surge SS 格式和 socks5:// 格式"""
        proxies = []
        with open(filepath, "r", encoding="utf-8") as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line or line.startswith("#"):
                    continue

                # 尝试 socks5:// 或 http:// 格式
                if line.startswith(("socks5://", "http://", "https://")):
                    proxy = self._parse_url_proxy(line, line_num)
                    if proxy:
                        proxies.append(proxy)
                    continue

                # 尝试 Surge SS 格式
                proxy = self._parse_surge_ss(line, line_num)
                if proxy:
                    proxies.append(proxy)

        return proxies

    def _parse_url_proxy(self, line: str, line_num: int) -> Optional[ProxyConfig]:
        """解析 socks5://host:port 格式"""
        match = re.match(r"(socks5|http|https)://([^:]+):(\d+)", line)
        if match:
            return ProxyConfig(
                name=f"proxy-{line_num}",
                type=match.group(1),
                host=match.group(2),
                port=int(match.group(3)),
            )
        return None

    def _parse_surge_ss(self, line: str, line_num: int) -> Optional[ProxyConfig]:
        """
        解析 Surge Shadowsocks 格式:
        🇭🇰 香港 01= ss, host, port, encrypt-method=xxx, password=xxx, ...
        """
        # 分离名称和配置
        parts = line.split("=", 1)
        if len(parts) != 2:
            return None

        name = parts[0].strip()
        config_str = parts[1].strip()

        # 解析配置项
        config_parts = [p.strip() for p in config_str.split(",")]
        if len(config_parts) < 3 or config_parts[0] != "ss":
            return None

        host = config_parts[1].strip()
        try:
            port = int(config_parts[2].strip())
        except ValueError:
            return None

        # 解析键值对参数
        params = {}
        for part in config_parts[3:]:
            if "=" in part:
                k, v = part.split("=", 1)
                params[k.strip()] = v.strip()

        return ProxyConfig(
            name=name,
            type="ss",
            host=host,
            port=port,
            method=params.get("encrypt-method", ""),
            password=params.get("password", ""),
            obfs=params.get("obfs", ""),
            obfs_host=params.get("obfs-host", ""),
        )

    def _start_mihomo(self, proxy: ProxyConfig, local_port: int) -> bool:
        """为 SS 代理启动 mihomo 进程"""
        mihomo_bin = self._ensure_mihomo()
        if not mihomo_bin:
            return False

        # 生成配置文件
        cfg_path = Path(mihomo_bin).parent / f"config_{local_port}.yaml"

        # 处理 obfs 插件配置
        plugin_yaml = ""
        if proxy.obfs:
            plugin_yaml = f"""
    plugin: obfs
    plugin-opts:
      mode: {proxy.obfs}
      host: {proxy.obfs_host}"""

        # 尝试绕过 macOS 的 Surge 增强模式
        # 探测当前物理网卡 (避开 VPN 的 utun 接口)
        interface_yaml = ""
        if platform.system().lower() == "darwin":
            try:
                out = subprocess.check_output(
                    "netstat -rn -f inet | grep -E '^default' | grep -v 'utun'",
                    shell=True, timeout=1
                ).decode()
                for line in out.strip().split("\n"):
                    parts = line.split()
                    if len(parts) >= 4:
                        interface_yaml = f"\n    interface-name: {parts[3]}"
                        break
            except Exception:
                pass

        yaml_config = f"""
mode: global
port: 0
mixed-port: {local_port}
bind-address: '127.0.0.1'

# 独立的 DNS 解析，避免 Surge 的 Fake IP 污染导致流量绕回 Surge
dns:
  enable: true
  nameserver:
    - https://223.5.5.5/dns-query
    - https://1.1.1.1/dns-query

proxies:
  - name: "{proxy.name}"
    type: {proxy.type}
    server: {proxy.host}
    port: {proxy.port}
    cipher: "{proxy.method}"
    password: "{proxy.password}"{plugin_yaml}{interface_yaml}

proxy-groups:
  - name: GLOBAL
    type: select
    proxies:
      - "{proxy.name}"
"""
        with open(cfg_path, 'w', encoding='utf-8') as f:
            f.write(yaml_config)

        cmd = [mihomo_bin, "-d", str(Path(mihomo_bin).parent), "-f", str(cfg_path)]

        try:
            proc = subprocess.Popen(
                cmd,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            time.sleep(1.5)  # 等待引擎启动绑定端口

            if proc.poll() is not None:
                if self.verbose:
                    print(f"  ❌ mihomo 启动失败: {proxy.name}", file=sys.stderr)
                return False

            proxy._process = proc
            proxy.local_port = local_port
            # 记录对应的配置文件路径以便停止时清理
            proxy._cfg_path = str(cfg_path)
            self._active_processes.append(proc)

            if self.verbose:
                print(
                    f"  🔗 已启动代理隧道: {proxy.name} -> 127.0.0.1:{local_port}",
                    file=sys.stderr,
                )
            return True

        except Exception as e:
            print(f"错误: 启动 mihomo 代理引擎失败: {e}", file=sys.stderr)
            return False

    def _stop_mihomo(self, proxy: ProxyConfig):
        """停止代理的 mihomo 进程并清理"""
        if proxy._process and proxy._process.poll() is None:
            proxy._process.terminate()
            try:
                proxy._process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                proxy._process.kill()
            if self.verbose:
                print(f"  🔌 已停止代理隧道: {proxy.name}", file=sys.stderr)

        # 清理配置文件
        if hasattr(proxy, '_cfg_path') and os.path.exists(proxy._cfg_path):
            try:
                os.remove(proxy._cfg_path)
            except Exception:
                pass

    def get_current_proxy_url(self) -> Optional[str]:
        """获取当前活跃的代理 URL"""
        if self.current_index < 0 or self.current_index >= len(self.proxies):
            return None

        proxy = self.proxies[self.current_index]

        if proxy.type == "ss":
            if proxy.local_port == 0:
                return None
            return f"socks5://127.0.0.1:{proxy.local_port}"
        elif proxy.type in ("socks5", "http", "https"):
            return f"{proxy.type}://{proxy.host}:{proxy.port}"

        return None

    def get_current_proxy_name(self) -> str:
        """获取当前代理名称"""
        if self.current_index < 0 or self.current_index >= len(self.proxies):
            return "直连"
        return self.proxies[self.current_index].name

    def record_download(self):
        """记录一次下载"""
        if 0 <= self.current_index < len(self.proxies):
            self.proxies[self.current_index].downloads_count += 1

    def should_rotate(self) -> bool:
        """是否需要轮换代理"""
        if self.current_index < 0:
            return len(self.proxies) > 0  # 有代理可用就应该开始用
        if self.current_index >= len(self.proxies):
            return False  # 已经用完了
        proxy = self.proxies[self.current_index]
        return proxy.downloads_count >= self.MAX_DOWNLOADS_PER_PROXY

    def has_proxy_pool(self) -> bool:
        """是否配置了可轮换的代理池"""
        return len(self.proxies) > 0

    def ensure_active(self) -> bool:
        """确保当前有活跃代理（有代理池时）"""
        if not self.proxies:
            return True
        if 0 <= self.current_index < len(self.proxies):
            return True
        return self.rotate()

    def rotate(self) -> bool:
        """
        轮换到下一个代理。
        返回 True 表示成功切换，False 表示没有更多代理可用。
        """
        # 停止当前代理
        if 0 <= self.current_index < len(self.proxies):
            current = self.proxies[self.current_index]
            if current.type == "ss":
                self._stop_mihomo(current)

        self.current_index += 1

        if self.current_index >= len(self.proxies):
            if self.verbose:
                print("  ⚠️ 所有代理已用完", file=sys.stderr)
            return False

        proxy = self.proxies[self.current_index]

        if proxy.type == "ss":
            local_port = 10800 + self.current_index
            success = self._start_mihomo(proxy, local_port)
            if not success:
                # 尝试下一个
                return self.rotate()
        else:
            if self.verbose:
                print(
                    f"  📡 切换到代理: {proxy.name} ({proxy.type}://{proxy.host}:{proxy.port})",
                    file=sys.stderr,
                )

        return True

    def remaining_proxies(self) -> int:
        """剩余可用代理数"""
        return max(0, len(self.proxies) - self.current_index - 1)

    def remaining_downloads(self) -> int:
        """当前代理剩余下载次数"""
        if self.current_index < 0 or self.current_index >= len(self.proxies):
            return 0
        proxy = self.proxies[self.current_index]
        return max(0, self.MAX_DOWNLOADS_PER_PROXY - proxy.downloads_count)

    def cleanup(self):
        """清理所有 proxy 相关的进程和配置"""
        for proxy in self.proxies:
            if proxy._process:
                self._stop_mihomo(proxy)


# ============================================================
# 智能匹配与推荐
# ============================================================

def normalize_text(text: str) -> str:
    """规范化文本用于比较（去除标点、统一大小写）"""
    text = unicodedata.normalize("NFKC", text)
    text = text.lower().strip()
    # 去除常见前后缀和标点
    text = re.sub(r"[^\w\s]", " ", text)
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def fuzzy_match(query: str, target: str, threshold: float = 0.6) -> float:
    """
    模糊匹配评分（0-1）。
    基于词汇重叠度。
    """
    if not query or not target:
        return 0.0

    q_norm = normalize_text(query)
    t_norm = normalize_text(target)

    # 完全匹配
    if q_norm == t_norm:
        return 1.0

    # 包含匹配
    if q_norm in t_norm or t_norm in q_norm:
        return 0.95

    # 词汇重叠
    q_words = set(q_norm.split())
    t_words = set(t_norm.split())

    if not q_words:
        return 0.0

    overlap = len(q_words & t_words)
    score = overlap / len(q_words)

    return score


def recommend_books(
    results: List[BookResult],
    title: str = "",
    author: str = "",
    language: Optional[str] = None,
) -> Dict[str, Optional[BookResult]]:
    """
    从结果中推荐最佳的 epub 和 pdf 版本。
    优先级: epub > pdf > mobi > 其他
    """
    FORMAT_PRIORITY = {"epub": 1, "pdf": 2, "mobi": 3, "azw3": 4, "fb2": 5, "djvu": 6}

    # 过滤和评分
    scored: List[Tuple[float, BookResult]] = []
    for book in results:
        score = 0.0

        # 标题匹配
        if title:
            title_score = fuzzy_match(title, book.title)
            if title_score < 0.5:
                continue  # 标题不匹配，跳过
            score += title_score * 50

        # 作者匹配
        if author:
            author_score = fuzzy_match(author, book.author)
            score += author_score * 30

        # 语言匹配（不做硬过滤，做加分）
        if language:
            if language.lower() in book.language.lower():
                score += 10

        # 质量评分
        try:
            q = float(book.quality) if book.quality else 0
            score += q * 2
        except ValueError:
            pass

        # 格式优先级（epub 和 pdf 加分）
        fmt_priority = FORMAT_PRIORITY.get(book.format.lower(), 10)
        score += max(0, 10 - fmt_priority * 2)

        scored.append((score, book))

    # 按评分排序
    scored.sort(key=lambda x: -x[0])

    # 按格式选出最佳
    recommended: Dict[str, Optional[BookResult]] = {
        "epub": None,
        "pdf": None,
        "mobi": None,
    }

    for score, book in scored:
        fmt = book.format.lower()
        if fmt in recommended and recommended[fmt] is None:
            recommended[fmt] = book

    return recommended


# ============================================================
# 搜索命令
# ============================================================

def build_query(args) -> Tuple[str, Dict[str, str]]:
    """从命令行参数构建搜索查询"""
    parts = []
    query_info = {}

    if hasattr(args, "query") and args.query:
        parts.append(args.query)
        query_info["raw_query"] = args.query
    else:
        if args.title:
            parts.append(args.title)
            query_info["title"] = args.title
        if args.author:
            parts.append(args.author)
            query_info["author"] = args.author
        if args.isbn:
            parts.append(args.isbn)
            query_info["isbn"] = args.isbn
        if args.publisher:
            parts.append(args.publisher)
            query_info["publisher"] = args.publisher

    query_str = " ".join(parts)
    query_info["combined_query"] = query_str
    return query_str, query_info


def cmd_search(args):
    """执行搜索命令"""
    query_str, query_info = build_query(args)

    if not query_str:
        print("错误: 请提供至少一个搜索条件 (--title, --author, --isbn, --query)", file=sys.stderr)
        sys.exit(1)

    if args.verbose:
        print(f"  🔍 搜索: {query_str}", file=sys.stderr)

    # 创建会话（搜索不需要代理）
    proxy_url = None
    if hasattr(args, "proxy") and args.proxy:
        proxy_url = args.proxy

    session = ZLibSession(proxy_url=proxy_url, verbose=args.verbose)

    try:
        results = session.search(query_str, max_results=args.max_results)
    except Exception as e:
        response = SearchResponse(
            status="error",
            query=query_info,
            total_results=0,
            results=[],
            recommended={},
            message=f"搜索出错: {e}",
        )
        if args.json:
            print(json.dumps(response.to_dict(), ensure_ascii=False, indent=2))
        else:
            print(f"❌ {response.message}", file=sys.stderr)
        sys.exit(1)

    if not results:
        response = SearchResponse(
            status="no_results",
            query=query_info,
            total_results=0,
            results=[],
            recommended={},
            message=f"未找到与 '{query_str}' 匹配的图书",
        )
        if args.json:
            print(json.dumps(response.to_dict(), ensure_ascii=False, indent=2))
        else:
            print(f"📭 {response.message}")
        return

    # 推荐
    recommended = recommend_books(
        results,
        title=args.title or "",
        author=args.author or "",
        language=getattr(args, "language", None),
    )

    # 构建消息
    rec_parts = []
    for fmt in ["epub", "pdf", "mobi"]:
        if recommended.get(fmt):
            r = recommended[fmt]
            rec_parts.append(f"{fmt} ({r.filesize}, Q:{r.quality}, {r.language})")
    rec_msg = ", ".join(rec_parts) if rec_parts else "无推荐"

    response = SearchResponse(
        status="success",
        query=query_info,
        total_results=len(results),
        results=results,
        recommended=recommended,
        message=f"找到 {len(results)} 个结果。推荐: {rec_msg}",
    )

    if args.json:
        print(json.dumps(response.to_dict(), ensure_ascii=False, indent=2))
    else:
        # 人类友好输出
        print(f"\n📚 搜索结果: {query_str}")
        print(f"   共 {len(results)} 个结果\n")

        # 推荐
        print("🏆 推荐版本:")
        for fmt in ["epub", "pdf", "mobi"]:
            r = recommended.get(fmt)
            if r:
                print(
                    f"   {fmt.upper():5s}  {r.title[:50]}  |  {r.author[:25]}  |  "
                    f"{r.filesize}  |  {r.language}  |  Q:{r.quality}  | dl={r.download_path}"
                )
            else:
                print(f"   {fmt.upper():5s}  (未找到)")

        print(f"\n{'─' * 120}")
        print(f"{'  #':>4s}  {'[格式]':>7s}  {'大小':>10s}  {'标题':<55s}  {'作者':<30s}  {'语言':<8s}  {'年份':<4s}")
        print(f"{'─' * 120}")

        for i, book in enumerate(results, 1):
            print(book.display_line(i))

        print(f"{'─' * 120}")
        print(f"\n💡 下载提示: zlib_dl.py download --dl-path \"/dl/xxx\" --title \"书名\"")


# ============================================================
# 下载命令
# ============================================================

def sanitize_filename(name: str) -> str:
    """清理文件名"""
    # 替换非法字符
    name = re.sub(r'[<>:"/\\|?*]', "_", name)
    name = re.sub(r"\s+", " ", name)
    return name.strip()[:200]


def _should_rotate_proxy(message: str) -> bool:
    """根据错误信息判断是否应该切换代理重试"""
    msg_lower = message.lower()
    return (
        "限" in message
        or "limit" in msg_lower
        or "429" in message
        or "503" in message
        or "504" in message
        or "403" in message
        or "502" in message
        or "timeout" in msg_lower
        or "请求失败" in message
        or "html" in msg_lower
    )


def _create_session_for_current_proxy(
    proxy_mgr: ProxyManager,
    explicit_proxy: Optional[str],
    verbose: bool,
) -> ZLibSession:
    """根据当前代理状态创建新的 ZLibSession"""
    proxy_url = explicit_proxy if explicit_proxy else proxy_mgr.get_current_proxy_url()
    return ZLibSession(proxy_url=proxy_url, verbose=verbose)


def cmd_download(args):
    """执行下载命令"""
    explicit_proxy = getattr(args, "proxy", None)
    proxy_file = getattr(args, "proxy_file", None)
    if explicit_proxy and proxy_file:
        print("错误: --proxy 与 --proxy-file 不能同时使用", file=sys.stderr)
        sys.exit(1)

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

    # 代理管理器
    proxy_mgr = ProxyManager(
        proxy_file=proxy_file, verbose=args.verbose
    )

    if proxy_mgr.has_proxy_pool() and not proxy_mgr.ensure_active():
        print("错误: 代理池不可用，无法启动任何代理", file=sys.stderr)
        sys.exit(1)

    download_results: List[DownloadResult] = []

    try:
        # 模式 1: 直接指定下载路径
        if args.dl_path:
            dl_paths = args.dl_path if isinstance(args.dl_path, list) else [args.dl_path]
            title = args.title or "unknown"
            author = args.author or "unknown"
            fmt = getattr(args, 'fmt', '') or ''

            for dl_path in dl_paths:
                result = _download_one(
                    dl_path=dl_path,
                    page_url="",
                    title=title,
                    author=author,
                    fmt=fmt or "",
                    output_dir=output_dir,
                    proxy_mgr=proxy_mgr,
                    explicit_proxy=explicit_proxy,
                    verbose=args.verbose,
                )
                download_results.append(result)

        # 模式 2: 搜索后自动选择并下载
        else:
            query_str, query_info = build_query(args)
            if not query_str:
                print(
                    "错误: 请提供搜索条件 (--title, --author) 或直接指定 --dl-path",
                    file=sys.stderr,
                )
                sys.exit(1)

            if args.verbose:
                print(f"  🔍 搜索: {query_str}", file=sys.stderr)

            # 搜索
            session = _create_session_for_current_proxy(
                proxy_mgr=proxy_mgr,
                explicit_proxy=explicit_proxy,
                verbose=args.verbose,
            )
            results = session.search(query_str)

            if not results:
                print(f"📭 未找到结果: {query_str}", file=sys.stderr)
                if args.json:
                    print(json.dumps({"status": "no_results", "message": f"未找到: {query_str}"}, ensure_ascii=False))
                return

            # 推荐
            recommended = recommend_books(
                results,
                title=args.title or "",
                author=args.author or "",
            )

            # 下载推荐的 epub 和 pdf
            for fmt in ["epub", "pdf"]:
                book = recommended.get(fmt)
                if book and book.download_path:
                    result = _download_one(
                        dl_path=book.download_path,
                        page_url=book.page_url,
                        title=book.title,
                        author=book.author,
                        fmt=book.format,
                        output_dir=output_dir,
                        proxy_mgr=proxy_mgr,
                        explicit_proxy=explicit_proxy,
                        verbose=args.verbose,
                    )
                    download_results.append(result)

            # 如果没有 epub 或 pdf, 尝试 mobi
            if not recommended.get("epub") and not recommended.get("pdf"):
                book = recommended.get("mobi")
                if book and book.download_path:
                    result = _download_one(
                        dl_path=book.download_path,
                        page_url=book.page_url,
                        title=book.title,
                        author=book.author,
                        fmt=book.format,
                        output_dir=output_dir,
                        proxy_mgr=proxy_mgr,
                        explicit_proxy=explicit_proxy,
                        verbose=args.verbose,
                    )
                    download_results.append(result)

    finally:
        proxy_mgr.cleanup()

    # 输出结果
    if args.json:
        print(
            json.dumps(
                {
                    "status": "complete",
                    "downloads": [r.to_dict() for r in download_results],
                    "proxy_stats": {
                        "remaining_proxies": proxy_mgr.remaining_proxies(),
                        "current_proxy": proxy_mgr.get_current_proxy_name(),
                    },
                },
                ensure_ascii=False,
                indent=2,
            )
        )
    else:
        print(f"\n📊 下载完成:")
        for r in download_results:
            icon = "✅" if r.status == "success" else "❌"
            print(f"  {icon} [{r.format}] {r.title} - {r.message}")


def _download_one(
    dl_path: str,
    page_url: str,
    title: str,
    author: str,
    fmt: str,
    output_dir: Path,
    proxy_mgr: ProxyManager,
    explicit_proxy: Optional[str],
    verbose: bool,
) -> DownloadResult:
    """下载单个文件，带代理轮询逻辑"""

    # 确定文件扩展名
    if not fmt:
        # 尝试从路径推断
        fmt = "bin"

    # 构建文件名
    filename = sanitize_filename(f"{title} - {author}.{fmt}")
    output_path = str(output_dir / filename)

    # 检查是否已存在
    if os.path.exists(output_path):
        if verbose:
            print(f"  ⏩ 已存在: {filename}", file=sys.stderr)
        return DownloadResult(
            status="skipped",
            title=title,
            author=author,
            format=fmt,
            filesize="",
            filepath=output_path,
            message="文件已存在，跳过",
        )

    uses_proxy_pool = proxy_mgr.has_proxy_pool()

    # 进入下载前，若当前代理已达配额，先轮换
    if uses_proxy_pool and proxy_mgr.should_rotate() and not proxy_mgr.rotate():
        return DownloadResult(
            status="error",
            title=title,
            author=author,
            format=fmt,
            filesize="",
            filepath="",
            message="所有代理已用完，且当前代理下载次数已达上限",
        )

    while True:
        if uses_proxy_pool and not proxy_mgr.ensure_active():
            return DownloadResult(
                status="error",
                title=title,
                author=author,
                format=fmt,
                filesize="",
                filepath="",
                message="所有代理已用完",
            )

        if uses_proxy_pool:
            proxy_name = proxy_mgr.get_current_proxy_name()
        elif explicit_proxy:
            proxy_name = explicit_proxy
        else:
            proxy_name = "直连"

        if verbose:
            print(f"  ⬇️  下载: {filename} (代理: {proxy_name})", file=sys.stderr)

        # 每次下载尝试都创建全新 Session；切代理后会自动带新出口重建会话
        session = _create_session_for_current_proxy(
            proxy_mgr=proxy_mgr,
            explicit_proxy=explicit_proxy,
            verbose=verbose,
        )
        success, message = session.download_file(dl_path, page_url, output_path)

        if success:
            if uses_proxy_pool:
                proxy_mgr.record_download()
            filesize = os.path.getsize(output_path)
            size_str = _format_size(filesize)
            return DownloadResult(
                status="success",
                title=title,
                author=author,
                format=fmt,
                filesize=size_str,
                filepath=output_path,
                message=f"下载成功 ({size_str})",
                proxy_used=proxy_name,
            )

        if not uses_proxy_pool or not _should_rotate_proxy(message):
            return DownloadResult(
                status="error",
                title=title,
                author=author,
                format=fmt,
                filesize="",
                filepath="",
                message=message,
            )

        if not proxy_mgr.rotate():
            return DownloadResult(
                status="error",
                title=title,
                author=author,
                format=fmt,
                filesize="",
                filepath="",
                message=f"{message}；且所有代理已用完",
            )

        if verbose:
            print(
                f"  🔄 下载受阻 ({message})，切换代理重试: {proxy_mgr.get_current_proxy_name()}",
                file=sys.stderr,
            )

def _format_size(size_bytes: int) -> str:
    """格式化文件大小"""
    for unit in ["B", "KB", "MB", "GB"]:
        if size_bytes < 1024:
            return f"{size_bytes:.1f} {unit}"
        size_bytes /= 1024
    return f"{size_bytes:.1f} TB"


# ============================================================
# CLI 入口
# ============================================================

def main():
    parser = argparse.ArgumentParser(
        prog="zlib_dl",
        description="ZLibrary 电子书搜索与下载工具 (AI-friendly)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 搜索（人类友好）
  %(prog)s search --title "Project Hail Mary" --author "Andy Weir"

  # 搜索（JSON输出，AI友好）
  %(prog)s search --title "Project Hail Mary" --author "Andy Weir" --json

  # 自由搜索
  %(prog)s search --query "我从达尔文那里学到的投资知识 普拉克·普拉萨德"

  # 自动搜索并下载最佳 epub + pdf
  %(prog)s download --title "Project Hail Mary" --author "Andy Weir"

  # 直接下载指定链接
  %(prog)s download --dl-path "/dl/O2pPQwkq2l" --title "Project Hail Mary" --format epub

  # 使用代理下载
  %(prog)s download --title "Project Hail Mary" --author "Andy Weir" --proxy-file proxies.txt
        """,
    )

    # 全局参数
    parser.add_argument("--json", action="store_true", help="JSON 格式输出（AI友好）")
    parser.add_argument("--verbose", "-v", action="store_true", help="详细输出")

    subparsers = parser.add_subparsers(dest="command", help="可用命令")

    # ---- search 子命令 ----
    search_p = subparsers.add_parser("search", help="搜索图书")
    search_p.add_argument("--title", "-t", default="", help="书名")
    search_p.add_argument("--author", "-a", default="", help="作者")
    search_p.add_argument("--isbn", default="", help="ISBN")
    search_p.add_argument("--publisher", default="", help="出版社")
    search_p.add_argument("--query", "-q", default="", help="自由搜索词")
    search_p.add_argument("--language", "-l", default=None, help="语言过滤 (English, Chinese)")
    search_p.add_argument("--max-results", type=int, default=50, help="最大结果数 (默认 50)")
    search_p.add_argument("--proxy", default=None, help="代理 URL (socks5://host:port)")

    # ---- download 子命令 ----
    dl_p = subparsers.add_parser("download", help="下载图书")
    dl_p.add_argument("--title", "-t", default="", help="书名")
    dl_p.add_argument("--author", "-a", default="", help="作者")
    dl_p.add_argument("--isbn", default="", help="ISBN")
    dl_p.add_argument("--publisher", default="", help="出版社")
    dl_p.add_argument("--query", "-q", default="", help="自由搜索词")
    dl_p.add_argument("--dl-path", nargs="*", default=None, help="直接指定下载路径 (/dl/xxx)")
    dl_p.add_argument("--format", dest="fmt", default="", help="指定格式 (epub, pdf)")
    dl_p.add_argument("--output-dir", "-o", default="./downloads", help="下载目录 (默认 ./downloads)")
    dl_p.add_argument("--proxy-file", "-p", default=None, help="代理列表文件")
    dl_p.add_argument("--proxy", default=None, help="单个代理 URL (socks5://host:port)")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(0)

    if args.command == "search":
        cmd_search(args)
    elif args.command == "download":
        cmd_download(args)


if __name__ == "__main__":
    main()