NeteaseUrl/services/batch.py

# -*- coding: utf-8 -*-
"""歌单/专辑批量打包 worker。

打包骨架（状态机 / 暂停 / 名额 / 终态收敛 / 结算）在 :mod:`services._zipcore`，
本模块只保留「音乐歌单」这一侧的策略：窗口化预取与并行下载、每首三条目
（音频 + 歌词 + 封面）入包、同名去重、包名与文案。
"""
import zipfile
from concurrent.futures import (
    Future,
    ThreadPoolExecutor,
)
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from core.utils import format_lyric as _format_lyric
from music_downloader import (
    fetch_cover_bytes,
)

from ._shared import (
    _build_download_base,
)
from ._taskcore import (
    _settle_batch_billing,
)
from ._zipcore import (
    ItemFailed,
    ZipContext,
    ZipJob,
    run_zip_job,
)
from .api_service import (
    api_service,
)


@dataclass(frozen=True)
class _SongPayload:
    """一首歌下载完成后的入包素材（下载线程产出，落包线程消费）。"""

    path: str
    name: str
    artists: str
    file_type: str
    lyric: str
    cover: bytes


class _SonglistZipJob(ZipJob):
    """音乐歌单打包：窗口化并行下载，每首写 音频 + 歌词 + 封面 三条目。"""

    tag = '批量ZIP'
    noun = '首'
    item_noun = '歌曲'
    tmp_prefix = 'batch_zip_'
    start_msg = '下载打包中...'
    record_duration = True

    def __init__(self) -> None:
        self._info_pool: ThreadPoolExecutor | None = None
        self._info_futures: dict[Any, Future[Any]] = {}

    def prepare(self, params: dict[str, Any], task_id: str, failed: list[str]) -> list[Any]:
        self.concurrency = int(params['concurrency'])
        return list(params['ids'])

    def prefetch(self, window: list[Any], ctx: ZipContext) -> None:
        """按窗口预取歌曲信息。

        原实现一次性提交全部（最多 500 首 ×3 接口）预取任务，大歌单瞬间以 10 并发
        连发上游约 1500 次请求打风控；改为每个下载窗口只预取本窗口的歌（与下载节奏
        对齐），``download_one`` 语义不变（等预取结果再下载）。
        """
        if self._info_pool is None:
            self._info_pool = ThreadPoolExecutor(max_workers=min(4, max(1, self.concurrency)))
        quality = ctx.params['quality']
        cookies = ctx.params['cookies']
        for mid in window:
            if mid not in self._info_futures:
                self._info_futures[mid] = self._info_pool.submit(
                    api_service.downloader.get_music_info, mid, quality, cookies
                )

    def download_one(self, item: Any, idx: int, ctx: ZipContext) -> _SongPayload:
        """并行下载单元：等待预取信息 → 下载 → 复用预取歌词与缓存封面（零重复请求）。"""
        mid = item
        info = self._info_futures[mid].result()
        params = ctx.params
        result = api_service.downloader.download_music_file(
            mid, params['quality'],
            write_meta=params['write_meta'],
            meta_fields=params['meta_fields'],
            cookies=params['cookies'],
            cancel_event=ctx.cancel_ev,
            music_info=info,
        )
        lyric = ''
        cover = b''
        if result.success:
            # 歌词已随 get_music_info 预取返回，无需再调一次歌词接口
            lyric = getattr(info, 'lyric', '') or ''
            # 封面走共享缓存（同一专辑只拉一次网络）；取消后不再发起新请求
            if not ctx.cancel_ev.is_set() and info.pic_url:
                cover = fetch_cover_bytes(info.pic_url) or b''
        if not result.success:
            raise ItemFailed(str(mid))
        assert result.file_path is not None
        if not Path(result.file_path).exists():
            # 文件可能已被同名并发任务入包删除，按失败处理而非炸掉整个任务
            raise ItemFailed(str(mid))
        mi = result.music_info
        if mi is None:
            # 与上面 exists() 同口径：成功但缺元信息（上游详情缺失）按单曲失败计，
            # 不让 mi.name 抛 AttributeError 把整批转 error
            raise ItemFailed(str(mid))
        return _SongPayload(path=result.file_path, name=mi.name, artists=mi.artists,
                            file_type=mi.file_type, lyric=lyric, cover=cover)

    def write_entry(self, item: Any, idx: int, payload: _SongPayload, ctx: ZipContext) -> None:
        mid = item
        base = _build_download_base(ctx.params['filename_format'],
                                    {'name': payload.name, 'artist_string': payload.artists})
        lyric_format = ctx.params['lyric_format']
        # 同名去重：歌单里同一首歌出现两次/同名不同版本时，追加序号避免 ZIP 内重名
        # 互相覆盖（用户被扣 2 次只拿到 1 首）。F-C：音频/歌词/封面三个 arcname 都要
        # 查重——同名同歌手但音源扩展不同（flac vs mp3）时歌词/封面也会互相覆盖
        final_name = f"{base}.{payload.file_type}"
        lrc_name = f"{base}.{lyric_format}"
        jpg_name = f"{base}.jpg"
        k = 2
        while (final_name in ctx.used_arcnames
               or lrc_name in ctx.used_arcnames
               or jpg_name in ctx.used_arcnames):
            final_name = f"{base} ({k}).{payload.file_type}"
            lrc_name = f"{base} ({k}).{lyric_format}"
            jpg_name = f"{base} ({k}).jpg"
            k += 1
        ctx.used_arcnames.update((final_name, lrc_name, jpg_name))
        # F-D：exists 与 zf.write 之间的缓存竞态（另一任务缓存 miss 删文件）
        # 按单曲失败计，不再让整批转 error
        try:
            ctx.zf.write(payload.path, arcname=final_name, compress_type=zipfile.ZIP_STORED)
        except OSError as e:
            raise ItemFailed(f"{mid}(缓存文件读取失败: {e})") from e
        # 缓存文件是跨任务共享资源，任务不得删除；磁盘回收交给监控线程的下载缓存
        # 清扫（按文件年龄）。歌词与封面已在并行下载线程中拉取，这里直接写包（无网络等待）
        if payload.lyric:
            try:
                ctx.zf.writestr(lrc_name, _format_lyric(payload.lyric, lyric_format),
                                compress_type=zipfile.ZIP_DEFLATED)
            except Exception:
                pass
        if payload.cover:
            try:
                ctx.zf.writestr(jpg_name, payload.cover, compress_type=zipfile.ZIP_STORED)
            except Exception:
                pass

    def zip_filename(self, ok_count: int) -> str:
        return f"批量下载_{ok_count}首.zip"

    def cleanup(self) -> None:
        if self._info_pool is not None:
            try:
                self._info_pool.shutdown(wait=False, cancel_futures=True)
            except TypeError:
                self._info_pool.shutdown(wait=False)


def _batch_zip_worker(task_id: str) -> None:
    """后台批量打包：窗口化下载 → 写 ZIP（音频免压缩），支持取消与断线暂停

    - 断线暂停：设备心跳超时后，当前窗口下载完即释放 BATCH_ZIP_SEM 名额给下一位，
      保存进度进入暂停等待；重新打开网站（心跳恢复）后自动回到队列继续下载
    - 音频入包即删，缓解磁盘占用
    """
    run_zip_job(_SonglistZipJob(), task_id, settle=_settle_batch_billing)