import audioop
import base64
from collections.abc import Iterable
from contextlib import contextmanager, redirect_stderr, redirect_stdout
from collections import deque
import copy
import funcy
from joblib import Parallel, delayed
from joblib.externals.loky import get_reusable_executor
import logging
import math
import multiprocessing
import os
import pathlib
import random
import shlex
import subprocess
import sys
import tempfile
import time
import uuid
import io
import asyncio
import wave
from typing import List, Literal, Optional

import ffmpeg
import numpy as np
import sox
from tinytag import TinyTag
from tinytag.tinytag import TinyTagException

# from ..utils.constants import N_CHANNELS
# from ..utils.display import suppress_logging
# from ..utils.s3 import open_from_s3, _download_s3_file
# from ..utils.io import pipe_through_process_asyncio
# from ..web.harvest import get_file_ext, get_filename, get_url, time_limit


@contextmanager
def suppress_logging(highest_level=logging.CRITICAL):
    previous_level = logging.root.manager.disable
    logging.disable(highest_level)
    try:
        yield
    finally:
        logging.disable(previous_level)


N_CHANNELS = 1


SUPPORTED_AUDIO_FORMATS = [
    {
        "sample_rate": 48_000,
        "byte_width": 2,
        "n_channels": 2,
    },
    {
        "sample_rate": 48_000,
        "byte_width": 2,
        "n_channels": 1,
    },
    {
        "sample_rate": 44_100,
        "byte_width": 2,
        "n_channels": 2,
    },
    {
        "sample_rate": 44_100,
        "byte_width": 2,
        "n_channels": 1,
    },
    {
        "sample_rate": 32_000,
        "byte_width": 2,
        "n_channels": 1,
    },
    {
        "sample_rate": 24_000,
        "byte_width": 2,
        "n_channels": 1,
    },
    {
        "sample_rate": 22_050,
        "byte_width": 2,
        "n_channels": 1,
    },
    {
        "sample_rate": 16_000,
        "byte_width": 2,
        "n_channels": 1,
    },
    {
        "sample_rate": 8_000,
        "byte_width": 2,
        "n_channels": 1,
    },
    # TODO: bit depth 8 is tricky cause wav riff header doesn't support it
    # {
    #     "sample_rate": 8_000,
    #     "byte_width": 1,
    #     "n_channels": 1,
    # },
]

FORMAT_STR_LOOKUP = {
    # 1: "u8",
    2: "s16le",
}

DEFAULT_LD_LIBRARY_PATH = "/usr/local/lib"
DEFAULT_OPUSENC_PATHS = (
    "/usr/local/bin/opusenc",
    "/usr/bin/opusenc",
    "/bin/opusenc",
)
DEFAULT_OPUSDEC_PATHS = (
    "/usr/local/bin/opusdec",
    "/usr/bin/opusdec",
    "/bin/opusdec",
)

MP3_FFMPEG_ARGS = ["-aq", "2", "-f", "mp3"]
OPUS_FFMPEG_ARGS = [
    "-c:a",
    "libopus",
    "-b:a",
    "128k",
    "-vbr",
    "1",
    "-flush_packets",
    "1",
    "-f",
    "ogg",
]

logger = logging.getLogger(__name__)


def _find_best_format(sample_rate, byte_width=None, n_channels=None):
    if byte_width is None:
        byte_width = 2
    if n_channels is None:
        n_channels = 1
    compatible_audio_formats = [d for d in SUPPORTED_AUDIO_FORMATS if d["n_channels"] == n_channels]
    if sample_rate not in set([d["sample_rate"] for d in compatible_audio_formats]):
        logger.warning(f"sample rate {sample_rate} not supported, will use infer best.")
    viable_formats = [
        d
        for d in compatible_audio_formats
        if d["sample_rate"] >= sample_rate and d["byte_width"] >= byte_width
    ]
    if len(viable_formats) == 0:
        # take highest
        viable_formats = [compatible_audio_formats[0]]
    best_format = viable_formats[-1]
    if byte_width > best_format["byte_width"]:
        logger.warning(f"down-conversion from byte width {byte_width} to {best_format['byte_width']}.")
    return best_format


def has_valid_header(filepath):
    """Fast read from header."""
    tag = TinyTag.get(filepath)
    if tag.samplerate is None:
        return False
    return True


def _read_wav_header(stream):
    header_start = stream.read(12)
    if len(header_start) < 12 or header_start[:4] != b"RIFF" or header_start[8:12] != b"WAVE":
        raise ValueError("Invalid WAV header")

    while True:
        chunk_header = stream.read(8)
        if len(chunk_header) < 8:
            raise ValueError("Incomplete WAV file: Couldn't find format chunk")

        chunk_id = chunk_header[:4]
        chunk_size = int.from_bytes(chunk_header[4:8], byteorder="little")

        if chunk_id == b"fmt ":
            fmt_data = stream.read(chunk_size)
            if len(fmt_data) < 16:
                raise ValueError("Incomplete format chunk in WAV header")

            n_channels = int.from_bytes(fmt_data[2:4], byteorder="little")
            sample_rate = int.from_bytes(fmt_data[4:8], byteorder="little")
            bits_per_sample = int.from_bytes(fmt_data[14:16], byteorder="little")
            byte_width = bits_per_sample // 8

            if chunk_size > 16:
                stream.read(chunk_size - 16)
            break
        else:
            stream.read(chunk_size)

    while True:
        chunk_header = stream.read(8)
        if len(chunk_header) < 8:
            raise ValueError("Incomplete WAV file: Couldn't find data chunk")

        chunk_id = chunk_header[:4]
        chunk_size = int.from_bytes(chunk_header[4:8], byteorder="little")

        if chunk_id == b"data":
            return sample_rate, byte_width, n_channels
        else:
            stream.read(chunk_size)


