"""Audio utilities for combining and tagging MP3s.

These utilities are useful for working with multiple MP3s generated by a TTS model from chunks of a larger body of text.
For example, combining the MP3s into a single MP3 and adding ID3 tags, including cover art.
"""

import io

from mutagen.mp3 import MP3, EasyMP3
from mutagen.id3 import ID3, APIC
from mutagen.id3._util import ID3NoHeaderError


def _is_valid_mp3(data: bytes) -> bool:
    """Quick check to see if this data at least starts like an MP3."""
    return data.startswith(b"ID3") or data.startswith(b"\xff")


def combine_mp3_buffers(
    mp3_buffers: list[io.BytesIO],
    tags: dict[str, str] | None = None,
    album_art: io.BytesIO | None = None,
    album_art_mime: str = "image/png",
) -> io.BytesIO:
    """Combines multiple in-memory MP3 buffers into a single MP3 file buffer and applies tags and album art.

    Use this function to combine TTS-generated MP3s that were created from chunked text input.

    Args:
        mp3_buffers: List of in-memory MP3 buffers to combine.
        tags: Dictionary of tags to apply to the output MP3.
        album_art: In-memory buffer for the album art image.
        album_art_mime: MIME type for the album art (typically 'image/png' or 'image/jpeg').

    Raises:
        ValueError: If no buffers are provided, if buffers are invalid MP3,
                    or if their codec parameters (bitrate, sample rate) differ.

    Examples:

     Combine two in-memory MP3 files and tag the result:

    ```python
    buffer1 = io.BytesIO(open("file1.mp3", "rb").read())
    buffer2 = io.BytesIO(open("file2.mp3", "rb").read())

    tags = {
        "title": "Combined Audio",
        "artist": "AI Composer",
        "album": "AI Album",
        "genre": "Books & Spoken",
    }

    # Load cover image PNG from disk
    with open("cover.png", "rb") as img_file:
        album_art_bytes = io.BytesIO(img_file.read())

    combined_mp3 = await combine_mp3_buffers(
        [buffer1, buffer2],
        tags=tags,
        album_art=album_art_bytes,
        album_art_mime="image/png",
    )

    # Write MP3 to file
    with open("output.mp3", "wb") as out_file:
        combined_mp3.seek(0)
        out_file.write(combined_mp3.read())
    ```
    """
    if not mp3_buffers:
        raise ValueError("No MP3 buffers provided for combination.")

    reference_info = None
    for index, mp3_buffer in enumerate(mp3_buffers):
        # Reset the buffer position just in case
        mp3_buffer.seek(0)
        data = mp3_buffer.read()

        # Sanity check to ensure it at least looks like MP3 data
        if not _is_valid_mp3(data):
            raise ValueError("Invalid MP3 buffer: does not start with ID3 or MPEG frame header.")

        # Attempt to parse the MP3 info - if this fails, we can't combine it with anything
        try:
            temp_audio = MP3(io.BytesIO(data))
        except Exception as exc:
            raise ValueError(f"Error parsing MP3 buffer at index {index}: {exc}") from exc

        current_info = (temp_audio.info.bitrate, temp_audio.info.sample_rate)
        if reference_info is None:
            reference_info = current_info
        else:
            if current_info != reference_info:
                raise ValueError(
                    f"Inconsistent MP3 parameters detected at index {index}. "
                    f"Expected {reference_info}, got {current_info}."
                )

    # Once validated, do the actual combination
    output_buffer = io.BytesIO()
    for index, mp3_buffer in enumerate(mp3_buffers):
        mp3_buffer.seek(0)
        data = mp3_buffer.read()

        if index == 0:
            # Write the entire first MP3, including headers
            output_buffer.write(data)
        else:
            # Skip the ID3 header for subsequent MP3s
            try:
                # Check if the file has ID3 tags and determine the audio frame offset
                tags = ID3(io.BytesIO(data))
                audio_offset = tags.size if tags else 0
            except ID3NoHeaderError:
                # No ID3 tags present, start from the beginning
                audio_offset = 0

            # Write the buffer starting from the offset
            output_buffer.write(data[audio_offset:])

    # Convert the combined bytes to an MP3
    output_buffer.seek(0)
    audio: EasyMP3 = EasyMP3(output_buffer)

    # Tag it with what we have so far
    if tags:
        # Overwrite or create tags
        audio.add_tags()
        for tag_key, tag_value in tags.items():
            audio[tag_key] = tag_value

    # Save the tags (does not save the file to disk)
    audio.save(output_buffer)
    output_buffer.seek(0)

    if album_art:
        album_art.seek(0)
        audio = ID3(output_buffer)
        audio.add(
            APIC(
                encoding=3,  # UTF-8
                mime=album_art_mime,
                type=3,  # Front cover
                desc="Cover",
                data=album_art.read(),
            )
        )
        # Save the tags (again) so the cover image is added
        audio.save(output_buffer)
        output_buffer.seek(0)

    # Return the in-memory MP3 as a BytesIO so caller can
    # do what they wish with it (like save it to a file)
    return output_buffer
