main.py

"""
音频压缩 - 基于 FFmpeg 的音频压缩脚本（不忙脚本盒子）
档位式压缩强度，支持 mp3/aac/ogg/opus/m4a 输出，可调采样率与声道

盒子契约：
- 启动时把参数 JSON 路径作为命令行第一个参数传入（sys.argv[1]）
- params_form 自动生成启动参数界面，所有运行选项由盒子注入 params
- 定时任务（invoke_mode == "scheduled"）：无人值守，按默认/表单参数运行
- 节点联动（invoke_mode == "node"）：把结果信封写入 environment.output_json
"""
import json
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path
from typing import List

BASE_DIR = Path(__file__).parent

# 压缩输出格式（有损格式）
AUDIO_FORMATS = ["mp3", "aac", "ogg", "opus", "m4a"]

AUDIO_CODECS = {
    "mp3": "libmp3lame",
    "aac": "aac",
    "ogg": "libvorbis",
    "opus": "libopus",
    "m4a": "aac",
}

# 压缩强度档位 → 目标码率 kbps（值越大体积越小、音质越低）
COMPRESSION_LEVELS = {
    "light": 256,     # 轻度压缩: 高音质
    "standard": 192,  # 标准压缩
    "compress": 128,  # 高效压缩(默认)
    "heavy": 96,      # 强力压缩
    "extreme": 64,    # 极致压缩
}

# 采样率 / 声道选项（"original" = 跟随输入）
ALLOWED_SAMPLE_RATES = {"original", "44100", "32000", "24000", "16000"}
ALLOWED_CHANNELS = {"original", "stereo", "mono"}

# 右键可选的音频扩展名（与 bm-scripts-box-rc.toml 的 filters 一致）
AUDIO_EXTS = {
    ".mp3", ".wav", ".flac", ".ogg", ".aac", ".wma",
    ".m4a", ".opus", ".ape", ".aiff", ".m4b", ".oga",
}


