import os
import re
import subprocess
import time


def recursive_ls_dir(directory):
    paths = [os.path.join(root, file) for root, dirs, files in os.walk(directory) for file in files]
    paths += [os.path.join(root, dir) for root, dirs, files in os.walk(directory) for dir in dirs]
    print(paths)


def download_models_to_dir(files: list[str], dir_path: str):
    """Download models to a directory using AWS CLI."""
    for f in files:
        full_path = os.path.abspath(os.path.join(dir_path, f))
        if not os.path.exists(full_path):
            parent_dir = os.path.dirname(full_path)
            os.makedirs(parent_dir, exist_ok=True)

            print("Downloading", full_path)
            cmd = [
                "aws",
                "s3",
                "cp",
                f"s3://suno-data/{f}",
                full_path,
            ]
            subprocess.run(cmd)
        else:
            print("Skipping", f)


def strip_square_brackets(text):
    return re.sub(r"\[.*?\]", "", text).strip()


def count_syllables_word(word):
    if not word:
        return 0

    word = word.lower()
    vowels = "aeiouy"
    count = 0

    if word[0] in vowels:
        count += 1

    for i in range(1, len(word)):
        if word[i] in vowels and word[i - 1] not in vowels:
            count += 1

    if word.endswith("e"):
        count -= 1

    if word.endswith("le"):
        count += 1

    if count == 0:
        count += 1

    return count


def count_syllables(text):
    word_arr = text.split()
    count = 0
    for word in word_arr:
        count += count_syllables_word(word)
    return count


def wrap_text(text, line_length=25):
    words = text.split()
    lines = []
    line = ""

    for word in words:
        if len(line) + len(word) + 1 <= line_length:
            if line:
                line += " " + word
            else:
                line += word
        else:
            lines.append(line)
            line = word

    lines.append(line)
    return "\n".join(lines)


def make_transcript_timestamps(audio, text):
    from suno_utils.tasks.asr.whisper import transcribe

    start_time = time.time()

    # Get whisper transcribe (retry it if it fails)
    history_text = transcribe(audio)
    print("whisper response ", str(history_text))
    if not history_text.get("text"):
        print("try whisper again")
        history_text = transcribe(audio)
        print("whisper response ", str(history_text))
    print(f"Made whisper transcription in {time.time() - start_time:.2f} seconds.")
    # Format whisper output
    segments = history_text["segments"]
    transcription = []
    for s in segments:
        for w in s["words"]:
            transcription.append(w)

    print(transcription)

    text = strip_square_brackets(text)
    text = text.split("\n")

    # 1. Store syllables up to each word in the original and transcribed text
    real_text = []
    for i in range(len(text)):
        word = text[i]
        val = {}
        val["word"] = word
        val["syllables"] = count_syllables(word) + (real_text[i - 1]["syllables"] if i > 0 else 0)
        real_text.append(val)

    transcribed_text = []
    for i in range(len(transcription)):
        word = transcription[i]["word"]
        val = {}
        val["word"] = word
        val["syllables"] = count_syllables(word) + (transcribed_text[i - 1]["syllables"] if i > 0 else 0)
        val["start"] = transcription[i]["start"]
        val["end"] = transcription[i]["end"]
        transcribed_text.append(val)

    # 2. Iterate through /n in original text.
    # Find first word in the transcribed text that has >= syllables and use that as start time for new line
    transcribe_word_count = 0
    for t in real_text:
        print(t)
        last_transcribed_syllable = transcribed_text[-1].get("syllables") if transcribed_text else 0
        print("last_transcribed_syllable is", last_transcribed_syllable)
        if t["syllables"] > last_transcribed_syllable:
            t["did_not_reach"] = True
            print("  did not reach")
        while transcribe_word_count < len(transcribed_text):
            transcribed_word = transcribed_text[transcribe_word_count]
            if transcribed_word["syllables"] >= t["syllables"]:
                t["end_time"] = transcribed_word["end"]
                next_transcribed_word = (
                    transcribed_text[transcribe_word_count + 1]
                    if transcribe_word_count + 1 < len(transcribed_text)
                    else transcribed_word
                )
                t["start_time"] = (
                    next_transcribed_word["start"]
                    if transcribe_word_count + 1 < len(transcribed_text)
                    else transcribed_word["end"]
                )
                print(
                    "start of next word",
                    next_transcribed_word["word"],
                    next_transcribed_word["start"],
                    next_transcribed_word["end"],
                    str(transcribe_word_count + 1 < len(transcribed_text)),
                )
                print(
                    "MATCH: ",
                    transcribe_word_count,
                    len(transcribed_text),
                    transcribed_word,
                )
                break
            transcribe_word_count += 1

    # Generate timestamps array (which is used by video). 't' stores timestamp, 'l' stores lyrics
    timestamps = []
    for i in range(len(real_text)):
        if real_text[i] and real_text[i].get("did_not_reach") is None or False:
            val = {}
            val["l"] = wrap_text(real_text[i]["word"])
            end_time_str = real_text[i - 1].get("end_time") if i > 0 else 0
            start_time_str = real_text[i - 1].get("start_time") if i > 0 else 0
            val["t"] = (
                round((float(end_time_str) + float(start_time_str)) / 2, 2)
                if end_time_str is not None and start_time_str is not None
                else timestamps[-1]["t"]
            )
            if i > 0 and val["t"] == timestamps[-1]["t"]:
                timestamps[-1]["l"] += "\n" + val["l"]
                timestamps[-1]["l"] = wrap_text(timestamps[-1]["l"])
            else:
                val["l"] = wrap_text(val["l"])
                timestamps.append(val)
    print(timestamps)
    return timestamps