def _probe_format_ffmpeg(filepath, codec_type="audio"):
    out = ffmpeg.probe(filepath)
    streams = [
        s for s in sorted(out["streams"], key=lambda x: x["index"]) if s["codec_type"] == codec_type
    ]
    if len(streams) >= 2:
        logger.warning("more than one streams found, taking first one")
    elif len(streams) == 0:
        raise ValueError("no streams found")
    main_stream = streams[0]
    if "duration" not in main_stream and "format" in out and "duration" in out["format"]:
        main_stream["duration"] = out["format"]["duration"]
    if "bit_rate" not in main_stream and "format" in out and "bit_rate" in out["format"]:
        main_stream["bit_rate"] = out["format"]["bit_rate"]
    return main_stream


def _load_wav_audio_bytes(filepath, sample_rate=None, byte_width=None, n_channels=None):
    # if get_file_ext(filepath) != "wav":
    #     raise NotImplementedError("wave module can only read wav files")
    f = wave.open(filepath, "r")
    file_info = f.getparams()
    in_sample_rate = int(file_info.framerate)
    in_byte_width = int(file_info.sampwidth)
    in_n_channels = int(file_info.nchannels)
    if sample_rate is not None:
        best_format_info = _find_best_format(sample_rate, byte_width=byte_width, n_channels=n_channels)
    else:
        best_format_info = _find_best_format(
            in_sample_rate, byte_width=in_byte_width, n_channels=n_channels
        )
    out_sample_rate = best_format_info["sample_rate"]
    out_byte_width = best_format_info["byte_width"]
    out_n_channels = best_format_info["n_channels"]
    if (
        out_sample_rate == in_sample_rate
        and out_byte_width == in_byte_width
        and out_n_channels == in_n_channels
    ):
        audio_bytes = f.readframes(f.getnframes())
    else:
        raise NotImplementedError("wave module can't do format conversions")
    f.close()
    return audio_bytes, (out_sample_rate, out_byte_width, out_n_channels)


def _load_sox_audio_bytes(filepath, sample_rate=None, byte_width=None, n_channels=None):
    # if get_file_ext(filepath) not in sox.core.VALID_FORMATS:
    #     raise NotImplementedError("unsupported filetype")
    if "sox WARN" in str(sox.file_info.stat(filepath)):
        # checking for warnings when processing with sox (this is not super fast)
        raise ValueError("warnings when analyzing via sox")
    if byte_width is None:
        with suppress_logging():
            in_bit_depth = sox.file_info.bitdepth(filepath)
        if in_bit_depth is not None:
            byte_width = int(in_bit_depth / 8)
    if sample_rate is not None:
        best_format_info = _find_best_format(sample_rate, byte_width=byte_width, n_channels=n_channels)
    else:
        in_sample_rate = int(sox.file_info.sample_rate(filepath))
        best_format_info = _find_best_format(
            in_sample_rate, byte_width=byte_width, n_channels=n_channels
        )
    out_sample_rate = best_format_info["sample_rate"]
    out_byte_width = best_format_info["byte_width"]
    out_bit_depth = int(out_byte_width * 8)
    out_n_channels = best_format_info["n_channels"]
    tfm = sox.Transformer()
    tfm.set_output_format(rate=out_sample_rate, bits=out_bit_depth, channels=out_n_channels)
    audio_arr = tfm.build_array(input_filepath=filepath)
    # fix if corrupt header (happens if eg recording device ran out of battery)
    if len(audio_arr) == 0:
        logger.warning("file seems corrupt, trying to fix.")
        tfm = sox.Transformer()
        tfm.set_input_format(ignore_length=True)
        tfm.set_output_format(rate=out_sample_rate, bits=out_bit_depth, channels=out_n_channels)
        audio_arr = tfm.build_array(input_filepath=filepath)
    if out_byte_width != audio_arr.dtype.itemsize:
        raise ValueError("inconsistent byte width during conversion")
    # if encoding was unsigned then we need to fix (sox bug always interprets as signed)
    encoding_type = sox.file_info.encoding(filepath)
    if "Unsigned Integer" in encoding_type and audio_arr.dtype == np.int8:
        audio_arr = audio_arr.view(np.uint8)
    elif "Unsigned Integer" in encoding_type:
        raise ValueError("unsigned fixup only tested for int8")
    return audio_arr.tobytes(), (out_sample_rate, out_byte_width, out_n_channels)


