jmcomic-bridge.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
JMComic-Crawler-Python 桥接脚本（供 Yunzai 插件经 child_process 调用）。
命令列表见 main()，直接运行可查看用法。

环境变量:
    JM_PROXY              - HTTP 代理地址，如 http://127.0.0.1:7890
    JM_BASE_URL           - 自定义 API 域名列表（JSON 数组字符串）
    JM_PDF_CACHE_ENABLED  - PDF 缓存开关，默认 true

约定: 成功 → stdout 输出 { "success": true, ... }
      失败 → stderr 输出 { "success": false, "error": "..." }，退出码非 0
"""

import sys
import os
import json
import math

sys.stdout.reconfigure(encoding='utf-8')
sys.stderr.reconfigure(encoding='utf-8')

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
if SCRIPT_DIR not in sys.path:
    sys.path.insert(0, SCRIPT_DIR)

try:
    import jmcomic  # noqa: E402
except ImportError as e:
    print(json.dumps({
        "success": False,
        "error": f"未安装 jmcomic 库，请先执行: pip install jmcomic -U（{e}）"
    }, ensure_ascii=False), file=sys.stderr)
    sys.exit(1)


# ==================== 基础工具 ====================

def json_out(data):
    """输出成功 JSON 到 stdout"""
    print(json.dumps(data, ensure_ascii=False, default=str))


def err_exit(msg, code=1):
    """输出错误 JSON 到 stderr 并退出"""
    print(json.dumps({"success": False, "error": msg}, ensure_ascii=False), file=sys.stderr)
    sys.exit(code)


def build_option_dict(cookies_dict=None, download_dir=None):
    """基于默认配置构造 Option 字典，支持 JM_PROXY / JM_BASE_URL 环境变量"""
    d = json.loads(json.dumps(jmcomic.JmModuleConfig.DEFAULT_OPTION_DICT))

    # 下载目录（绝对路径）
    base_dir = os.path.abspath(download_dir or './jmcomic-downloads')
    d['dir_rule']['base_dir'] = base_dir

    # 使用 API 实现（移动端接口，IP 兼容性好）
    d['client']['impl'] = 'api'

    # 默认配置中 photo 线程数为 None，修复为 1
    if d['download']['threading']['photo'] is None:
        d['download']['threading']['photo'] = 1

    if cookies_dict:
        d['client']['postman']['meta_data']['cookies'] = cookies_dict

    proxy = os.environ.get('JM_PROXY', '')
    if proxy:
        d['client']['postman']['meta_data']['proxies'] = proxy

    base_url = os.environ.get('JM_BASE_URL', '')
    if base_url:
        try:
            domains = json.loads(base_url)
            if isinstance(domains, list):
                d['client']['domain'] = domains
        except json.JSONDecodeError:
            pass

    return d


def build_client(cookies_dict=None, download_dir=None):
    """构建 JmApiClient 实例"""
    d = build_option_dict(cookies_dict, download_dir)
    opt = jmcomic.JmOption.construct(d)
    return opt.build_jm_client()


def get_cookies(client):
    """从客户端提取登录后的 cookies"""
    rp = client.get_root_postman()
    return rp.meta_data.get('cookies', {})


# ==================== 命令实现 ====================

def cmd_login(args):
    if len(args) < 2:
        err_exit('用法: bridge.py login <用户名> <密码>')

    username, password = args[0], args[1]

    try:
        client = build_client()
        client.login(username, password)
        cookies = get_cookies(client)

        if not cookies or len(cookies) == 0:
            err_exit('登录失败：未能获取 Cookie（请检查用户名/密码或网络）')

        json_out({
            "success": True,
            "username": username,
            "cookie": cookies,
        })
    except Exception as e:
        err_exit(f"登录异常: {e}")


def cmd_favorites(args):
    """单页收藏列表（每页 20 条，对应禁漫收藏夹分页）"""
    if len(args) < 1:
        err_exit('用法: bridge.py favorites <cookies_json> [page]')

    cookies = json.loads(args[0])
    page = int(args[1]) if len(args) > 1 else 1

    try:
        client = build_client(cookies)
        fpage = client.favorite_folder(page=page)

        favorites_list = [{
            "id": item[0],
            "title": item[1],
            "coverUrl": jmcomic.JmcomicText.get_album_cover_url(item[0]),
        } for item in fpage.iter_id_title()]

        folders = [{"id": f[0], "name": f[1]} for f in fpage.iter_folder_id_name()]

        json_out({
            "success": True,
            "favorites": favorites_list,
            "folders": folders,
            "page": page,
            "total": getattr(fpage, 'total', 0),
            "totalPages": getattr(fpage, 'page_count', 1),
        })
    except Exception as e:
        err_exit(f"获取收藏异常: {e}")


def cmd_add_favorite(args):
    """添加收藏（官方 add_favorite_album 接口）"""
    if len(args) < 2:
        err_exit('用法: bridge.py add <cookies_json> <album_id> [folder_id]')

    cookies = json.loads(args[0])
    album_id = args[1]
    folder_id = args[2] if len(args) > 2 else '0'

    try:
        client = build_client(cookies)
        client.add_favorite_album(album_id, folder_id)
        json_out({"success": True})
    except Exception as e:
        err_exit(f"添加收藏异常: {e}")


def cmd_search(args):
    """站内搜索（search_site，main_tag=0）"""
    if len(args) < 1:
        err_exit('用法: bridge.py search <关键词> [page]')

    keyword = args[0]
    page = int(args[1]) if len(args) > 1 else 1

    try:
        client = build_client()
        result = client.search_site(search_query=keyword, page=page)

        results_list = [{
            "id": item[0],
            "title": item[1],
            "coverUrl": jmcomic.JmcomicText.get_album_cover_url(item[0]),
        } for item in result.iter_id_title()]

        json_out({
            "success": True,
            "results": results_list,
            "total": getattr(result, 'total', 0),
            "totalPages": getattr(result, 'page_count', 1),
            "page": page,
        })
    except Exception as e:
        err_exit(f"搜索异常: {e}")


def cmd_detail(args):
    """漫画详情"""
    if len(args) < 1:
        err_exit('用法: bridge.py detail <album_id>')

    album_id = args[0]

    try:
        client = build_client()
        detail = client.get_album_detail(album_id)

        episode_count = len(detail) if hasattr(detail, '__len__') else 1

        json_out({
            "success": True,
            "id": detail.id,
            "title": detail.title,
            "author": getattr(detail, 'author', ''),
            "authors": getattr(detail, 'authors', []),
            "coverUrl": jmcomic.JmcomicText.get_album_cover_url(album_id),
            "description": getattr(detail, 'description', None),
            "likes": getattr(detail, 'likes', None),
            "views": getattr(detail, 'views', None),
            "commentCount": getattr(detail, 'comment_count', None),
            "pageCount": getattr(detail, 'page_count', 0),
            "episodeCount": episode_count,
            "tags": getattr(detail, 'tags', []),
            "works": getattr(detail, 'works', []),
            "actors": getattr(detail, 'actors', []),
            "pubDate": getattr(detail, 'pub_date', None),
            "updateDate": getattr(detail, 'update_date', None),
        })
    except Exception as e:
        err_exit(f"获取详情异常: {e}")


def cmd_categories(args):
    """分类列表（categories_filter）"""
    try:
        from jmcomic import JmMagicConstants

        page = int(args[0]) if len(args) > 0 else 1
        time = args[1] if len(args) > 1 else JmMagicConstants.TIME_ALL
        category = args[2] if len(args) > 2 else JmMagicConstants.CATEGORY_ALL
        order_by = args[3] if len(args) > 3 else JmMagicConstants.ORDER_BY_VIEW

        client = build_client()
        result = client.categories_filter(
            page=page, time=time, category=category, order_by=order_by
        )

        items = [{
            "id": item[0],
            "title": item[1],
            "coverUrl": jmcomic.JmcomicText.get_album_cover_url(item[0]),
        } for item in result.iter_id_title()]

        json_out({
            "success": True,
            "categories": items,
            "page": page,
            "total": getattr(result, 'total', 0),
            "totalPages": getattr(result, 'page_count', 1),
        })
    except Exception as e:
        err_exit(f"获取分类异常: {e}")


def cmd_ranking(args):
    """排行榜：日榜/周榜/月榜（底层均为 categories_filter + ORDER_BY_VIEW）"""
    try:
        from jmcomic import JmMagicConstants

        page = int(args[0]) if len(args) > 0 else 1
        rank_type = args[1] if len(args) > 1 else 'week'

        time_map = {
            'day': JmMagicConstants.TIME_TODAY,
            'week': JmMagicConstants.TIME_WEEK,
            'month': JmMagicConstants.TIME_MONTH,
        }
        time_val = time_map.get(rank_type, JmMagicConstants.TIME_TODAY)

        client = build_client()
        result = client.categories_filter(
            page=page,
            time=time_val,
            category=JmMagicConstants.CATEGORY_ALL,
            order_by=JmMagicConstants.ORDER_BY_VIEW,
        )

        items = [{
            "id": item[0],
            "title": item[1],
            "coverUrl": jmcomic.JmcomicText.get_album_cover_url(item[0]),
        } for item in result.iter_id_title()]

        json_out({
            "success": True,
            "ranking": items,
            "page": page,
            "type": rank_type,
            "total": getattr(result, 'total', 0),
            "totalPages": getattr(result, 'page_count', 1),
        })
    except Exception as e:
        err_exit(f"获取排行榜异常: {e}")


def cmd_download(args):
    """下载漫画到服务器本地（支持登录态下载）"""
    if len(args) < 2:
        err_exit('用法: bridge.py download <cookies_json> <album_id> [download_dir]')

    cookies = json.loads(args[0]) if args[0] else {}
    comic_id = args[1]
    download_dir = args[2] if len(args) > 2 else None

    try:
        d = build_option_dict(cookies, download_dir)
        base_dir = d['dir_rule']['base_dir']
        opt = jmcomic.JmOption.construct(d)

        from jmcomic import JmDownloader
        dler = JmDownloader(opt)
        album = dler.download_album(comic_id)

        # 统计下载的图片数量
        image_count = 0
        for root, _dirs, files in os.walk(base_dir):
            for f in files:
                if f.lower().endswith(('.jpg', '.jpeg', '.png', '.gif', '.webp')):
                    image_count += 1

        json_out({
            "success": True,
            "albumId": album.id if hasattr(album, 'id') else comic_id,
            "title": album.title if hasattr(album, 'title') else str(comic_id),
            "author": getattr(album, 'author', ''),
            "totalPages": getattr(album, 'page_count', 0),
            "downloadedFiles": image_count,
            "downloadDir": base_dir,
        })
    except Exception as e:
        err_exit(f"下载异常: {e}")


# ==================== 图片转 PDF ====================

def images_to_pdf(images, pdf_path):
    """将图片列表转为单个 PDF（优先 img2pdf，未安装时回退 Pillow）"""
    if not images:
        return False
    try:
        import img2pdf
        with open(pdf_path, 'wb') as f:
            f.write(img2pdf.convert(images))
        return True
    except ImportError:
        from PIL import Image
        pil_images = []
        for img_path in images:
            try:
                img = Image.open(img_path)
                if img.mode in ('RGBA', 'P', 'LA'):
                    img = img.convert('RGB')
                pil_images.append(img)
            except Exception:
                pass
        if pil_images:
            pil_images[0].save(pdf_path, save_all=True, append_images=pil_images[1:], optimize=True)
            for img in pil_images:
                img.close()
            return True
        return False


def cleanup_images(base_dir):
    """递归删除 base_dir 下的所有图片文件，并清理空目录（保留 PDF 产物）"""
    for root, _dirs, files in os.walk(base_dir, topdown=False):
        for f in files:
            if f.lower().endswith(('.jpg', '.jpeg', '.png', '.webp', '.gif')):
                try:
                    os.remove(os.path.join(root, f))
                except Exception:
                    pass
    for root, dirs, files in os.walk(base_dir, topdown=False):
        if root == base_dir:
            continue
        try:
            if not os.listdir(root):
                os.rmdir(root)
        except Exception:
            pass


def enforce_cache_limit(base_dir, max_pdf):
    """限制 PDF 与封面缓存数量（同一上限，0 = 不限制），删除最早的文件"""
    if max_pdf > 0:
        # PDF 清理
        pdfs = []
        for root, _dirs, files in os.walk(base_dir):
            for f in files:
                if f.lower().endswith('.pdf'):
                    fp = os.path.join(root, f)
                    pdfs.append((os.path.getmtime(fp), fp))
        pdfs.sort()
        while len(pdfs) > max_pdf:
            _, old = pdfs.pop(0)
            try:
                os.remove(old)
            except Exception:
                pass

        # 封面清理（与 PDF 同一上限）
        covers = []
        for root, _dirs, files in os.walk(base_dir):
            for f in files:
                if f.lower().startswith('cover_') and f.lower().endswith('.jpg'):
                    fp = os.path.join(root, f)
                    covers.append((os.path.getmtime(fp), fp))
        covers.sort()
        while len(covers) > max_pdf:
            _, old = covers.pop(0)
            try:
                os.remove(old)
            except Exception:
                pass


def get_album_title(comic_id):
    """快速获取漫画标题（用于缓存命中时补标题）"""
    try:
        client = build_client()
        detail = client.get_album_detail(comic_id)
        return getattr(detail, 'title', f'#{comic_id}'), getattr(detail, 'author', '')
    except Exception:
        return f'#{comic_id}', ''


def cached_json_out(comic_id, **extra):
    """缓存命中时返回，自动补全标题"""
    title, author = get_album_title(comic_id)
    json_out({"success": True, "albumId": comic_id, "title": title, "author": author, "cached": True, **extra})


def cmd_pdf(args):
    """下载漫画并生成 PDF（保存到本地）。

    规则: 单章→单个 PDF；多章→每章一个；单章过大(>200页或>1GB)→分片；
          PDF 已存在则返回缓存；生成后清理图片；缓存上限默认 20（0 = 不限制）。
    """
    if len(args) < 2:
        err_exit('用法: bridge.py pdf <cookies_json> <album_id> [download_dir] [max_pdfs]')

    cookies = json.loads(args[0]) if args[0] else {}
    comic_id = args[1]
    download_dir = args[2] if len(args) > 2 else None
    max_pdfs = int(args[3]) if len(args) > 3 and args[3] else 20

    try:
        d = build_option_dict(cookies, download_dir)
        base_dir = d['dir_rule']['base_dir']
        os.makedirs(base_dir, exist_ok=True)
        opt = jmcomic.JmOption.construct(d)

        # === 缓存检查（命中已有 PDF 直接返回，不重新下载） ===
        cache_enabled = os.environ.get('JM_PDF_CACHE_ENABLED', 'true').lower() == 'true'
        if cache_enabled:
            import glob as g
            patterns = [
                f'{comic_id}.pdf',
                f'{comic_id}_p*.pdf',
                f'{comic_id}_ch*.pdf',
            ]
            cached = {}
            for pat in patterns:
                for fp in g.glob(os.path.join(base_dir, pat)):
                    if os.path.getsize(fp) > 0:
                        cached[os.path.basename(fp)] = fp

            if cached:
                pdfs = sorted([p for n, p in cached.items() if n.endswith('.pdf')])
                if pdfs:
                    for p in pdfs:
                        os.utime(p, None)
                    files = [{"path": p, "size": os.path.getsize(p),
                              "chapterTitle": os.path.basename(p) if len(pdfs) > 1 else None} for p in pdfs]
                    cached_json_out(comic_id, chapterCount=len(pdfs), totalImages=0,
                                    files=files, isZip=False, zipPath=None, zipSize=0)
                    return

        from jmcomic import JmDownloader
        dler = JmDownloader(opt)
        album = dler.download_album(comic_id)

        # 章节→图片 映射（download_success_dict: {Album: {Photo: [(path, img), ...]}}）
        chapter_images = {}   # {章节序号: [图片路径]}
        chapter_titles = {}   # {章节序号: 章节名}
        for _album, photo_dict in dler.download_success_dict.items():
            for photo, image_list in photo_dict.items():
                ch_idx = getattr(photo, 'index', 0) or 0
                paths = [p for p, _ in sorted(image_list, key=lambda x: os.path.basename(x[0]))]
                if paths:
                    chapter_images[ch_idx] = paths
                    chapter_titles[ch_idx] = getattr(photo, 'name', f'第{ch_idx}章')

        # 若没有章节信息，回退到目录遍历
        if not chapter_images:
            images = []
            for root, _dirs, files in os.walk(base_dir):
                for f in sorted(files):
                    if f.lower().endswith(('.jpg', '.jpeg', '.png', '.webp', '.gif')):
                        images.append(os.path.join(root, f))
            if images:
                chapter_images[1] = images

        chapter_count = len(chapter_images)
        total_images = sum(len(v) for v in chapter_images.values())

        # 估算图片总大小
        total_img_size = 0
        for imgs in chapter_images.values():
            for p in imgs:
                try:
                    total_img_size += os.path.getsize(p)
                except Exception:
                    pass

        MAX_PAGES = 200        # 每 PDF 最多 200 张图
        SPLIT_SIZE = 200 * 1024 * 1024   # 200MB
        LARGE_THRESHOLD = 1024 * 1024 * 1024  # 1GB

        def should_split(imgs):
            """是否需要分包：页数>200 或 单章总大小>1GB"""
            if len(imgs) > MAX_PAGES:
                return True
            if chapter_count == 1 and total_img_size > LARGE_THRESHOLD:
                return True
            return False

        def chunk_images(imgs):
            """按 200MB 且 200 页为上限分片"""
            chunks = []
            current = []
            cur_size = 0
            for p in imgs:
                try:
                    sz = os.path.getsize(p)
                except Exception:
                    sz = 0
                if (cur_size + sz > SPLIT_SIZE or len(current) >= MAX_PAGES) and current:
                    chunks.append(current)
                    current = []
                    cur_size = 0
                current.append(p)
                cur_size += sz
            if current:
                chunks.append(current)
            return chunks

        # 生成每章 PDF
        pdf_files = []
        for ch_idx in sorted(chapter_images.keys()):
            imgs = chapter_images[ch_idx]

            if chapter_count == 1 and should_split(imgs):
                chunks = chunk_images(imgs)
                for ci, chunk in enumerate(chunks):
                    pdf_name = f'{comic_id}_p{ci + 1:02d}.pdf'
                    pdf_path = os.path.join(base_dir, pdf_name)
                    images_to_pdf(chunk, pdf_path)
                    pdf_files.append({
                        'path': pdf_path,
                        'size': os.path.getsize(pdf_path) if os.path.exists(pdf_path) else 0,
                        'chapterTitle': f'分片{ci + 1}/{len(chunks)}',
                    })
            elif chapter_count == 1:
                pdf_name = f'{comic_id}.pdf'
                pdf_path = os.path.join(base_dir, pdf_name)
                images_to_pdf(imgs, pdf_path)
                pdf_files.append({
                    'path': pdf_path,
                    'size': os.path.getsize(pdf_path) if os.path.exists(pdf_path) else 0,
                    'chapterTitle': None,
                })
            else:
                pdf_name = f'{comic_id}_ch{ch_idx:02d}.pdf'
                pdf_path = os.path.join(base_dir, pdf_name)
                images_to_pdf(imgs, pdf_path)
                ch_title = chapter_titles.get(ch_idx, f'第{ch_idx}章')
                pdf_files.append({
                    'path': pdf_path,
                    'size': os.path.getsize(pdf_path) if os.path.exists(pdf_path) else 0,
                    'chapterTitle': str(ch_title),
                })

        # 清除原始图片，保留 PDF
        cleanup_images(base_dir)

        # 缓存上限（PDF 与封面共用）
        cache_enabled = os.environ.get('JM_PDF_CACHE_ENABLED', 'true').lower() == 'true'
        if cache_enabled:
            enforce_cache_limit(base_dir, max_pdfs)

        json_out({
            "success": True,
            "albumId": album.id if hasattr(album, 'id') else comic_id,
            "title": album.title if hasattr(album, 'title') else str(comic_id),
            "author": getattr(album, 'author', ''),
            "chapterCount": chapter_count,
            "totalImages": total_images,
            "files": pdf_files,
            "isZip": False,
            "zipPath": None,
            "zipSize": 0,
            "downloadDir": base_dir,
        })
    except Exception as e:
        err_exit(f"PDF 生成异常: {e}")


def cmd_cover(args):
    """下载封面到本地，返回本地路径（供聊天内发图）"""
    if len(args) < 1:
        err_exit('用法: bridge.py cover <album_id> [download_dir]')

    comic_id = args[0]
    download_dir = args[1] if len(args) > 1 else None

    try:
        base_dir = os.path.abspath(download_dir or './jmcomic-downloads')
        os.makedirs(base_dir, exist_ok=True)
        local_path = os.path.join(base_dir, f'cover_{comic_id}.jpg')

        if not os.path.exists(local_path) or os.path.getsize(local_path) == 0:
            client = build_client(download_dir=download_dir)
            client.download_album_cover(comic_id, local_path)

        json_out({"success": True, "localPath": local_path})
    except Exception as e:
        err_exit(f"封面下载异常: {e}")


# ==================== 命令映射 ====================

COMMANDS = {
    "login": cmd_login,
    "favorites": cmd_favorites,
    "fav": cmd_favorites,
    "add": cmd_add_favorite,
    "search": cmd_search,
    "detail": cmd_detail,
    "categories": cmd_categories,
    "ranking": cmd_ranking,
    "rank": cmd_ranking,
    "download": cmd_download,
    "dl": cmd_download,
    "pdf": cmd_pdf,
    "cover": cmd_cover,
}


def main():
    if len(sys.argv) < 2:
        err_exit(
            "JMComic Bridge Script\n\n"
            "Commands:\n"
            "  login <username> <password>          - 登录并返回 cookie\n"
            "  favorites <cookies_json> [page]      - 获取收藏列表（单页）\n"
            "  add <cookies_json> <album_id>        - 添加收藏\n"
            "  search <关键词> [page]                - 搜索漫画\n"
            "  detail <album_id>                    - 获取漫画详情\n"
            "  categories [page] [time] [cat] [ord] - 分类列表\n"
            "  ranking [page] [day|week|month]      - 排行榜\n"
            "  download <cookies_json> <id> [dir]   - 下载漫画到本地\n"
            "  pdf <cookies_json> <id> [dir] [max]  - 下载并转 PDF（单章/多章/分片）\n"
            "  cover <album_id> [dir]               - 下载封面到本地"
        )

    cmd = sys.argv[1].lower()
    rest = sys.argv[2:]

    if cmd not in COMMANDS:
        err_exit(f"未知命令: {cmd}. 可用: {', '.join(COMMANDS.keys())}")

    COMMANDS[cmd](rest)


if __name__ == "__main__":
    main()