class AudioCompressor:
    """音频压缩器（基于 FFmpeg）：档位式压缩强度，承载通用 subprocess 执行，只产数据。"""

    @staticmethod
    def _run(cmd, **kw):
        """执行命令，默认隐藏控制台窗口、按 UTF-8 容错解码。"""
        kw.setdefault("creationflags", getattr(subprocess, "CREATE_NO_WINDOW", 0))
        kw.setdefault("encoding", "utf-8")
        return subprocess.run(cmd, text=True, errors="replace", **kw)

    @staticmethod
    def _require_binaries():
        """校验 ffmpeg，缺则直接报错（供开局预检）。"""
        if not shutil.which("ffmpeg"):
            raise FileNotFoundError("未找到 FFmpeg，请确认盒子已安装 ffmpeg 依赖")

    def __init__(self, audios, config):
        """audios: 音频文件路径列表；config: 配置 dict（键见 _build_config）。"""
        self.audios = [a for a in audios if Path(a).exists()]
        c = config

        self.output_dir = str(c.get("output_dir") or "").strip()
        self.output_format = str(c.get("output_format") or "mp3").strip().lstrip(".").lower()
        if self.output_format not in AUDIO_FORMATS:
            raise ValueError(f"不支持的格式：{self.output_format}，支持：{', '.join(AUDIO_FORMATS)}")

        # 非法参数值回退默认
        self.compression_level = c.get("compression_level") if c.get("compression_level") in COMPRESSION_LEVELS else "compress"
        self.sample_rate = c.get("sample_rate") if c.get("sample_rate") in ALLOWED_SAMPLE_RATES else "original"
        self.channels = c.get("channels") if c.get("channels") in ALLOWED_CHANNELS else "original"
        self.overwrite = bool(c.get("overwrite"))

        self._require_binaries()
        self._ffmpeg = shutil.which("ffmpeg")

        if self.output_dir:
            os.makedirs(self.output_dir, exist_ok=True)

    def _get_output_path(self, input_path: str) -> str:
        """生成输出文件路径；同目录同扩展名时加 _compressed 后缀避免覆盖原文件"""
        stem = Path(input_path).stem
        ext = self.output_format

        if self.output_dir:
            return os.path.join(self.output_dir, f"{stem}.{ext}")

        parent = Path(input_path).parent
        out = os.path.join(parent, f"{stem}.{ext}")
        if os.path.abspath(out) == os.path.abspath(input_path):
            out = os.path.join(parent, f"{stem}_compressed.{ext}")
        return out

    def _build_command(self, input_path: str, output_path: str) -> List[str]:
        """构建 ffmpeg 命令"""
        cmd = [self._ffmpeg, "-y", "-i", input_path]

        # 编码器
        cmd += ["-c:a", AUDIO_CODECS[self.output_format]]

        # 采样率
        if self.sample_rate != "original":
            cmd += ["-ar", self.sample_rate]

        # 声道
        if self.channels != "original":
            cmd += ["-ac", "1" if self.channels == "mono" else "2"]

        # 压缩强度 → 目标码率
        cmd += ["-b:a", f"{COMPRESSION_LEVELS[self.compression_level]}k"]

        # 保留元数据，失败时输出错误信息
        cmd += ["-map_metadata", "0", "-loglevel", "error"]
        cmd.append(output_path)

        return cmd

    def _compress_single(self, audio_path: str, on_start=None) -> tuple:
        """压缩单个音频，返回 (路径, 状态, 信息)，状态: success/skipped/failed。"""
        try:
            output_path = self._get_output_path(audio_path)

            # 压缩总是重编码，无"同格式跳过"；仅按覆盖模式处理已存在输出
            if os.path.exists(output_path) and not self.overwrite:
                return audio_path, "skipped", "文件已存在"

            if on_start:
                on_start(audio_path)

            cmd = self._build_command(audio_path, output_path)
            proc = self._run(cmd, capture_output=True)
            if proc.returncode != 0:
                return audio_path, "failed", (proc.stderr or "压缩失败").strip()
            return audio_path, "success", output_path

        except Exception as e:
            return audio_path, "failed", str(e)

    def compress(self, on_start=None, on_done=None) -> dict:
        """顺序压缩全部音频，回调供展示；返回分组结果 dict。"""
        results = {"success": [], "skipped": [], "failed": [], "total": len(self.audios)}

        for path in self.audios:
            path, status, info = self._compress_single(path, on_start=on_start)
            if status == "success":
                results["success"].append((path, info))
            elif status == "skipped":
                results["skipped"].append((path, info))
            else:
                results["failed"].append((path, info))
            if on_done:
                on_done(path, status, info)

        return results

    @staticmethod
    def get_file_size(path: str) -> str:
        """获取文件大小（人性化显示）"""
        size = os.path.getsize(path)
        for unit in ["B", "KB", "MB", "GB"]:
            if size < 1024:
                return f"{size:.1f} {unit}"
            size /= 1024
        return f"{size:.1f} TB"


# ==================== 盒子契约接入 ====================

def _fix_encoding():
    """统一输出编码，避免 GBK 控制台下 emoji/中文报错（盒子已设 PYTHONUTF8=1）。"""
    for _s in (sys.stdout, sys.stderr):
        try:
            _s.reconfigure(encoding="utf-8", errors="replace")
        except (AttributeError, ValueError):
            pass


def _read_payload(param_path):
    """读取盒子传入的 JSON 参数文件；缺失/损坏返回 None。"""
    if not (param_path and Path(param_path).exists()):
        return None
    try:
        with open(param_path, "r", encoding="utf-8") as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError):
        return None


def _build_config(params) -> dict:
    """params → 压缩配置（仅取 TOML 声明的键，缺失回退默认值）。"""
    return {
        "compression_level": params.get("compression_level", "compress"),
        "output_format": params.get("output_format", "mp3"),
        "sample_rate": params.get("sample_rate", "original"),
        "channels": params.get("channels", "original"),
        "output_dir": params.get("output_dir", ""),
        "overwrite": params.get("overwrite", False),
    }


def _pick_audios(paths):
    """按扩展名过滤音频，返回 (audios, skipped)。"""
    audios, skipped = [], []
    for p in paths:
        if Path(p).suffix.lower() in AUDIO_EXTS:
            audios.append(p)
        else:
            skipped.append(p)
    return audios, skipped