def _load_ffmpeg_audio_bytes(filepath, sample_rate=None, byte_width=None, n_channels=None):
    if byte_width is None:
        # TODO: this can often be 0
        byte_width = int(_probe_format_ffmpeg(filepath)["bits_per_sample"] / 8)
        if byte_width == 0:
            byte_width = None
    if sample_rate is not None:
        best_format_info = _find_best_format(sample_rate, byte_width=byte_width, n_channels=n_channels)
    else:
        in_sample_rate = int(_probe_format_ffmpeg(filepath)["sample_rate"])
        best_format_info = _find_best_format(
            in_sample_rate, byte_width=byte_width, n_channels=n_channels
        )
    out_sample_rate = best_format_info["sample_rate"]
    out_byte_width = best_format_info["byte_width"]
    out_n_channels = best_format_info["n_channels"]
    with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir:
        tmp_fp = os.path.join(temp_dir, "converted.wav")
        _convert_audio_ffmpeg(
            filepath,
            tmp_fp,
            sample_rate=out_sample_rate,
            byte_width=out_byte_width,
            n_channels=out_n_channels,
        )
        audio_bytes, _ = _load_wav_audio_bytes(
            tmp_fp,
            sample_rate=out_sample_rate,
            byte_width=out_byte_width,
            n_channels=out_n_channels,
        )
    return audio_bytes, (out_sample_rate, out_byte_width, out_n_channels)


def _load_audio_as_bytes(filepath, sample_rate=None, byte_width=None, n_channels=None):
    if not os.path.exists(filepath):
        raise FileNotFoundError(f"{filepath} does not exists")
    try:
        audio_bytes, (sample_rate, byte_width, n_channels) = _load_wav_audio_bytes(
            filepath, sample_rate=sample_rate, byte_width=byte_width, n_channels=n_channels
        )
    except (NotImplementedError, wave.Error):
        try:
            audio_bytes, (sample_rate, byte_width, n_channels) = _load_sox_audio_bytes(
                filepath, sample_rate=sample_rate, byte_width=byte_width, n_channels=n_channels
            )
        except (NotImplementedError, ValueError, sox.core.SoxiError):
            audio_bytes, (sample_rate, byte_width, n_channels) = _load_ffmpeg_audio_bytes(
                filepath, sample_rate=sample_rate, byte_width=byte_width, n_channels=n_channels
            )
    return audio_bytes, (sample_rate, byte_width, n_channels)


def _convert_audio_sox(
    in_filepath: str,
    out_filepath: str,
    sample_rate: Optional[int] = None,
    byte_width: Optional[int] = None,
    n_channels: Optional[int] = None,
    encoding: Optional[str] = None,
) -> None:
    """Convert audio using sox, such as mp3 -> wav."""
    tfm = sox.Transformer()
    kwargs = {}
    if sample_rate is not None:
        kwargs["rate"] = sample_rate
    if byte_width is not None:
        kwargs["bits"] = 8 * byte_width
    if n_channels is not None:
        kwargs["channels"] = n_channels
    if encoding is not None:
        kwargs["encoding"] = encoding
    tfm.set_output_format(**kwargs)
    tfm.build_file(in_filepath, out_filepath)


def _convert_audio_ffmpeg(
    in_filepath: str,
    out_filepath: str,
    offset_s: Optional[float] = None,
    duration_s: Optional[float] = None,
    sample_rate: Optional[int] = None,
    byte_width: Optional[int] = None,
    n_channels: Optional[int] = None,
    limit_threads: bool = False,
):
    in_kwargs = {}
    if limit_threads:
        in_kwargs["threads"] = 1
    if offset_s is not None:
        in_kwargs["ss"] = offset_s
    if duration_s is not None:
        in_kwargs["t"] = duration_s
    out_kwargs = {}
    if sample_rate is not None:
        out_kwargs["ar"] = sample_rate
    if byte_width is not None:
        format_str = FORMAT_STR_LOOKUP[byte_width]
        out_kwargs["acodec"] = f"pcm_{format_str}"
    if n_channels is not None:
        out_kwargs["ac"] = n_channels
    if limit_threads:
        out_kwargs["threads"] = 1
    global_args = ["-nostdin"]
    stream = ffmpeg.input(in_filepath, **in_kwargs)
    stream = stream.audio
    stream = ffmpeg.output(stream, out_filepath, **out_kwargs)
    stream = stream.global_args(*global_args)
    _ = ffmpeg.run(stream, overwrite_output=True, quiet=True)


def convert_audio_file(
    in_filepath: str,
    out_filepath: str,
    offset_s: Optional[float] = None,
    duration_s: Optional[float] = None,
    sample_rate: Optional[int] = None,
    byte_width: Optional[int] = None,
    n_channels: Optional[int] = None,
    limit_threads: bool = False,
    debug: bool = False,
):
    try:
        with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir:
            _convert_audio_ffmpeg(
                in_filepath,
                out_filepath,
                offset_s=offset_s,
                duration_s=duration_s,
                sample_rate=sample_rate,
                byte_width=byte_width,
                n_channels=n_channels,
                limit_threads=limit_threads,
            )
    except Exception as e:
        if debug:
            raise e
        return False
    return True


def convert_audio_files(
    in_filepaths: List[str],
    out_filepaths: List[str],
    n_cores: Optional[int] = None,
    force_threads: bool = False,
    offset_s: Optional[float] = None,
    duration_s: Optional[float] = None,
    sample_rate: Optional[int] = None,
    byte_width: Optional[int] = None,
    n_channels: Optional[int] = None,
    debug: bool = False,
):
    assert len(in_filepaths) == len(out_filepaths)
    p_convert_audio = funcy.partial(
        convert_audio_file,
        offset_s=offset_s,
        duration_s=duration_s,
        sample_rate=sample_rate,
        byte_width=byte_width,
        n_channels=n_channels,
        limit_threads=True,
        debug=debug,
    )
    if n_cores is None:
        n_cores = np.max([1, multiprocessing.cpu_count() - 1])
    n_cores = np.min([n_cores, len(in_filepaths)])
    if len(in_filepaths) == 1:
        confirmed_list = [p_convert_audio(in_filepaths[0], out_filepaths[0])]
    else:
        process_type = "threads" if force_threads else "processes"
        confirmed_list = Parallel(n_jobs=n_cores, prefer=process_type, batch_size=1)(
            delayed(p_convert_audio)(in_fp, out_fp) for in_fp, out_fp in zip(in_filepaths, out_filepaths)
        )
        get_reusable_executor().shutdown(wait=True)
    return confirmed_list


