import uuid
import ffmpeg
from .progress_socket import progress_socket
import tempfile

CODEC_OPTS = {
    "mp3": {
        "codec": "libmp3lame",
        "f": "mp3",
        "qscale:a": "0",
    },
    "flac": {
        "codec": "flac",
        "f": "flac",
        "compression_level": "12",
    },
    "ogg": {
        "codec": "libopus",
        "f": "ogg",
        "b:a": "160K",
        "application": "audio",
    },
}

# x realtime, measured on deployment hardware (g4dn.xlarge)
CODEC_SPEEDS = {"mp3": 26.0, "ogg": 29.7, "flac": 9.4}


def download_audio(
    s3,
    s3_bucket,
    s3_key,
    out_dir,
    progress_handler=None,
    out_basename=None,
    sample_rate=48000,
):
    if out_basename is None:
        random_id = uuid.uuid4()
        out_path = f"{out_dir}/{random_id}.wav"
    else:
        out_path = f"{out_dir}/{out_basename}.wav"

    if progress_handler is None:
        progress_handler = lambda k, v: None

    with progress_socket(progress_handler) as socket_filename:
        process = (
            ffmpeg.input("pipe:")
            .output(
                out_path,
                acodec="pcm_f32le",
                af="aresample=resampler=soxr:precision=28",
                ar=sample_rate,
            )
            .global_args("-progress", f"unix://{socket_filename}")
            .run_async(pipe_stdin=True, quiet=True)
        )
        s3.download_fileobj(s3_bucket, s3_key, process.stdin)
        process.stdin.close()
        process.wait()

    return out_path


def upload_audio(
    s3,
    s3_bucket,
    s3_key,
    in_path,
    codec,
    progress_handler=None,
):
    if progress_handler is None:
        progress_handler = lambda k, v: None

    with tempfile.TemporaryDirectory() as temp_dir, progress_socket(
        progress_handler
    ) as socket_filename:
        (
            ffmpeg.input(in_path)
            .output(f"{temp_dir}/output", **CODEC_OPTS[codec])
            .global_args("-progress", f"unix://{socket_filename}")
            .run(quiet=True)
        )
        s3.upload_file(f"{temp_dir}/output", s3_bucket, s3_key)


def resample_audio(
    in_path,
    out_path,
    sample_rate,
    progress_handler=None,
):
    if progress_handler is None:
        progress_handler = lambda k, v: None

    with progress_socket(progress_handler) as socket_filename:
        (
            ffmpeg.input(in_path)
            .output(
                out_path,
                af="aresample=resampler=soxr:precision=28",
                acodec="pcm_f32le",
                ar=sample_rate,
            )
            .global_args("-progress", f"unix://{socket_filename}")
            .run(quiet=True)
        )