AstrBot/astrbot/core/utils/media_utils.py

"""媒体文件处理工具

提供音视频格式转换、时长获取等功能。
"""

import asyncio
import base64
import io
import os
import subprocess
import uuid
from pathlib import Path

from PIL import Image as PILImage

from astrbot import logger
from astrbot.core.utils.astrbot_path import get_astrbot_temp_path
from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav

IMAGE_COMPRESS_DEFAULT_MAX_SIZE = 1280
IMAGE_COMPRESS_DEFAULT_QUALITY = 95
IMAGE_COMPRESS_DEFAULT_OPTIMIZE = True
IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_MB = 1.0


async def get_media_duration(file_path: str) -> int | None:
    """使用ffprobe获取媒体文件时长

    Args:
        file_path: 媒体文件路径

    Returns:
        时长（毫秒），如果获取失败返回None
    """
    try:
        # 使用ffprobe获取时长
        process = await asyncio.create_subprocess_exec(
            "ffprobe",
            "-v",
            "error",
            "-show_entries",
            "format=duration",
            "-of",
            "default=noprint_wrappers=1:nokey=1",
            file_path,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        stdout, stderr = await process.communicate()

        if process.returncode == 0 and stdout:
            duration_seconds = float(stdout.decode().strip())
            duration_ms = int(duration_seconds * 1000)
            logger.debug(f"[Media Utils] 获取媒体时长: {duration_ms}ms")
            return duration_ms
        else:
            logger.warning(f"[Media Utils] 无法获取媒体文件时长: {file_path}")
            return None

    except FileNotFoundError:
        logger.warning(
            "[Media Utils] ffprobe未安装或不在PATH中，无法获取媒体时长。请安装ffmpeg: https://ffmpeg.org/"
        )
        return None
    except Exception as e:
        logger.warning(f"[Media Utils] 获取媒体时长时出错: {e}")
        return None


async def convert_audio_to_opus(audio_path: str, output_path: str | None = None) -> str:
    """将音频转换为opus格式。"""
    return await convert_audio_format(
        audio_path=audio_path,
        output_format="opus",
        output_path=output_path,
    )


async def convert_video_format(
    video_path: str, output_format: str = "mp4", output_path: str | None = None
) -> str:
    """使用ffmpeg转换视频格式

    Args:
        video_path: 原始视频文件路径
        output_format: 目标格式，默认mp4
        output_path: 输出文件路径，如果为None则自动生成

    Returns:
        转换后的视频文件路径

    Raises:
        Exception: 转换失败时抛出异常
    """
    # 如果已经是目标格式，直接返回
    if video_path.lower().endswith(f".{output_format}"):
        return video_path

    # 生成输出文件路径
    if output_path is None:
        temp_dir = get_astrbot_temp_path()
        os.makedirs(temp_dir, exist_ok=True)
        output_path = os.path.join(
            temp_dir,
            f"media_video_{uuid.uuid4().hex}.{output_format}",
        )

    try:
        # 使用ffmpeg转换视频格式
        process = await asyncio.create_subprocess_exec(
            "ffmpeg",
            "-y",
            "-i",
            video_path,
            "-c:v",
            "libx264",
            "-c:a",
            "aac",
            output_path,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        stdout, stderr = await process.communicate()

        if process.returncode != 0:
            # 清理可能已生成但无效的临时文件
            if output_path and os.path.exists(output_path):
                try:
                    os.remove(output_path)
                    logger.debug(
                        f"[Media Utils] 已清理失败的{output_format}输出文件: {output_path}"
                    )
                except OSError as e:
                    logger.warning(
                        f"[Media Utils] 清理失败的{output_format}输出文件时出错: {e}"
                    )

            error_msg = stderr.decode() if stderr else "未知错误"
            logger.error(f"[Media Utils] ffmpeg转换视频失败: {error_msg}")
            raise Exception(f"ffmpeg conversion failed: {error_msg}")

        logger.debug(f"[Media Utils] 视频转换成功: {video_path} -> {output_path}")
        return output_path

    except FileNotFoundError:
        logger.error(
            "[Media Utils] ffmpeg未安装或不在PATH中，无法转换视频格式。请安装ffmpeg: https://ffmpeg.org/"
        )
        raise Exception("ffmpeg not found")
    except Exception as e:
        logger.error(f"[Media Utils] 转换视频格式时出错: {e}")
        raise


async def convert_audio_format(
    audio_path: str,
    output_format: str = "amr",
    output_path: str | None = None,
) -> str:
    """使用ffmpeg将音频转换为指定格式。

    Args:
        audio_path: 原始音频文件路径
        output_format: 目标格式，例如 amr / ogg / opus / wav
        output_path: 输出文件路径，如果为None则自动生成

    Returns:
        转换后的音频文件路径
    """
    if audio_path.lower().endswith(f".{output_format}"):
        return audio_path

    if output_path is None:
        temp_dir = Path(get_astrbot_temp_path())
        temp_dir.mkdir(parents=True, exist_ok=True)
        output_path = str(temp_dir / f"media_audio_{uuid.uuid4().hex}.{output_format}")

    args = ["ffmpeg", "-y", "-i", audio_path]
    if output_format == "amr":
        args.extend(
            [
                "-ac",
                "1",
                "-ar",
                "8000",
                "-ab",
                "12.2k",
                "-af",
                (
                    "highpass=f=310:poles=2,"
                    "lowpass=f=3720:poles=2,"
                    "equalizer=f=3150:width_type=h:width=1000:g=7.5,"
                    "loudnorm=I=-18.5:TP=-1.5:LRA=6,"
                    "aresample=8000"
                ),
            ]
        )
    elif output_format == "ogg":
        args.extend(["-acodec", "libopus", "-ac", "1", "-ar", "16000"])
    elif output_format == "opus":
        args.extend(["-acodec", "libopus", "-ac", "1", "-ar", "16000"])
    args.append(output_path)

    try:
        process = await asyncio.create_subprocess_exec(
            *args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        _, stderr = await process.communicate()
        if process.returncode != 0:
            if output_path and os.path.exists(output_path):
                try:
                    os.remove(output_path)
                except OSError as e:
                    logger.warning(f"[Media Utils] 清理失败的音频输出文件时出错: {e}")
            error_msg = stderr.decode() if stderr else "未知错误"
            raise Exception(f"ffmpeg conversion failed: {error_msg}")
        logger.debug(f"[Media Utils] 音频转换成功: {audio_path} -> {output_path}")
        return output_path
    except FileNotFoundError:
        raise Exception("ffmpeg not found")


async def convert_audio_to_amr(audio_path: str, output_path: str | None = None) -> str:
    """将音频转换为amr格式。"""
    return await convert_audio_format(
        audio_path=audio_path,
        output_format="amr",
        output_path=output_path,
    )


async def convert_audio_to_wav(audio_path: str, output_path: str | None = None) -> str:
    """将音频转换为wav格式。"""
    return await convert_audio_format(
        audio_path=audio_path,
        output_format="wav",
        output_path=output_path,
    )


async def ensure_wav(audio_path: str, output_path: str | None = None) -> str:
    """Ensure the audio path points to wav format by extension/guess and convert when needed.

    If the file appears to already be wav, return it directly to avoid extra conversion.
    """

    if not audio_path:
        return audio_path

    if not os.path.exists(audio_path):
        # File not available yet (e.g. napcat race condition);
        # return the path as-is so upstream retry logic can handle it later.
        return audio_path

    audio_type = _get_audio_magic_type(audio_path)
    if audio_type == "wav":
        return audio_path

    if audio_type == "silk":
        if output_path is None:
            temp_dir = get_astrbot_temp_path()
            os.makedirs(temp_dir, exist_ok=True)
            output_path = os.path.join(temp_dir, f"media_audio_{uuid.uuid4().hex}.wav")
        return await tencent_silk_to_wav(audio_path, output_path)

    return await convert_audio_to_wav(audio_path, output_path)


def _get_audio_magic_type(audio_path: str) -> str:
    """Detect common audio formats from magic bytes."""
    try:
        with open(audio_path, "rb") as f:
            header = f.read(64)
    except FileNotFoundError:
        logger.warning(f"[Media Utils] wav check file not found: {audio_path}")
        return ""
    except Exception as e:
        logger.warning(f"[Media Utils] wav check failed: {audio_path}, error: {e}")
        return ""

    if len(header) < 12:
        return ""

    if header[:4] == b"RIFF" and header[8:12] == b"WAVE":
        return "wav"

    if header[:4] == b"#!AM":
        return "amr"

    if header[:4] == b"OggS":
        if b"OpusHead" in header:
            return "opus"
        return "ogg"

    if header[:3] == b"fLa":
        return "flac"

    if header[:3] == b"ID3" or header[:2] == b"\xff\xfb":
        return "mp3"

    if header[:4] == b"ftyp" and b"mp4" in header[:8]:
        return "mp4"

    if header.startswith(b"#!SILK_V3"):
        return "silk"

    # Tencent SILK: leading \x02 byte before #!SILK_V3
    if header.startswith(b"\x02#!SILK_V3"):
        return "silk"

    return ""


async def extract_video_cover(
    video_path: str,
    output_path: str | None = None,
) -> str:
    """从视频中提取封面图(JPG)"""
    if output_path is None:
        temp_dir = Path(get_astrbot_temp_path())
        temp_dir.mkdir(parents=True, exist_ok=True)
        output_path = str(temp_dir / f"media_cover_{uuid.uuid4().hex}.jpg")

    try:
        process = await asyncio.create_subprocess_exec(
            "ffmpeg",
            "-y",
            "-i",
            video_path,
            "-ss",
            "00:00:00",
            "-frames:v",
            "1",
            output_path,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        _, stderr = await process.communicate()
        if process.returncode != 0:
            if output_path and os.path.exists(output_path):
                try:
                    os.remove(output_path)
                except OSError as e:
                    logger.warning(f"[Media Utils] 清理失败的视频封面文件时出错: {e}")
            error_msg = stderr.decode() if stderr else "未知错误"
            raise Exception(f"ffmpeg extract cover failed: {error_msg}")
        return output_path
    except FileNotFoundError:
        raise Exception("ffmpeg not found")


def _compress_image_sync(
    data: bytes,
    temp_dir: Path,
    max_size: int,
    quality: int,
    optimize: bool,
) -> str:
    """Run image compression synchronously via ``asyncio.to_thread``."""
    with PILImage.open(io.BytesIO(data)) as opened_img:
        img = opened_img
        converted_img: PILImage.Image | None = None

        try:
            if img.mode != "RGB":
                converted_img = img.convert("RGB")
                img = converted_img

            if max(img.size) > max_size:
                img.thumbnail((max_size, max_size), PILImage.Resampling.LANCZOS)

            new_uuid = uuid.uuid4().hex
            save_path = temp_dir / f"compressed_{new_uuid}.jpg"
            img.save(save_path, "JPEG", quality=quality, optimize=optimize)
            logger.debug(f"Image compressed successfully: {save_path}")
            return str(save_path)
        finally:
            if converted_img is not None:
                converted_img.close()


async def compress_image(
    url_or_path: str,
    max_size: int = IMAGE_COMPRESS_DEFAULT_MAX_SIZE,
    quality: int = IMAGE_COMPRESS_DEFAULT_QUALITY,
) -> str:
    """Compress large user-uploaded images.

    Args:
        url_or_path: Image path or URL.
        max_size: Longest edge of the compressed image in pixels.
        quality: JPEG output quality in the range 1-100.

    Returns:
        The compressed image path. Returns the original path if compression
        fails or the source does not need compression.
    """
    max_size = max(int(max_size), 1)
    quality = min(max(int(quality), 1), 100)
    optimize = IMAGE_COMPRESS_DEFAULT_OPTIMIZE
    min_file_size_bytes = int(IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_MB * 1024 * 1024)
    data = None

    def _exceeds_max_size(source: bytes | Path) -> bool:
        try:
            fp = io.BytesIO(source) if isinstance(source, bytes) else source
            with PILImage.open(fp) as opened_img:
                return max(opened_img.size) > max_size
        except Exception:  # noqa: BLE001
            return False

    # Skip compression for remote images and return the original value.
    if url_or_path.startswith("http"):
        return url_or_path
    elif url_or_path.startswith("data:image"):
        _header, encoded = url_or_path.split(",", 1)
        data = base64.b64decode(encoded)
        if len(data) < min_file_size_bytes and not _exceeds_max_size(data):
            return url_or_path
    else:
        local_path = Path(url_or_path)
        if not local_path.exists():
            return url_or_path
        if local_path.stat().st_size < min_file_size_bytes and not _exceeds_max_size(
            local_path
        ):
            return url_or_path
        with local_path.open("rb") as f:
            data = f.read()

    if not data:
        return url_or_path

    temp_dir = Path(get_astrbot_temp_path())
    temp_dir.mkdir(parents=True, exist_ok=True)

    # Offload the blocking image processing task to a thread.
    return await asyncio.to_thread(
        _compress_image_sync,
        data,
        temp_dir,
        max_size,
        quality,
        optimize,
    )