def _convert_audio_bytes(
    in_bytes: bytes,
    in_sample_rate: int,
    in_byte_width: int,
    in_n_channels: int,
    out_sample_rate: int,
    out_byte_width: int,
    out_n_channels: int,
) -> bytes:
    """Convert audio using sox."""
    if (
        in_sample_rate == out_sample_rate
        and in_byte_width == out_byte_width
        and in_n_channels == out_n_channels
    ):
        return in_bytes
    if in_byte_width == 1:
        array_dtype = np.int8
    elif in_byte_width == 2:
        array_dtype = np.int16
    else:
        raise NotImplementedError("only byte width 1 & 2 supported")
    in_array = np.frombuffer(in_bytes, dtype=array_dtype).reshape(-1, in_n_channels)  # T, ch
    tfm = sox.Transformer()
    tfm.set_output_format(rate=out_sample_rate, bits=out_byte_width * 8, channels=out_n_channels)
    out_array = tfm.build_array(input_array=in_array, sample_rate_in=in_sample_rate).T  # ch, T
    # convention is unsigned for int8 (even though sox doesn't like that)
    if out_array.dtype == np.int8:
        out_array = (out_array + 128).astype(np.uint8)
    return out_array.T.reshape(
        -1,
    ).tobytes()


def _write_bytes_as_wav(
    fp: str,
    audio_bytes: bytes,
    sample_rate: int,
    byte_width: int,
    n_channels: int,
) -> None:
    """Save wav file to disk with correct header information."""
    with wave.open(fp, "wb") as f:
        f.setnchannels(n_channels)
        f.setframerate(sample_rate)
        f.setsampwidth(byte_width)
        f.writeframes(audio_bytes)
        return f


def _get_audio_bytes_slice(
    audio_bytes, sample_rate, byte_width, n_channels=N_CHANNELS, from_s=None, to_s=None
):
    if to_s is not None:
        end_b = int(round(to_s * byte_width * sample_rate * n_channels))
        # make sure to respect byte_width
        if end_b % byte_width != 0:
            end_b -= end_b % byte_width
        audio_bytes = audio_bytes[:end_b]
    if from_s is not None:
        start_b = int(round(from_s * byte_width * sample_rate * n_channels))
        # make sure to respect byte_width
        if (len(audio_bytes) - start_b) % byte_width != 0:
            start_b += (len(audio_bytes) - start_b) % byte_width
        audio_bytes = audio_bytes[start_b:]
    return audio_bytes


def _collapse_to_numpy_array(arr):
    if isinstance(arr, np.ndarray):
        pass
    if "torch" in str(type(arr)):
        arr = arr.detach().cpu().numpy()
    # squash extra dimensions
    arr = arr.squeeze()
    return arr


def _convert_to_int_audio_array(arr, auto_compress=True, max_allowed_val=1.2):
    arr = _collapse_to_numpy_array(arr)
    # enforce int representation
    if "int" in str(arr.dtype):
        pass
    elif "float" in str(arr.dtype):
        # alert if signal too high
        max_val = np.abs(arr).max()
        if max_val > max_allowed_val:
            raise ValueError("signal overflow")
        if max_val > 1.0:
            logger.info("signal overflow")
            if auto_compress:
                logger.info("compressing signal")
                arr = arr / max_val
        # convert to int16 for playback
        arr = (
            (arr * np.iinfo(np.int16).max)
            .clip(np.iinfo(np.int16).min, np.iinfo(np.int16).max)
            .astype(np.int16)
        )
    else:
        raise NotImplementedError("unknown array format")
    return arr


##############
# Public API #
##############


def get_audio_properties(filepath, attempt_using_header=True):
    """Fast read from header."""
    filetype = filepath.split(".")[-1]
    if attempt_using_header:
        # note that header base info is much faster but potentially incorrect
        try:
            tag = TinyTag.get(filepath)
            filetype = filepath.split(".")[-1]
            if tag.bitrate is None:
                raise ValueError("header parsing failed")
            est_duration_s = tag.filesize * 8 / tag.bitrate / 1_000
            if tag.duration / est_duration_s >= 1.5:
                raise ValueError("header duration seems corrupt")
            audio_info = {
                "sample_rate": tag.samplerate,
                "byte_width": (
                    int(tag.bitrate * 1_000 / tag.samplerate / 8) if filetype == "wav" else None
                ),
                "bitrate": tag.bitrate * 1_000,
                "n_channels": tag.channels,
                "duration_s": tag.duration,
            }
            return audio_info
        except (ValueError, TinyTagException):
            pass
    # use ffmpeg
    ffmpeg_info = _probe_format_ffmpeg(filepath)
    sample_rate = int(ffmpeg_info["sample_rate"])
    bit_rate = int(ffmpeg_info["bit_rate"])
    audio_info = {
        "sample_rate": sample_rate,
        "byte_width": int(bit_rate / sample_rate / 8) if filetype == "wav" else None,
        "bitrate": bit_rate,
        "n_channels": int(ffmpeg_info["channels"]),
        "duration_s": float(ffmpeg_info["duration"]),
    }
    return audio_info


