import apache_beam as beam
import apache_beam.io.fileio as beam_io
from pathlib import Path
from apache_beam.testing.test_pipeline import TestPipeline as BeamTestPipeline
from apache_beam.testing.util import assert_that, equal_to

from klay_beam.transforms import LoadWithLibrosa


def test_LoadWithLibrosa():
    this_file = Path(__file__)
    data_dir = this_file.parent / "test_data" / "audio"
    wav_pattern = str(data_dir / "*.wav")
    ogg_pattern = str(data_dir / "*.ogg")
    mp3_pattern = str(data_dir / "*.mp3")
    patterns = [wav_pattern, ogg_pattern, mp3_pattern]

    # This list of filenames is a hard-coded copy of the list of files in the
    # tests/test_data/audio directory. These files were generated by the
    # bin/create_test_audio_files.py script. Each file is 0.5 seconds at
    # 44.1kHz.
    data = {
        "test_mono.mp3": (22050,),  # Annoyingly, this is what librosa returns for mono
        "test_quad.ogg": (4, 22050),
        "test_quad.wav": (4, 22050),
        "test_stereo.mp3": (2, 22050),
        "test_stereo.ogg": (2, 22050),
        "test_stereo_16bit.wav": (2, 22050),
        "test_stereo_24bit.wav": (2, 22050),
        "test_stereo_32bit.wav": (2, 22050),
        "test_stereo_64bit.wav": (2, 22050),
    }

    def verify_data(element):
        filename, audio, sr = element
        basename = Path(filename).name
        expected_shape = data[basename]

        assert sr == 44100
        assert audio.shape == expected_shape

    expected_filenames = [str(Path(data_dir / datum)) for datum in data.keys()]

    with BeamTestPipeline() as p:
        # Apply the custom transform
        readable_files = p | beam.Create(patterns) | beam.io.fileio.MatchAll()

        # Extract file names from the metadata for assertion
        file_names = readable_files | "Extract File Names" >> beam.Map(
            lambda metadata: metadata.path
        )

        # Check if the readable_files match the expected file names
        assert_that(file_names, equal_to(expected_filenames))

        # Load the audio files using librosa
        (
            readable_files
            | beam_io.ReadMatches()
            | beam.ParDo(LoadWithLibrosa(target_sr=None, mono=False))
            | beam.Map(verify_data)
        )


def test_fail_LoadWithLibrosa():
    this_file = Path(__file__)
    data_dir = this_file.parent / "test_data" / "audio" / "corrupt"
    wav_pattern = str(data_dir / "*.wav")
    ogg_pattern = str(data_dir / "*.ogg")
    mp3_pattern = str(data_dir / "*.mp3")
    patterns = [wav_pattern, ogg_pattern, mp3_pattern]

    # This list of filenames is a hard-coded copy of the list of files in the
    # tests/test_data/audio directory. These files were generated by the
    # bin/create_test_audio_files.py script. Each file is 0.5 seconds at
    # 44.1kHz.
    data = {"test_corrupt.mp3": ()}

    expected_filenames = [str(Path(data_dir / datum)) for datum in data.keys()]

    with BeamTestPipeline() as p:
        # Apply the custom transform
        readable_files = p | beam.Create(patterns) | beam.io.fileio.MatchAll()

        # Extract file names from the metadata for assertion
        file_names = readable_files | "Extract File Names" >> beam.Map(
            lambda metadata: metadata.path
        )

        # Check if the readable_files match the expected file names
        assert_that(file_names, equal_to(expected_filenames))

        # Load the audio files using librosa
        audio, failed = (
            readable_files
            | beam_io.ReadMatches()
            | beam.ParDo(LoadWithLibrosa(target_sr=None, mono=False)).with_outputs(
                "failed", main="audio"
            )
        )

        def verify_failed(element):
            assert element == 1

        (
            failed
            | "Count Failed" >> beam.combiners.Count.Globally()
            | "Check Failed" >> beam.Map(verify_failed)
        )

        def verify_audio_failed(element):
            assert element == 0

        (
            audio
            | "Count Audio" >> beam.combiners.Count.Globally()
            | "Check Audio Failed" >> beam.Map(verify_audio_failed)
        )
