Trixy7/trixy_core/music/formats/wav.py

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

Dekodiert WAV-Dateien zu PCM.
"""

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

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


class WAVDecoder(AudioDecoder):
    """
    WAV-Decoder.

    WAV ist bereits PCM, daher ist die Dekodierung einfach.
    Unterstützt Resampling und Channel-Konvertierung.
    """

    FORMAT = AudioFormat.WAV
    EXTENSIONS = {"wav", "wave"}
    MIME_TYPES = {"audio/wav", "audio/x-wav", "audio/wave"}

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

    def _check_dependencies(self) -> bool:
        """WAV wird von Python nativ unterstützt."""
        return True

    async def get_info(self, source: Path | bytes) -> AudioInfo:
        """Liest WAV-Informationen."""
        try:
            if isinstance(source, bytes):
                wav_file = wave.open(io.BytesIO(source), "rb")
            else:
                if not source.exists():
                    raise DecoderError(f"Datei nicht gefunden: {source}")
                wav_file = wave.open(str(source), "rb")

            try:
                frames = wav_file.getnframes()
                sample_rate = wav_file.getframerate()
                channels = wav_file.getnchannels()
                sample_width = wav_file.getsampwidth()

                duration_ms = int((frames / sample_rate) * 1000)
                bit_rate = sample_rate * channels * sample_width * 8

                return AudioInfo(
                    format=AudioFormat.WAV,
                    duration_ms=duration_ms,
                    sample_rate=sample_rate,
                    channels=channels,
                    bit_depth=sample_width * 8,
                    bit_rate=bit_rate,
                    file_size=source.stat().st_size if isinstance(source, Path) else len(source),
                )
            finally:
                wav_file.close()

        except wave.Error as e:
            raise DecoderError(f"Ungültige WAV-Datei: {e}")
        except Exception as e:
            raise DecoderError(f"Kann WAV-Info nicht lesen: {e}")

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

        def do_decode():
            if isinstance(source, bytes):
                wav_file = wave.open(io.BytesIO(source), "rb")
            else:
                wav_file = wave.open(str(source), "rb")

            try:
                frames = wav_file.getnframes()
                sample_rate = wav_file.getframerate()
                channels = wav_file.getnchannels()
                sample_width = wav_file.getsampwidth()

                # Rohdaten lesen
                raw_data = wav_file.readframes(frames)

                # Konvertierung wenn nötig
                pcm_data = self._convert_pcm(
                    raw_data,
                    sample_rate=sample_rate,
                    channels=channels,
                    sample_width=sample_width,
                    target_sr=target_sr,
                    target_ch=target_ch,
                )

                return DecodedAudio(
                    pcm_data=pcm_data,
                    sample_rate=target_sr,
                    channels=target_ch,
                    bit_depth=16,
                )
            finally:
                wav_file.close()

        try:
            return await asyncio.to_thread(do_decode)
        except Exception as e:
            raise DecoderError(f"WAV Dekodierung fehlgeschlagen: {e}")

    def _convert_pcm(
        self,
        data: bytes,
        sample_rate: int,
        channels: int,
        sample_width: int,
        target_sr: int,
        target_ch: int,
    ) -> bytes:
        """
        Konvertiert PCM-Daten.

        Args:
            data: Rohe PCM-Daten
            sample_rate: Ursprüngliche Sample-Rate
            channels: Ursprüngliche Kanalzahl
            sample_width: Bytes pro Sample (1, 2, 3, 4)
            target_sr: Ziel-Sample-Rate
            target_ch: Ziel-Kanalzahl

        Returns:
            Konvertierte PCM-Daten (16-bit)
        """
        # Schritt 1: Zu 16-bit konvertieren falls nötig
        if sample_width == 1:
            # 8-bit unsigned → 16-bit signed
            samples = struct.unpack(f"{len(data)}B", data)
            samples = [(s - 128) << 8 for s in samples]
            data = struct.pack(f"{len(samples)}h", *samples)
            sample_width = 2
        elif sample_width == 3:
            # 24-bit → 16-bit (einfach obere 16 bit nehmen)
            samples = []
            for i in range(0, len(data), 3):
                # Little-endian 24-bit
                val = data[i] | (data[i+1] << 8) | (data[i+2] << 16)
                if val & 0x800000:
                    val -= 0x1000000
                samples.append(val >> 8)
            data = struct.pack(f"{len(samples)}h", *samples)
            sample_width = 2
        elif sample_width == 4:
            # 32-bit → 16-bit
            samples = struct.unpack(f"{len(data)//4}i", data)
            samples = [s >> 16 for s in samples]
            data = struct.pack(f"{len(samples)}h", *samples)
            sample_width = 2

        # Schritt 2: Channel-Konvertierung
        if channels != target_ch:
            samples = struct.unpack(f"{len(data)//2}h", data)

            if channels == 2 and target_ch == 1:
                # Stereo → Mono (Mittelwert)
                mono = []
                for i in range(0, len(samples), 2):
                    mono.append((samples[i] + samples[i+1]) // 2)
                samples = mono
                channels = 1

            elif channels == 1 and target_ch == 2:
                # Mono → Stereo (duplizieren)
                stereo = []
                for s in samples:
                    stereo.extend([s, s])
                samples = stereo
                channels = 2

            data = struct.pack(f"{len(samples)}h", *samples)

        # Schritt 3: Resampling falls nötig
        if sample_rate != target_sr:
            data = self._resample(data, channels, sample_rate, target_sr)

        return data

    def _resample(
        self,
        data: bytes,
        channels: int,
        src_rate: int,
        dst_rate: int,
    ) -> bytes:
        """
        Einfaches lineares Resampling.

        Für Produktionsqualität sollte scipy.signal.resample verwendet werden.
        """
        samples = struct.unpack(f"{len(data)//2}h", data)
        frame_count = len(samples) // channels

        # Ratio
        ratio = dst_rate / src_rate
        new_frame_count = int(frame_count * ratio)

        resampled = []
        for i in range(new_frame_count):
            # Position im Original
            src_pos = i / ratio
            src_idx = int(src_pos)
            frac = src_pos - src_idx

            for ch in range(channels):
                idx1 = src_idx * channels + ch
                idx2 = min(idx1 + channels, len(samples) - 1)

                if idx1 < len(samples):
                    # Lineare Interpolation
                    val1 = samples[idx1]
                    val2 = samples[idx2] if idx2 < len(samples) else val1
                    val = int(val1 * (1 - frac) + val2 * frac)
                    resampled.append(max(-32768, min(32767, val)))

        return struct.pack(f"{len(resampled)}h", *resampled)

    async def decode_stream(
        self,
        source: Path | bytes,
        chunk_duration_ms: int = 100,
    ) -> AsyncIterator[DecodedAudio]:
        """Dekodiert WAV 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

        try:
            if isinstance(source, bytes):
                wav_file = wave.open(io.BytesIO(source), "rb")
            else:
                wav_file = wave.open(str(source), "rb")

            try:
                source_sr = wav_file.getframerate()
                source_ch = wav_file.getnchannels()
                source_sw = wav_file.getsampwidth()

                # Frames pro Chunk (in Quell-Rate)
                source_chunk_frames = int(samples_per_chunk * source_sr / self.OUTPUT_SAMPLE_RATE)

                while True:
                    raw_data = wav_file.readframes(source_chunk_frames)
                    if not raw_data:
                        break

                    # Konvertieren
                    pcm_data = self._convert_pcm(
                        raw_data,
                        sample_rate=source_sr,
                        channels=source_ch,
                        sample_width=source_sw,
                        target_sr=self.OUTPUT_SAMPLE_RATE,
                        target_ch=self.OUTPUT_CHANNELS,
                    )

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

            finally:
                wav_file.close()

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