NeteaseUrl/services/batch.py

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

打包骨架（状态机 / 暂停 / 名额 / 终态收敛 / 结算）在 :mod:`services._zipcore`，
本模块只保留「音乐歌单」这一侧的策略：窗口化预取与并行下载、每首三条目
（音频 + 歌词 + 封面）入包、同名去重、包名与文案。
"""
import asyncio
import zipfile
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_sem: asyncio.Semaphore | None = None
        self._info_tasks: dict[Any, asyncio.Task[Any]] = {}

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

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

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

    async def _load_info(self, mid: Any, quality: str, cookies: dict[str, str]) -> Any:
        """取单曲信息；并发由 _info_sem 钳制（避免预取风暴打风控）。"""
        assert self._info_sem is not None
        async with self._info_sem:
            return await api_service.downloader.get_music_info(mid, quality, cookies)

    async def download_one(self, item: Any, idx: int, ctx: ZipContext) -> _SongPayload:
        """并行下载单元：等待预取信息 → 下载 → 复用预取歌词与缓存封面（零重复请求）。"""
        mid = item
        info = await self._info_tasks[mid]
        params = ctx.params
        result = await 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 = await 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)

    async 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。音频条目可达上百 MB，落盘经 to_thread
        # 卸载，避免在事件循环里做整段文件拷贝
        try:
            await asyncio.to_thread(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"

    async def cleanup(self) -> None:
        """取消在途预取任务（窗口因暂停/取消被跳过时它们没人 await）。

        取消后不再 await：asyncio 任务被 cancel 不会产生"异常未被取用"告警，
        而 await 会在此处引入最长一次上游请求的等待。
        """
        for task in self._info_tasks.values():
            if not task.done():
                task.cancel()
        self._info_tasks.clear()


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

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