def get_duration_s(filepath, attempt_using_header=True):
    audio_info = get_audio_properties(filepath, attempt_using_header=attempt_using_header)
    return audio_info["duration_s"]


def convert_audio(
    in_filepath: str,
    out_filepath: str,
    sample_rate: Optional[int] = None,
    byte_width: Optional[int] = None,
    n_channels: Optional[int] = None,
):
    try:
        # with suppress_logging():
        _convert_audio_sox(
            in_filepath,
            out_filepath,
            sample_rate=sample_rate,
            byte_width=byte_width,
            n_channels=n_channels,
        )
    except:
        _convert_audio_ffmpeg(
            in_filepath,
            out_filepath,
            sample_rate=sample_rate,
            byte_width=byte_width,
            n_channels=n_channels,
        )


global AUDIO_STREAM_LINKS
AUDIO_STREAM_LINKS = []

AUDIO_PLAYER_HTML_PTN = """
<audio controls="controls" autobuffer="autobuffer" style="{width_style_str}">
  <source src="{audio_src}"/>
  Your browser does not support the audio element.
</audio>
"""


def _get_stream_links_dir():
    if "JUPYTER_NOTEBOOK_DIR" not in os.environ:
        raise ValueError("env var `JUPYTER_NOTEBOOK_DIR` not specified")
    jupyter_notebook_dir = pathlib.Path(os.environ["JUPYTER_NOTEBOOK_DIR"])
    stream_links_dir = os.path.join(jupyter_notebook_dir, "suno_stream_links")
    os.makedirs(stream_links_dir, exist_ok=True)
    # make this relative
    cwd = pathlib.Path(os.getcwd())
    relative_stream_links_dir = (
        "".join(["../"] * len(cwd.relative_to(jupyter_notebook_dir).parts)) + "suno_stream_links"
    )
    return relative_stream_links_dir


def clean_stream_links(force_global=False):
    global AUDIO_STREAM_LINKS
    if force_global:
        stream_links_dir = _get_stream_links_dir()
        for symlink_filename in os.listdir(stream_links_dir):
            os.remove(os.path.join(stream_links_dir, symlink_filename))
    else:
        for _, symlink_filepath in AUDIO_STREAM_LINKS:
            os.remove(symlink_filepath)
    AUDIO_STREAM_LINKS = []


def fade_out_audio(
    audio_filepath: str,
    fade_out_len: float,
    fade_out_shape: Optional[Literal["q", "l", "t"]],
    sample_rate: int,
    byte_width: int,
    n_channels: int,
):
    tfm = sox.Transformer()
    best_format_info = _find_best_format(sample_rate, byte_width=byte_width, n_channels=n_channels)
    out_sample_rate = best_format_info["sample_rate"]
    out_byte_width = best_format_info["byte_width"]
    out_bit_depth = int(out_byte_width * 8)
    out_n_channels = best_format_info["n_channels"]

    tfm.set_output_format(rate=out_sample_rate, bits=out_bit_depth, channels=out_n_channels)
    tfm.fade(fade_out_len=fade_out_len, fade_shape=fade_out_shape or "q")
    array_out = tfm.build_array(input_filepath=audio_filepath, sample_rate_in=sample_rate)
    return Audio.from_array(array_out.T.reshape(out_n_channels, -1), sample_rate=out_sample_rate)


def change_audio_speed(
    audio_filepath: str,
    speed_factor: float,
    tempo_only: bool,
    sample_rate: int,
    byte_width: int,
    n_channels: int,
):
    tfm = sox.Transformer()
    best_format_info = _find_best_format(sample_rate, byte_width=byte_width, n_channels=n_channels)
    out_sample_rate = best_format_info["sample_rate"]
    out_byte_width = best_format_info["byte_width"]
    out_bit_depth = int(out_byte_width * 8)
    out_n_channels = best_format_info["n_channels"]

    tfm.set_output_format(rate=out_sample_rate, bits=out_bit_depth, channels=out_n_channels)
    if tempo_only:
        tfm.tempo(speed_factor, audio_type="m")
    else:
        tfm.speed(speed_factor)

    array_out = tfm.build_array(input_filepath=audio_filepath, sample_rate_in=sample_rate)
    return Audio.from_array(array_out.T.reshape(out_n_channels, -1), sample_rate=out_sample_rate)


def trim_audio_silence(
    audio: "Audio",
    trim_start: bool = True,
    trim_end: bool = True,
    silence_threshold: float = 0.1,
    output_original_start_end_time: bool = False,
):
    """Given an input audio, trim the start and end of it if they are at low volumns."""
    step_duration = 1
    audio_loudness = deque(
        [
            audio.get_segment(from_s=i, to_s=i + step_duration).loudness
            for i in range(0, math.floor(audio.duration_s), step_duration)
        ]
    )
    start_time = 0
    end_time = math.ceil(audio.duration_s)
    loundess_cut = np.quantile(audio_loudness, silence_threshold)
    while trim_end and audio_loudness and audio_loudness[-1] < loundess_cut:
        audio_loudness.pop()
        end_time -= step_duration
    while trim_start and audio_loudness and audio_loudness[0] < loundess_cut and start_time < end_time:
        audio_loudness.popleft()
        start_time += step_duration

    output_audio = audio.get_segment(from_s=start_time, to_s=end_time)
    if output_original_start_end_time:
        return output_audio, start_time, end_time
    return output_audio


