import boto3
import os
from suno_utils.audio import Audio
from suno_utils.utils.s3 import _parse_s3_filepath, get_oracle_client_config

try:
    from suno_utils.utils.opusfile import OpusFile, S3ByteRangeReader, ffi, lib
except Exception as e:
    print(f"Could not load libopusfile: {e}")

oracle_client = None

os.umask(0o003)  # set umask to 0o003 to allow group write for created directories


def _is_corrupt_opus(local_filepath):
    error_ = ffi.new("int *")
    c_path = ffi.new("char[]", local_filepath.encode("utf-8"))
    of_ = lib.op_open_file(c_path, error_)
    is_corrupt = of_ == ffi.NULL
    lib.op_free(of_)
    del of_, error_, c_path
    return is_corrupt


def get_sample_oracle_file_segment(
    local_filepath,
    s3_filepath="s3://webdataset/bundles/v0/discogs_subset_50k/audio/DEJm1igan2Y.opus",
    cloud_type="oracle",  # "oracle" or "aws"
    max_duration_s=60 * 60,  # 1 hour limit by default
    start_s=0,
) -> Audio:
    if not os.path.exists(local_filepath):
        global oracle_client
        assert cloud_type in ["oracle", "aws"]
        assert s3_filepath.startswith("s3://")
        assert s3_filepath.endswith(".opus")
        assert start_s >= 0

        source_bucket, source_key = _parse_s3_filepath(s3_filepath)
        if oracle_client is None:
            client_config = {}
            if cloud_type == "oracle":
                client_config = get_oracle_client_config()
            oracle_client = boto3.client(
                "s3",
                endpoint_url=client_config.get("endpoint_url"),
                aws_access_key_id=client_config.get("aws_access_key_id"),
                aws_secret_access_key=client_config.get("aws_secret_access_key"),
                region_name=client_config.get("region_name"),
            )

        # check if local_filepath is a file
        if local_filepath is not None:
            exists = os.path.exists(local_filepath)
            is_corrupt = False if not exists else _is_corrupt_opus(local_filepath)
            if not exists or is_corrupt:
                # copy from s3 to local_filepath
                if is_corrupt:
                    print(f"Repairing corrupt file, downloading {s3_filepath} to {local_filepath}")
                oracle_client.download_file(source_bucket, source_key, local_filepath)
                os.chmod(local_filepath, 0o774)

        assert os.path.exists(local_filepath)

    if local_filepath.endswith(".opus"):
        audio = Audio.from_array_float(
            OpusFile(path=local_filepath).read(
                buf_size=int(48000 * max_duration_s),
                float_samples=True,
                from_position=int(48000 * start_s),
            ),
            sample_rate=48000,
            max_allowed_val=12,
        )
    else:
        audio = Audio.from_file(local_filepath, sample_rate=48000, n_channels=2)
    return audio