def retry_decorator(max_retries=3, wait_seconds=1, silent_return: bool = False):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for attempt in range(1, max_retries + 1):
                try:
                    result = func(*args, **kwargs)
                    return result  # If successful, return the result
                except Exception as e:
                    print(
                        f"Attempt {attempt} failed: {str(e)}. \
                        Input args are {[str(x) for x in args]}. \
                        Input kwargs are {[str(k) + ':' + str(v) for k, v in kwargs.items()]}"
                    )
                    if attempt < max_retries:
                        print(f"Retrying in {wait_seconds} seconds...")
                        time.sleep(wait_seconds)
                    else:
                        print(f"All {max_retries} attempts failed. Giving up.")
                        if silent_return:
                            return None
                        else:
                            raise FileNotFoundError(f"Timed out retry because of {str(e)}")

        return wrapper

    return decorator


def ffmpeg_stream_encode(gain_adjust: float = 0):
    return subprocess.Popen(
        [
            "ffmpeg",
            "-f",
            "f32le",
            "-acodec",
            "pcm_f32le",
            "-ar",
            "48000",
            "-ac",
            "2",
            "-i",
            "pipe:",
            "-af",
            f"volume={str(float(gain_adjust))}dB,alimiter=limit=0.95:attack=5:release=50:level=0",
            "-aq",
            "2",
            "-flush_packets",
            "1",
            "-f",
            "mp3",
            "pipe:",
        ],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        bufsize=0,
    )


def ffmpeg_stream_encode_opus_webm(gain_adjust: float = 0):
    return subprocess.Popen(
        [
            "ffmpeg",
            "-f",
            "f32le",
            "-acodec",
            "pcm_f32le",
            "-ar",
            "48000",
            "-ac",
            "2",
            "-i",
            "pipe:",
            "-af",
            f"volume={str(float(gain_adjust))}dB,alimiter=limit=0.95:attack=5:release=50:level=0",
            "-c:a",
            "libopus",
            "-b:a",
            "128k",
            "-vbr",
            "1",
            "-flush_packets",
            "1",
            "-f",
            "webm",
            "pipe:",
        ],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        bufsize=0,
    )


def print_gpu_memory_usage(class_name: str = None):
    import torch

    t = torch.cuda.get_device_properties(0).total_memory
    r = torch.cuda.memory_reserved(0)
    a = torch.cuda.memory_allocated(0)
    print(
        f"{class_name} GPU memory usage: allocated {a / 1e9:.2f} GB, reserved {r / 1e9:.2f} GB, total {t / 1e9:.2f} GB"
    )