Trixy7/trixy_core/music/formats/mp3.py

# -*- coding: utf-8 -*-
"""
MP3-Decoder.

Dekodiert MP3-Dateien zu PCM mit pydub oder ffmpeg.
"""

import asyncio
import io
import logging
import struct
from pathlib import Path
from typing import AsyncIterator

from trixy_core.music.formats.base import (
    AudioDecoder,
    AudioFormat,
    AudioInfo,
    DecodedAudio,
    DecoderError,
)


class MP3Decoder(AudioDecoder):
    """
    MP3-Decoder.

    Verwendet pydub (mit ffmpeg Backend) für das Dekodieren.
    Fallback: Direktes ffmpeg-Subprocess.
    """

    FORMAT = AudioFormat.MP3
    EXTENSIONS = {"mp3"}
    MIME_TYPES = {"audio/mpeg", "audio/mp3"}

    def __init__(self) -> None:
        super().__init__()
        self.logger = logging.getLogger(__name__)
        self._use_pydub = False
        self._use_ffmpeg = False

    def _check_dependencies(self) -> bool:
        """Prüft verfügbare Decoder-Backends."""
        # Versuche pydub
        try:
            from pydub import AudioSegment
            self._use_pydub = True
            return True
        except ImportError:
            pass

        # Versuche ffmpeg direkt
        import shutil
        if shutil.which("ffmpeg"):
            self._use_ffmpeg = True
            return True

        return False

    async def get_info(self, source: Path | bytes) -> AudioInfo:
        """Liest MP3-Informationen."""
        if isinstance(source, bytes):
            return await self._get_info_from_bytes(source)
        return await self._get_info_from_file(source)

    async def _get_info_from_file(self, path: Path) -> AudioInfo:
        """Liest Informationen aus MP3-Datei."""
        if not path.exists():
            raise DecoderError(f"Datei nicht gefunden: {path}")

        file_size = path.stat().st_size

        # Versuche Mutagen für Metadaten
        try:
            import mutagen.mp3
            audio = mutagen.mp3.MP3(path)

            return AudioInfo(
                format=AudioFormat.MP3,
                duration_ms=int(audio.info.length * 1000),
                sample_rate=audio.info.sample_rate,
                channels=2 if audio.info.mode < 3 else 1,  # 0-2=Stereo, 3=Mono
                bit_rate=audio.info.bitrate,
                file_size=file_size,
            )
        except ImportError:
            pass
        except Exception as e:
            self.logger.debug(f"Mutagen-Fehler: {e}")

        # Fallback: Basis-Info aus Header
        try:
            info = await asyncio.to_thread(self._parse_mp3_header, path)
            info.file_size = file_size
            return info
        except Exception as e:
            raise DecoderError(f"Kann MP3-Info nicht lesen: {e}")

    async def _get_info_from_bytes(self, data: bytes) -> AudioInfo:
        """Liest Informationen aus MP3-Bytes."""
        try:
            info = self._parse_mp3_header_bytes(data)
            info.file_size = len(data)
            return info
        except Exception as e:
            raise DecoderError(f"Kann MP3-Info nicht lesen: {e}")

    def _parse_mp3_header(self, path: Path) -> AudioInfo:
        """Parst MP3-Frame-Header."""
        with open(path, "rb") as f:
            data = f.read(4096)
        return self._parse_mp3_header_bytes(data)

    def _parse_mp3_header_bytes(self, data: bytes) -> AudioInfo:
        """Parst MP3-Frame-Header aus Bytes."""
        # ID3v2 Tag überspringen
        offset = 0
        if data[:3] == b"ID3":
            size_bytes = data[6:10]
            size = (
                (size_bytes[0] & 0x7f) << 21 |
                (size_bytes[1] & 0x7f) << 14 |
                (size_bytes[2] & 0x7f) << 7 |
                (size_bytes[3] & 0x7f)
            )
            offset = 10 + size

        # MP3 Frame Header finden (0xFF 0xFB, 0xFF 0xFA, 0xFF 0xF3, 0xFF 0xF2)
        while offset < len(data) - 4:
            if data[offset] == 0xFF and (data[offset + 1] & 0xE0) == 0xE0:
                break
            offset += 1

        if offset >= len(data) - 4:
            return AudioInfo(format=AudioFormat.MP3)

        # Header parsen
        header = struct.unpack(">I", data[offset:offset+4])[0]

        # Version (Bits 19-20)
        version_bits = (header >> 19) & 0x03
        version_map = {0: 2.5, 2: 2, 3: 1}
        version = version_map.get(version_bits, 1)

        # Layer (Bits 17-18)
        layer_bits = (header >> 17) & 0x03
        layer = 4 - layer_bits if layer_bits > 0 else 0

        # Bitrate (Bits 12-15)
        bitrate_index = (header >> 12) & 0x0F
        bitrate_table = {
            # V1 L3
            (1, 3): [0, 32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320],
            # V2 L3
            (2, 3): [0, 8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160],
        }
        bitrate_list = bitrate_table.get((int(version), layer), bitrate_table[(1, 3)])
        bitrate = bitrate_list[bitrate_index] * 1000 if bitrate_index < len(bitrate_list) else 128000

        # Sample Rate (Bits 10-11)
        sample_rate_index = (header >> 10) & 0x03
        sample_rate_table = {
            1: [44100, 48000, 32000],
            2: [22050, 24000, 16000],
            2.5: [11025, 12000, 8000],
        }
        sample_rate_list = sample_rate_table.get(version, sample_rate_table[1])
        sample_rate = sample_rate_list[sample_rate_index] if sample_rate_index < len(sample_rate_list) else 44100

        # Channel Mode (Bits 6-7)
        channel_mode = (header >> 6) & 0x03
        channels = 1 if channel_mode == 3 else 2

        return AudioInfo(
            format=AudioFormat.MP3,
            sample_rate=sample_rate,
            channels=channels,
            bit_rate=bitrate,
            bit_depth=16,
        )

    async def decode(
        self,
        source: Path | bytes,
        target_sample_rate: int | None = None,
        target_channels: int | None = None,
    ) -> DecodedAudio:
        """Dekodiert MP3 vollständig zu PCM."""
        target_sr = target_sample_rate or self.OUTPUT_SAMPLE_RATE
        target_ch = target_channels or self.OUTPUT_CHANNELS

        if self._use_pydub:
            return await self._decode_pydub(source, target_sr, target_ch)
        elif self._use_ffmpeg:
            return await self._decode_ffmpeg(source, target_sr, target_ch)
        else:
            raise DecoderError("Kein MP3-Decoder verfügbar (pydub oder ffmpeg erforderlich)")

    async def _decode_pydub(
        self,
        source: Path | bytes,
        target_sr: int,
        target_ch: int,
    ) -> DecodedAudio:
        """Dekodiert mit pydub."""
        try:
            from pydub import AudioSegment

            def decode():
                if isinstance(source, bytes):
                    audio = AudioSegment.from_mp3(io.BytesIO(source))
                else:
                    audio = AudioSegment.from_mp3(str(source))

                # Resampling falls nötig
                if audio.frame_rate != target_sr:
                    audio = audio.set_frame_rate(target_sr)

                # Channel-Anpassung
                if audio.channels != target_ch:
                    if target_ch == 1:
                        audio = audio.set_channels(1)
                    else:
                        audio = audio.set_channels(2)

                return DecodedAudio(
                    pcm_data=audio.raw_data,
                    sample_rate=audio.frame_rate,
                    channels=audio.channels,
                    bit_depth=audio.sample_width * 8,
                )

            return await asyncio.to_thread(decode)

        except Exception as e:
            raise DecoderError(f"pydub Dekodierung fehlgeschlagen: {e}")

    async def _decode_ffmpeg(
        self,
        source: Path | bytes,
        target_sr: int,
        target_ch: int,
    ) -> DecodedAudio:
        """Dekodiert mit ffmpeg Subprocess."""
        try:
            if isinstance(source, bytes):
                # Pipe-Input
                cmd = [
                    "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
                    "-i", "pipe:0",
                    "-f", "s16le",
                    "-acodec", "pcm_s16le",
                    "-ar", str(target_sr),
                    "-ac", str(target_ch),
                    "pipe:1"
                ]
                process = await asyncio.create_subprocess_exec(
                    *cmd,
                    stdin=asyncio.subprocess.PIPE,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                stdout, stderr = await process.communicate(input=source)
            else:
                # Datei-Input
                cmd = [
                    "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
                    "-i", str(source),
                    "-f", "s16le",
                    "-acodec", "pcm_s16le",
                    "-ar", str(target_sr),
                    "-ac", str(target_ch),
                    "pipe:1"
                ]
                process = await asyncio.create_subprocess_exec(
                    *cmd,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                stdout, stderr = await process.communicate()

            if process.returncode != 0:
                raise DecoderError(f"ffmpeg Fehler: {stderr.decode()}")

            return DecodedAudio(
                pcm_data=stdout,
                sample_rate=target_sr,
                channels=target_ch,
                bit_depth=16,
            )

        except Exception as e:
            if isinstance(e, DecoderError):
                raise
            raise DecoderError(f"ffmpeg Dekodierung fehlgeschlagen: {e}")

    async def decode_stream(
        self,
        source: Path | bytes,
        chunk_duration_ms: int = 100,
    ) -> AsyncIterator[DecodedAudio]:
        """Dekodiert MP3 als Stream."""
        # Berechne Chunk-Größe
        bytes_per_sample = 2  # 16-bit
        samples_per_chunk = int(self.OUTPUT_SAMPLE_RATE * chunk_duration_ms / 1000)
        chunk_size = samples_per_chunk * bytes_per_sample * self.OUTPUT_CHANNELS

        if self._use_pydub:
            # Bei pydub alles dekodieren und in Chunks liefern
            decoded = await self.decode(source)
            data = decoded.pcm_data

            for i in range(0, len(data), chunk_size):
                yield DecodedAudio(
                    pcm_data=data[i:i + chunk_size],
                    sample_rate=decoded.sample_rate,
                    channels=decoded.channels,
                    bit_depth=decoded.bit_depth,
                )

        elif self._use_ffmpeg:
            # ffmpeg streaming
            async for chunk in self._stream_ffmpeg(source, chunk_size):
                yield chunk

    async def _stream_ffmpeg(
        self,
        source: Path | bytes,
        chunk_size: int,
    ) -> AsyncIterator[DecodedAudio]:
        """Streaming-Dekodierung mit ffmpeg."""
        try:
            if isinstance(source, bytes):
                cmd = [
                    "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
                    "-i", "pipe:0",
                    "-f", "s16le",
                    "-acodec", "pcm_s16le",
                    "-ar", str(self.OUTPUT_SAMPLE_RATE),
                    "-ac", str(self.OUTPUT_CHANNELS),
                    "pipe:1"
                ]
                process = await asyncio.create_subprocess_exec(
                    *cmd,
                    stdin=asyncio.subprocess.PIPE,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                # Input schreiben
                if process.stdin:
                    process.stdin.write(source)
                    await process.stdin.drain()
                    process.stdin.close()
            else:
                cmd = [
                    "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
                    "-i", str(source),
                    "-f", "s16le",
                    "-acodec", "pcm_s16le",
                    "-ar", str(self.OUTPUT_SAMPLE_RATE),
                    "-ac", str(self.OUTPUT_CHANNELS),
                    "pipe:1"
                ]
                process = await asyncio.create_subprocess_exec(
                    *cmd,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )

            # Output streamen
            while True:
                chunk = await process.stdout.read(chunk_size)
                if not chunk:
                    break

                yield DecodedAudio(
                    pcm_data=chunk,
                    sample_rate=self.OUTPUT_SAMPLE_RATE,
                    channels=self.OUTPUT_CHANNELS,
                    bit_depth=16,
                )

            await process.wait()

        except Exception as e:
            raise DecoderError(f"ffmpeg Streaming fehlgeschlagen: {e}")