def _write_envelope(output_json, code, msg, **extra):
    """把结果信封写入 output_json（code/msg 为保留字段，业务键顶层平铺）。"""
    if not output_json:
        return
    data = {"code": code, "msg": msg}
    data.update(extra)
    with open(output_json, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def _config_summary(config) -> str:
    """压缩参数一行摘要（manual / scheduled 展示用）。"""
    lv = {"light": "轻度(256k)", "standard": "标准(192k)", "compress": "高效(128k)",
          "heavy": "强力(96k)", "extreme": "极致(64k)"}
    sr = {"original": "原始", "44100": "44100Hz", "32000": "32000Hz",
          "24000": "24000Hz", "16000": "16000Hz"}
    ch = {"original": "原始", "stereo": "立体声", "mono": "单声道"}
    line1 = (f"压缩强度 {lv.get(config['compression_level'], config['compression_level'])}"
             f" · 输出 {config['output_format']}"
             f" · 采样率 {sr.get(config['sample_rate'], config['sample_rate'])}"
             f" · 声道 {ch.get(config['channels'], config['channels'])}")
    out = config["output_dir"] or "源文件所在目录"
    mode = "允许覆盖" if config["overwrite"] else "跳过已存在文件"
    return f"{line1}\n  📁 输出目录 {out} · {mode}"


def _wait_exit():
    """手动模式倒计时停留，便于查看结果（终端窗口由盒子保持）。"""
    width, total = 10, 5
    for i in range(total, 0, -1):
        filled = round(width * (total - i + 1) / total)
        bar = "█" * filled + "░" * (width - filled)
        print(f"\r  ⏳ {i}s {bar}  窗口即将关闭", end="")
        time.sleep(1)
    print("\r" + " " * 60, end="\r")


def main():
    _fix_encoding()

    if len(sys.argv) < 2:
        print("用法：python main.py <参数JSON路径>\n参数 JSON 由不忙脚本盒子生成并传入。")
        return

    payload = _read_payload(sys.argv[1])
    if payload is None:
        print("❌ 无法读取盒子参数文件")
        return

    env = payload.get("environment", {})
    params = payload.get("params", {})
    invoke_mode = env.get("invoke_mode", "manual")
    output_json = env.get("output_json")

    audios, skipped = _pick_audios(payload.get("data", {}).get("target_paths", []) or [])
    if skipped:
        print("  ⏭️ 忽略非音频: " + ", ".join(Path(p).name for p in skipped))

    if not audios:
        msg = "未选择有效的音频文件"
        print(f"❌ {msg}")
        if invoke_mode == "node":
            _write_envelope(output_json, 1, msg, output_paths=[], summary=msg)
        return

    config = _build_config(params)
    try:
        compressor = AudioCompressor(audios, config)
    except (FileNotFoundError, ValueError) as e:
        msg = str(e)
        print(f"❌ {msg}")
        if invoke_mode == "node":
            _write_envelope(output_json, 1, msg, output_paths=[], summary=msg)
        return

    if invoke_mode != "node":
        print("🗜️ 音频压缩 · 批量压缩")
        print(f"  {_config_summary(config)}")

    total = len(audios)
    started = [0]

    def on_start(path):
        started[0] += 1
        print(f"  ▶ ({started[0]}/{total}) 正在压缩: {Path(path).name}")

    def on_done(path, status, info):
        name = Path(path).name
        if status == "success":
            size = AudioCompressor.get_file_size(info)
            print(f"  ✅ {name} → {Path(info).name}（{size}）")
        elif status == "skipped":
            print(f"  ⏭️ {name}  {info}")
        else:
            print(f"  ❌ {name}  {(info or '未知错误').strip().splitlines()[0]}")

    result = compressor.compress(on_start=on_start, on_done=on_done)

    success_paths = [info for _, info in result["success"]]
    parts = [f"成功 {len(result['success'])} 个"]
    if result["skipped"]:
        parts.append(f"跳过 {len(result['skipped'])} 个")
    if result["failed"]:
        parts.append(f"失败 {len(result['failed'])} 个")
    summary = " · ".join(parts)

    # 信封回传：节点联动必须写；定时任务有 output_json 时记录结果
    if output_json:
        code = 0 if (success_paths or result["skipped"]) else 1
        msg = summary if code == 0 else "全部压缩失败"
        _write_envelope(output_json, code, msg, output_paths=success_paths, summary=summary)

    if invoke_mode == "node":
        return  # 节点模式无人值守，跑完即退

    print("  " + summary)
    if invoke_mode == "manual":
        _wait_exit()


if __name__ == "__main__":
    main()