import av

from towhee import register
from towhee.operator import PyOperator, OperatorFlag
from towhee.types.audio_frame import AudioFrame


@register(output_schema=['audio_frame'],
          flag=OperatorFlag.REUSEABLE)
class AudioDecoderFFmpeg(PyOperator):
    """
    """

    def __init__(self) -> None:
        super().__init__()

    def __call__(self, audio_path: str):
        in_container = av.open(audio_path)
        stream = in_container.streams.get(audio=0)[0]
        for frame in in_container.decode(stream):
            timestamp = int(frame.time * 1000)
            sample_rate = frame.sample_rate
            ndarray = frame.to_ndarray()
            yield AudioFrame(ndarray, sample_rate, timestamp)