class Audio:
    def __init__(self, audio_bytes, sample_rate, byte_width, n_channels, metadata=None):
        supported_format_tuples = set(
            [(d["sample_rate"], d["byte_width"], d["n_channels"]) for d in SUPPORTED_AUDIO_FORMATS]
        )
        if (sample_rate, byte_width, n_channels) not in supported_format_tuples:
            raise NotImplementedError("not supported audio format")
        self.bytes = audio_bytes
        self.sample_rate = int(sample_rate)
        self.byte_width = int(byte_width)
        self.n_channels = int(n_channels)
        self.metadata = copy.deepcopy(metadata) if metadata is not None else {}

    def __repr__(self):
        return (
            f"Audio({self.duration_s:.1f} secs, {self.sample_rate} Hz, "
            f"{self.bit_depth} bits, {self.n_channels} ch)"
        )

    @property
    def bit_depth(self):
        return self.byte_width * 8

    @property
    def duration_s(self):
        return len(self.bytes) / self.sample_rate / self.byte_width / self.n_channels

    @property
    def duration_ms(self):
        return int(round(self.duration_s * 1_000))

    @property
    def samples(self):
        return len(self.bytes) // self.byte_width // self.n_channels

    @property
    def array(self):
        assert self.byte_width == 2
        arr = np.frombuffer(self.bytes, dtype=np.int16).reshape(-1, self.n_channels).T
        if self.n_channels == 1:
            arr = arr.mean(0).astype(np.int16)
        return arr

    @property
    def array_float(self):
        arr = np.frombuffer(self.bytes, dtype=np.int16).reshape(-1, self.n_channels).T
        arr = arr.astype(np.float32) / np.iinfo(arr.dtype).max
        if self.n_channels == 1:
            arr = arr.mean(0).astype(np.float32)
        return arr

    @property
    def loudness(self):
        import pyloudnorm as pyln

        m = pyln.Meter(self.sample_rate)  # create BS.1770 meter
        arr = self.array_float
        if len(arr.shape) == 1:
            arr = arr[None]
        lufs_db = m.integrated_loudness(arr.T)
        return lufs_db

    @classmethod
    def empty(cls, sample_rate, byte_width, n_channels):
        return cls(b"", sample_rate, byte_width, n_channels)

    def is_zero(self):
        return np.all(self.array == 0)

    @classmethod
    def from_file(cls, filepath, sample_rate=None, byte_width=None, n_channels=1, metadata=None):
        audio_bytes, (out_sample_rate, out_byte_width, out_n_channels) = _load_audio_as_bytes(
            filepath,
            sample_rate=sample_rate,
            byte_width=byte_width,
            n_channels=n_channels,
        )
        return cls(audio_bytes, out_sample_rate, out_byte_width, out_n_channels, metadata=metadata)

    @classmethod
    def from_array(cls, audio_arr, sample_rate, metadata=None):
        if audio_arr.dtype != np.int16:
            raise NotImplementedError("only int16 (byte width 2) supported")
        if audio_arr.ndim == 1:
            audio_arr = audio_arr[None]
        elif audio_arr.ndim == 2:
            pass
        else:
            raise NotImplementedError(f"only single 1 & 2 channel supported, got {audio_arr.shape}")
        n_channels = audio_arr.shape[0]

        return cls(
            audio_arr.T.reshape(
                -1,
            ).tobytes(),
            sample_rate,
            2,
            n_channels,
            metadata=metadata,
        )

    @classmethod
    def from_array_float(
        cls, audio_arr, sample_rate, metadata=None, auto_compress=True, max_allowed_val=1.2
    ):
        audio_arr = _convert_to_int_audio_array(
            audio_arr, auto_compress=auto_compress, max_allowed_val=max_allowed_val
        )
        return cls.from_array(audio_arr, sample_rate, metadata=metadata)

    @classmethod
    def from_silence(cls, duration_s, sample_rate, byte_width=2, n_channels=1, metadata=None):
        n_frames = int(round(duration_s * sample_rate))
        audio = cls.from_array(np.zeros(n_frames, dtype=np.int16), sample_rate, metadata=metadata)
        return audio.convert(sample_rate, byte_width, n_channels)

    @classmethod
    def from_beep(
        cls,
        sample_rate,
        byte_width=2,
        n_channels=1,
        freq_khz=0.1,
        duration_s=0.25,
        amp=0.5,
        metadata=None,
    ):
        arr = np.sin(np.arange(0, 1, 1 / sample_rate / duration_s) * freq_khz * 1_000 * 2 * np.pi) * amp
        audio = cls.from_array_float(arr, sample_rate=sample_rate, metadata=metadata)
        return audio.convert(sample_rate, byte_width, n_channels)

    @classmethod
    def concatenate(cls, audios):
        if len(audios) == 0:
            raise ValueError("audio list empty")
        out_sample_rate = audios[0].sample_rate
        out_byte_width = audios[0].byte_width
        out_n_channels = audios[0].n_channels
        if not (
            all(audio.sample_rate == out_sample_rate for audio in audios)
            and all(audio.byte_width == out_byte_width for audio in audios)
            and all(audio.n_channels == out_n_channels for audio in audios)
        ):
            raise ValueError("audio elements have different formats")
        out_bytes = b"".join([audio.bytes for audio in audios])
        return cls(out_bytes, out_sample_rate, out_byte_width, out_n_channels)

    @staticmethod
    def get_details(filepath, attempt_using_header=False):
        return get_audio_properties(filepath, attempt_using_header=attempt_using_header)

    @staticmethod
    def get_duration_s(filepath, attempt_using_header=False):
        return get_audio_properties(filepath, attempt_using_header=attempt_using_header)["duration_s"]

    # @staticmethod
    # def play_audio(audio_data):
    #     play_audio(audio_data)

    @staticmethod
    def convert_file(in_filepath, out_filepath, sample_rate=None, byte_width=None, n_channels=None):
        convert_audio(
            in_filepath,
            out_filepath,
            sample_rate=sample_rate,
            byte_width=byte_width,
            n_channels=n_channels,
        )

    def convert(self, sample_rate, byte_width, n_channels, metadata=None):
        new_audio_bytes = _convert_audio_bytes(
            in_bytes=self.bytes,
            in_sample_rate=self.sample_rate,
            in_byte_width=self.byte_width,
            in_n_channels=self.n_channels,
            out_sample_rate=sample_rate,
            out_byte_width=byte_width,
            out_n_channels=n_channels,
        )
        return Audio(new_audio_bytes, sample_rate, byte_width, n_channels, metadata=metadata)

    def mono(self):
        return self.convert(self.sample_rate, self.byte_width, 1)

    def stereo(self):
        return self.convert(self.sample_rate, self.byte_width, 2)

    def reduce_stereo_width(self, target_width=0.5):
        # 1 is retain full, 0 is mono-like
        assert self.n_channels == 2
        left_arr = self.array_float[0]
        right_arr = self.array_float[1]
        new_left_arr = left_arr / 2 * (1 + target_width) + right_arr / 2 * (1 - 1 * target_width)
        new_right_arr = right_arr / 2 * (1 + target_width) + left_arr / 2 * (1 - 1 * target_width)
        new_arr = np.stack([new_left_arr, new_right_arr])
        return Audio.from_array_float(new_arr, sample_rate=self.sample_rate, metadata=self.metadata)

    def normalize_volume(self, target_db=-16):
        gain_factor = np.log(10) / 20
        gain = target_db - self.loudness
        gain = np.exp(gain * gain_factor)
        norm_arr = self.array_float * gain
        if np.abs(norm_arr).max() > 1:
            norm_arr = norm_arr / np.abs(norm_arr).max()
        if self.n_channels == 1:
            norm_arr = norm_arr[0]
        return Audio.from_array_float(norm_arr, self.sample_rate, metadata=self.metadata)

    def apply_gain(self, gain_db: float):
        gain_factor = np.log(10) / 20
        gain = gain_db * gain_factor
        norm_arr = self.array_float * np.exp(gain)
        actual_gain_db = gain_db
        max_val = np.abs(norm_arr).max()
        if max_val > 1:
            # Calculate the actual gain applied after clipping protection
            clip_factor = 1.0 / max_val
            norm_arr = norm_arr * clip_factor
            actual_gain_db = 20 * np.log10(clip_factor * np.exp(gain))
        return (
            Audio.from_array_float(norm_arr, self.sample_rate, metadata=self.metadata),
            actual_gain_db,
        )

    def normalize_loudness(self):
        raise NotImplementedError("needs to be done via ffmpeg")

    def append(self, audio):
        compat_audio = audio.convert(self.sample_rate, self.byte_width, self.n_channels)
        total_bytes = self.bytes + compat_audio.bytes
        return Audio(total_bytes, self.sample_rate, self.byte_width, self.n_channels)

    def prepend(self, audio):
        compat_audio = audio.convert(self.sample_rate, self.byte_width, self.n_channels)
        total_bytes = compat_audio.bytes + self.bytes
        return Audio(total_bytes, self.sample_rate, self.byte_width, self.n_channels)

    async def pipe_to_ffmpeg_async(self, ffmpeg_args: list[str], receive_stdout: bool = True):
        wav_data_io = io.BytesIO()
        with wave.open(wav_data_io, "wb") as wf:
            wf.setnchannels(self.n_channels)
            wf.setframerate(self.sample_rate)
            wf.setsampwidth(self.byte_width)
            wf.writeframes(self.bytes)

        reader = asyncio.StreamReader()

        class BytesIOProtocol(asyncio.Protocol):
            def __init__(self, reader: asyncio.StreamReader):
                self.reader = reader

            def connection_made(self, transport):
                pass

            def data_received(self, data):
                self.reader.feed_data(data)

            def connection_lost(self, exc):
                self.reader.feed_eof()

        protocol = BytesIOProtocol(reader)

        data = wav_data_io.getvalue()
        protocol.data_received(data)
        protocol.connection_lost(None)

        raise NotImplementedError("Not implemented")

        # return await pipe_through_process_asyncio(
        #     reader,
        #     [
        #         "ffmpeg",
        #         "-hide_banner",
        #         "-loglevel",
        #         "error",
        #         "-y",
        #         "-f",
        #         "wav",
        #         "-t",
        #         str(self.duration_s),
        #         "-i",
        #         "pipe:0",
        #     ]
        #     + ffmpeg_args,
        #     receive_stdout=receive_stdout,
        # )

    async def pipe_opus(self):
        return await self.pipe_to_ffmpeg_async(
            [
                *OPUS_FFMPEG_ARGS,
                "-flush_packets",
                "1",
                "pipe:1",
            ]
        )

    async def write_opus_async(self, filepath):
        return await self.pipe_to_ffmpeg_async(
            [
                *OPUS_FFMPEG_ARGS,
                filepath,
            ],
            receive_stdout=False,
        )

    def get_slice(self, from_s=None, to_s=None):
        new_audio_bytes = _get_audio_bytes_slice(
            self.bytes,
            self.sample_rate,
            self.byte_width,
            n_channels=self.n_channels,
            from_s=from_s,
            to_s=to_s,
        )
        return Audio(new_audio_bytes, self.sample_rate, self.byte_width, self.n_channels)

    def mix(self, other_audio, weight_self=1, weight_other=1):
        """Mix this audio with another audio file.

        Args:
            other_audio (Audio): Another Audio object to mix with
            weight_self (float): Weight for the current audio (default: 0.5)
            weight_other (float): Weight for the other audio (default: 0.5)

        Returns:
            Audio: A new Audio object containing the mixed audio
        """
        # Convert other audio to match current format
        other_audio = other_audio.convert(self.sample_rate, self.byte_width, self.n_channels)

        # Convert both to float arrays for mixing
        arr1 = self.array_float
        arr2 = other_audio.array_float

        # Make sure arrays are the same length by padding the shorter one with zeros
        max_length = max(len(arr1), len(arr2))
        if len(arr1) < max_length:
            pad_length = max_length - len(arr1)
            arr1 = np.pad(arr1, ((0, pad_length), (0, 0)) if arr1.ndim > 1 else (0, pad_length))
        elif len(arr2) < max_length:
            pad_length = max_length - len(arr2)
            arr2 = np.pad(arr2, ((0, pad_length), (0, 0)) if arr2.ndim > 1 else (0, pad_length))

        # Mix the arrays using the specified weights
        mixed_array = weight_self * arr1 + weight_other * arr2

        # Create new Audio object from mixed array
        return Audio.from_array_float(
            mixed_array, self.sample_rate, metadata=self.metadata, max_allowed_val=10
        )

    @staticmethod
    def sum(audios: list["Audio"], auto_compress: bool = False):
        """Sum multiple audio files together.

        Args:
            audios (list[Audio]): List of Audio objects to sum together

        Returns:
            Audio: A new Audio object containing the summed audio

        Raises:
            ValueError: If the audio files have different sample rates or channel counts
        """
        if not audios:
            raise ValueError("Cannot sum an empty list of audio files")

        # Check that all audio files have the same sample rate and channel count
        sample_rate = audios[0].sample_rate
        n_channels = audios[0].n_channels

        for i, audio in enumerate(audios[1:], 1):
            if audio.sample_rate != sample_rate:
                raise ValueError(
                    f"Audio at index {i} has sample rate {audio.sample_rate}, expected {sample_rate}"
                )
            if audio.n_channels != n_channels:
                raise ValueError(
                    f"Audio at index {i} has {audio.n_channels} channels, expected {n_channels}"
                )

        # Get all arrays and find the maximum length
        arrays = [audio.array_float for audio in audios]
        max_length = max(arr.shape[-1] for arr in arrays)

        # Pad arrays to the same length if needed
        padded_arrays = []
        for arr in arrays:
            if arr.shape[-1] < max_length:
                pad_length = max_length - arr.shape[-1]
                padded_arr = np.pad(arr, ((0, 0), (0, pad_length)) if arr.ndim > 1 else (0, pad_length))
                padded_arrays.append(padded_arr)
            else:
                padded_arrays.append(arr)

        return Audio.from_array_float(
            np.sum(padded_arrays, axis=0),
            audios[0].sample_rate,
            metadata=audios[0].metadata,
            max_allowed_val=10,  # dont check for overflow
            auto_compress=auto_compress,
        )

    def __add__(self, other):
        return Audio.sum([self, other])

    def pad_to_length(self, length_s: float):
        """Pad the audio to a specific length.

        Args:
            length_s (float): The desired length of the audio in seconds
        """
        if length_s <= self.duration_s:
            return self
        else:
            return self.concatenate(
                [
                    self,
                    Audio.from_silence(length_s, self.sample_rate, self.byte_width, self.n_channels),
                ]
            ).get_slice(from_s=0, to_s=length_s)

    def resample(self, target_sample_rate: int):
        """Resample the audio to a different sample rate.

        Args:
            target_sample_rate (int): The desired output sample rate in Hz

        Returns:
            Audio: A new Audio object with the resampled audio
        """
        if target_sample_rate == self.sample_rate:
            return self

        return self.convert(
            sample_rate=target_sample_rate,
            byte_width=self.byte_width,
            n_channels=self.n_channels,
            metadata=self.metadata,
        )

    # make alias methods
    get_properties = get_details
    get_segment = get_slice
    # to_wav = write_wav
    # to_wav_async = write_wav_async
    # to_mp3 = write_mp3
    # to_hq_mp3 = write_hq_mp3
    # to_m4a = write_m4a
    # to_opus = write_opus
    # to_webm_opus = write_webm_opus
