# Code that produces suno videos
import math
import os
import shutil
import subprocess
import tempfile
import threading
import time
from pathlib import Path
import bisect
import numpy as np

from PIL import Image, ImageDraw, ImageFilter, ImageFont, ImageOps, ImageFile

from suno_utils.audio.conversion import Audio
from suno_utils.worker.settings import s3_client
from suno_utils.worker.utils import retry_decorator

ImageFile.LOAD_TRUNCATED_IMAGES = True

retry_s3_download = retry_decorator(3, wait_seconds=10)(s3_client.download_file)

ASSETS_PATH = Path(str((Path(__file__).parent / "assets")))


def contains_chinese(text):
    return any("\u4e00" <= char <= "\u9fff" for char in text)


def contains_hebrew(text):
    return any("\u0590" <= char <= "\u05ff" for char in text)


def contains_arabic(text):
    return any("\u0600" <= char <= "\u06ff" for char in text)


def contains_japanese(text):
    for char in text:
        if (
            "\u3000" <= char <= "\u303f"
            or "\u3040" <= char <= "\u309f"  # Japanese-style punctuation
            or "\u30a0" <= char <= "\u30ff"  # Hiragana
            or "\u4e00" <= char <= "\u9faf"  # Katakana
            or "\u3400" <= char <= "\u4dbf"  # Kanji  # Rare Kanji
        ):
            return True
    return False


def contains_korean(text):
    # Hangul Syllables
    return any("\uac00" <= char <= "\ud7af" for char in text)


def get_font(text, is_title=False, is_bold=False):
    if text is None:
        text = ""
    is_english = text.isascii()

    is_cjk = not is_english and (
        contains_chinese(text) or contains_japanese(text) or contains_korean(text)
    )

    is_hebrew = (not is_english) and contains_hebrew(text)
    is_arabic = (not is_english) and contains_arabic(text)
    if is_hebrew:
        return str((ASSETS_PATH / "NotoSansHebrew-Medium.ttf").resolve())
    elif is_cjk:
        return str((ASSETS_PATH / "NotoSansCJK-Medium.ttc").resolve())
    elif is_arabic:
        return str((ASSETS_PATH / "NotoSansCJK-Medium.ttc").resolve())
    else:  # English
        if is_title:
            return str((ASSETS_PATH / "PPEditorialNew-Regular.woff").resolve())
        if is_bold:
            return str((ASSETS_PATH / "PPNeueMontreal-SemiBold.woff").resolve())
        return str((ASSETS_PATH / "PPNeueMontreal-Regular.woff").resolve())
        # return str((ASSETS_PATH / "Roobert-SemiBold.ttf").resolve())


# Creates a PIL image from a given string + font
def create_text_image(
    text, font, text_color=(255, 255, 255, 255), bg_color=(0, 0, 0, 0), max_width=None
):
    if text is None:
        text = ""
    bbox = font.getbbox(text)
    width = bbox[2] - bbox[0]

    if max_width is not None and width > max_width:
        chars_to_truncate = int(round((width - max_width / width) * len(text)))
        text = text[: -1 * chars_to_truncate]
        bbox = font.getbbox(text)
        width = bbox[2] - bbox[0]
    height = bbox[3] - bbox[1]
    size = (width, height)  # Calculate text size based on the bounding box
    image = Image.new("RGBA", size, bg_color)

    draw = ImageDraw.Draw(image)
    draw.text(
        (-bbox[0], -bbox[1]),
        text,
        font=font,
        fill=text_color,
    )
    return image


def add_gradient_to_image(image, start_opacity, end_opacity):
    mask = np.array(image)
    gradient = np.linspace(start_opacity, end_opacity, image.height)
    gradient = gradient[:, np.newaxis, np.newaxis]
    gradient = np.repeat(gradient, image.width, axis=1)
    gradient = np.repeat(gradient[:, :, np.newaxis], 4, axis=2)
    gradient[:, :, :3] = 1  # Set RGB channels to 1
    gradient = gradient.astype(np.float32)  # Ensure correct data type
    gradient = gradient.reshape(image.height, image.width, 4)
    mask = mask * gradient
    return Image.fromarray(mask.astype("uint8"), "RGBA")


def generate_video(
    audio: Audio,
    image_path: Path | str,
    aligned_lyrics: list[dict[str, str | float]],
    out_path: Path | str,
    song_title: str,
    song_user_handle: str | None = None,
    tags: str | None = None,
    fps=5,  # based on performance, if the text is rolling, video is less noticable
    resolution_factor: int = 1,
    video_mode: str | None = None,
) -> str:
    print("Generating video")
    tags = tags or ""
    handle = f"by @{song_user_handle or 'unknown'}"

    with tempfile.TemporaryDirectory() as temp_dir:
        start_time_total = time.time()

        temp_dir_string = str(temp_dir)
        # Save audio to MP3 for putting into video
        audio.write_wav(os.path.join(temp_dir_string, "tmp.wav"))
        duration = audio.duration_s

        # Run frame generation
        start_time = time.time()

        # NOTE the video size is 412x732 defined here
        size = (412, 732)
        width, height = size
        upsample = int(resolution_factor)
        width = int(width * upsample)
        height = int(height * upsample)
        safe_margin_x = 88 * upsample

        image_size_x = 140 * upsample
        image_size_y = 190 * upsample

        image = Image.open(image_path)

        new_w = int(image.height * (image_size_x / image_size_y))
        new_left = (image.width - new_w) // 2
        image = image.crop((new_left, 0, new_left + new_w, image.height))
        image = image.resize((image_size_x, image_size_y))

        # Logo
        logo = Image.open((ASSETS_PATH / "Logo-6.png").resolve())
        logo = logo.resize((164 * upsample, 40 * upsample))  # Hard-coded size
        logo_2 = Image.open((ASSETS_PATH / "alz1.png").resolve())
        logo_2 = logo_2.resize((101 * upsample, 95 * upsample))  # Hard-coded size
        logo_3 = Image.open((ASSETS_PATH / "alz2.png").resolve())
        logo_3 = logo_3.resize((114 * upsample, 48 * upsample))  # Hard-coded size

        # Render song name
        name_font = ImageFont.truetype(get_font(song_title, is_title=True), 26 * upsample)
        small_handle_font = ImageFont.truetype(get_font(handle), 11 * upsample)
        handle_font = ImageFont.truetype(get_font(handle), 12 * upsample)
        timbaland_font = ImageFont.truetype(get_font(handle), 14 * upsample)

        song_name = create_text_image(song_title, name_font, max_width=width - safe_margin_x)
        song_handle = create_text_image(handle, handle_font)
        small_song_handle = create_text_image(handle, small_handle_font)

        logo_line_font = ImageFont.truetype(get_font("MADE WITH"), 15 * upsample)
        logo_line = create_text_image("MADE WITH", logo_line_font)
        logo_small = logo.resize((45 * upsample, 11 * upsample))  # Hard-coded size

        # Song genre
        # genre_font = ImageFont.truetype(get_font(tags), 18)
        # song_genre = create_text_image(f"{tags[0:30]}..." if len(tags) > 30 else tags, genre_font)

        timbaland_text_1 = create_text_image("TIMBALAND'S", timbaland_font)
        timbaland_text_2 = create_text_image("'LOVE AGAIN' REMIX", timbaland_font)

        # Song lyrics
        # The song lyrics are parsed on a per line basis
        # First, parse the song lyrics into lines
        lines, lyrics_line_timestamps = _process_lyrics(aligned_lyrics, duration)

        # Make an image for each line
        lyrics_images = []
        bold_lyrics_images = []
        lyrics_font = ImageFont.truetype(get_font(" ".join(lines)), 16 * upsample)
        bold_lyrics_font = ImageFont.truetype(get_font(" ".join(lines), is_bold=True), 16 * upsample)
        for i, line in enumerate(lines):
            line_image = create_text_image(
                line,
                lyrics_font,
                text_color=(255, 255, 255, 255),
                max_width=(width - (safe_margin_x * 2)),
            )
            lyrics_images.append(line_image)
            bold_line_image = create_text_image(
                line,
                bold_lyrics_font,
                max_width=(width - (safe_margin_x * 2)),
            )
            bold_lyrics_images.append(bold_line_image)

        timbaland_bg = (
            Image.open((ASSETS_PATH / "SUNO_VIDEOSHARE_TIM.jpg").resolve())
            .resize((width, height))
            .convert("RGBA")
        )

        def generate_video_chunk(frame_start, frame_end, image):
            # Define the command to start ffmpeg accepting frames from a pipe
            command = [
                "ffmpeg",
                "-loglevel",
                "error",
                "-y",
                "-r",
                str(fps),  # Frame rate
                "-f",
                "image2pipe",
                "-s",
                f"{width}x{height}",  # Resolution of the input images
                "-i",
                "-",  # Video input from the pipe
                "-c:v",
                "libx264",  # Video codec
                "-qp",
                "28",  # Quality/speed balance for NVENC
                "-pix_fmt",
                "yuv420p",
                f"{temp_dir_string}/out{frame_start}.ts",
            ]

            # Start the ffmpeg process
            process = subprocess.Popen(command, stdin=subprocess.PIPE)

            # Add rounded corners to cover
            rounded_mask = Image.new("L", (image.size[0], image.size[1]), 0)
            image_masked = ImageDraw.Draw(rounded_mask)
            image_masked.rounded_rectangle([(0, 0), (image_size_x, image_size_y)], radius=25, fill=255)

            # Apply the rounded mask to the original image
            rounded_image = ImageOps.fit(image, (image_size_x, image_size_y), centering=(0.5, 0.5))
            rounded_image.putalpha(rounded_mask)

            timbaland_img = timbaland_bg.copy()  # noqa: F821  ruff false positive but todo: delete at some point

            if video_mode != "timbaland":
                # Create blurry background
                blur_radius = 30
                blurred_background = image.copy()
                blurred_background = blurred_background.filter(
                    ImageFilter.GaussianBlur(radius=blur_radius)
                )
                blurred_background = blurred_background.resize(
                    (width, math.ceil(height * 2))
                )  # Make background very large

            current_line = 0
            lines_shown = 3 if video_mode == "timbaland" else 5
            image_y_pos = 240 if video_mode == "timbaland" else 190 * upsample

            for i in range(frame_start, frame_end):
                # this is going to be 20 second per cycle
                y_crop = int(math.sin(2 * math.pi * i / 20 / fps) * height // 2 + height // 2)
                if video_mode == "timbaland":
                    frame = timbaland_img.resize((width, height)).convert("RGBA")
                else:
                    frame = blurred_background.crop((0, y_crop, width, height + y_crop)).convert("RGBA")

                frame.paste(
                    rounded_image, (width // 2 - image_size_x // 2, image_y_pos), mask=rounded_image
                )

                title_y_pos = (185 if video_mode == "timbaland" else 120) * upsample
                frame.paste(
                    song_name, (width // 2 - song_name.size[0] // 2, title_y_pos), mask=song_name
                )

                if video_mode != "timbaland":
                    frame.paste(
                        song_handle,
                        (width // 2 - song_handle.size[0] // 2, 160 * upsample),
                        mask=song_handle,
                    )
                else:
                    frame.paste(
                        timbaland_text_1,
                        (width // 2 - timbaland_text_1.size[0] // 2, 120 * upsample),
                        mask=timbaland_text_1,
                    )
                    frame.paste(
                        timbaland_text_2,
                        (width // 2 - timbaland_text_2.size[0] // 2, 140 * upsample),
                        mask=timbaland_text_2,
                    )
                    frame.paste(
                        small_song_handle,
                        (width // 2 - small_song_handle.size[0] // 2, 160 * upsample),
                        mask=small_song_handle,
                    )

                # frame.paste(song_genre, (width // 2 - song_genre.size[0] // 2, 480), mask=song_genre)

                # Advance to next line if we are not at the end of the lyrics and next line timestamp has been reached
                current_time = i / fps
                right_current_line = bisect.bisect_left(lyrics_line_timestamps, current_time)
                current_line = max(right_current_line - 2, 0)  # backup two lines

                # Render the lines of lyrics
                lyrics_curr_height = (450 if video_mode == "timbaland" else 400) * upsample
                for j in range(current_line, min(current_line + lines_shown, len(lyrics_images))):
                    # TODO: to be adjusted here
                    if j == right_current_line or j == right_current_line - 1:
                        image_to_paste = bold_lyrics_images[j]
                    else:
                        image_to_paste = lyrics_images[j]

                    if j == current_line + lines_shown - 1:
                        image_to_paste = add_gradient_to_image(image_to_paste, 0.6, 0.3)
                    elif j == current_line + lines_shown - 2:
                        image_to_paste = add_gradient_to_image(image_to_paste, 1, 0.6)

                    frame.paste(
                        image_to_paste,
                        ((width - image_to_paste.size[0]) // 2, lyrics_curr_height),
                        mask=image_to_paste,
                    )

                    lyrics_curr_height += 26 * upsample

                frame.paste(
                    logo_line,
                    (
                        width // 2
                        - (logo_line.size[0] // 2)
                        - (logo_small.size[0] // 2)
                        - (3 * upsample),
                        540 * upsample,
                    ),
                    mask=logo_line,
                )
                frame.paste(
                    logo_small,
                    (
                        width // 2
                        + (logo_line.size[0] // 2)
                        - (logo_small.size[0] // 2)
                        + (3 * upsample),
                        540 * upsample,
                    ),
                    mask=logo_small,
                )
                frame.save(process.stdin, "bmp")

            process.stdin.close()
            process.wait()
            print("Time to encode video", time.time() - start_time)

            if process.returncode != 0:
                raise subprocess.CalledProcessError(process.returncode, command)

        num_threads = 4
        total_frames = fps * duration
        i_ranges = [
            (
                math.floor(i * (total_frames // num_threads)),
                math.floor((i + 1) * (total_frames // num_threads)),
            )
            for i in range(num_threads)
        ]

        # Run frame generation in parallel
        threads = []
        for i in range(num_threads):  # frames, cover_path, fps, size, duration, lyrics, i_range
            thread = threading.Thread(
                target=generate_video_chunk, args=(int(i_ranges[i][0]), int(i_ranges[i][1]), image)
            )
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        print("Time to generate frames", time.time() - start_time)

        # Combine the video chunks created into the large video file
        input_files = [(str(temp_dir_string + f"/out{r[0]}.ts")) for r in i_ranges]

        # Write the file paths to a temporary text file
        with open(str(temp_dir_string + "/concat_list.txt"), "w") as f:
            for filepath in input_files:
                f.write(f"file '{filepath}'\n")

        command = [
            "ffmpeg",
            "-loglevel",
            "error",
            "-y",
            "-f",
            "concat",
            "-safe",
            "0",
            "-i",
            f"{temp_dir_string}/concat_list.txt",
            "-i",
            f"{os.path.join(temp_dir_string, 'tmp.wav')}",
            "-c:v",
            "copy",
            "-c:a",
            "aac",
            "-b:a",
            "192k",
            "-map",
            "0:v",
            "-map",
            "1:a",
            "-shortest",
            f"{temp_dir_string}/out.mp4",
        ]
        # Start the ffmpeg process
        process = subprocess.Popen(command, stdin=subprocess.PIPE)

        process.stdin.close()
        process.wait()

        print("Time to genereate total video", time.time() - start_time_total)
        shutil.copy(f"{temp_dir_string}/out.mp4", out_path)
        del image, logo, logo_2, logo_3, timbaland_bg
        return out_path


def generate_video_for_video_to_song(
    video_s3id: str,
    audio: Audio,
    image_path: Path | str,
    aligned_lyrics: list[dict[str, str | float]],
    out_path: Path | str,
    song_title: str,
    tags: str | None,
    fps=5,  # based on performance, if the text is rolling, video is less noticable
    video_mode: str | None = None,
) -> str:
    print("Generating video for video to song")
    start = now = time.time()
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_dir_string = str(temp_dir)
        # Save audio to MP3 for putting into video
        audio.write_mp3(os.path.join(temp_dir_string, "tmp.mp3"))
        audio_duration = audio.duration_s

        video_path = os.path.join(temp_dir_string, "tmp_video.mp4")
        video_s3_path = f"studio/uploads/{video_s3id}.mp4"
        retry_s3_download("suno-data-uploads", video_s3_path, video_path)
        print(f"Downloaded video from {video_s3_path} to {video_path}")

        print(f"Time to download video {time.time() - now} seconds")
        now = time.time()

        video_duration = float(
            subprocess.check_output(
                [
                    "ffprobe",
                    "-v",
                    "error",
                    "-show_entries",
                    "format=duration",
                    "-of",
                    "default=noprint_wrappers=1:nokey=1",
                    video_path,
                ]
            ).strip()
        )

        # Calculate the number of times the video needs to be repeated
        repeat_count = math.ceil(audio_duration / video_duration)
        duration = min(audio_duration, video_duration * repeat_count)

        # Create a text file with repeated video segments for concatenation
        with open(os.path.join(temp_dir_string, "video_list.txt"), "w") as f:
            for _ in range(repeat_count):
                f.write(f"file '{video_path}'\n")

        # Concatenate the repeated video segments
        subprocess.run(
            [
                "ffmpeg",
                "-f",
                "concat",
                "-safe",
                "0",
                "-i",
                os.path.join(temp_dir_string, "video_list.txt"),
                "-c",
                "copy",
                os.path.join(temp_dir_string, "repeated_video.mp4"),
            ],
            check=True,
        )

        print(f"Time to concatenate video {time.time() - now} seconds")
        now = time.time()

        lines, lyrics_line_timestamps = _process_lyrics(aligned_lyrics, duration)

        def create_srt_file(lines, timestamps, total_duration, output_path):
            with open(output_path, "w") as f:
                for i, (line, start_time) in enumerate(zip(lines, timestamps)):
                    end_time = (
                        timestamps[i + 1] if i + 1 < len(timestamps) else total_duration
                    )  # Default end time 5 seconds later
                    f.write(f"{i + 1}\n")
                    f.write(f"{format_timestamp(start_time)} --> {format_timestamp(end_time)}\n")
                    f.write(f"{line}\n\n")

        def format_timestamp(seconds):
            seconds = float(seconds)
            milliseconds = int((seconds % 1) * 1000)
            seconds = int(seconds)
            minutes, seconds = divmod(seconds, 60)
            hours, minutes = divmod(minutes, 60)
            return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"

        create_srt_file(
            lines, lyrics_line_timestamps, duration, os.path.join(temp_dir_string, "lyrics.srt")
        )

        print(f"Time to create SRT file {time.time() - now} seconds")
        now = time.time()

        # # Combine the concatenated video with the audio without subtitles
        # subprocess.run(
        #     [
        #         "ffmpeg",
        #         "-i",
        #         os.path.join(temp_dir_string, "repeated_video.mp4"),
        #         "-i",
        #         os.path.join(temp_dir_string, "tmp.mp3"),
        #         "-c:v",
        #         "copy",
        #         "-c:a",
        #         "aac",
        #         "-strict",
        #         "experimental",
        #         "-shortest",  # Ensure the output duration matches the shortest input (the MP3 audio)
        #         os.path.join(temp_dir_string, "out.mp4"),
        #     ],
        #     check=True,
        # )

        # Combine the concatenated video with the audio and subtitles
        subprocess.run(
            [
                "ffmpeg",
                "-i",
                os.path.join(temp_dir_string, "repeated_video.mp4"),
                "-i",
                os.path.join(temp_dir_string, "tmp.mp3"),
                "-c:v",
                "libx264",  # Re-encode the video using libx264 codec
                "-c:a",
                "aac",
                "-strict",
                "experimental",
                "-shortest",  # Ensure the output duration matches the shortest input (the MP3 audio)
                "-vf",
                f"subtitles={os.path.join(temp_dir_string, 'lyrics.srt')}",  # Add subtitles filter
                os.path.join(temp_dir_string, "out.mp4"),
            ],
            check=True,
        )

        shutil.copy(os.path.join(temp_dir_string, "out.mp4"), out_path)

        print(f"Time to combine video and audio {time.time() - now} seconds")
        print(f"Total time to generate video {time.time() - start} seconds")
        return out_path


def _process_lyrics(aligned_lyrics, duration):
    lines = []
    # this is each line's start time
    lyrics_line_timestamps = []
    current_line = ""
    current_line_timestamps = []
    max_line_length = 24
    aligned_lyrics = [word for word in aligned_lyrics if "word" in word.keys()]
    for index, word in enumerate(aligned_lyrics):
        if len(current_line) < max_line_length:
            current_line += word["word"].split("\n")[0]
            if "start_s" in word.keys():
                current_line_timestamps.append(word["start_s"])
            if "end_s" in word.keys():
                current_line_timestamps.append(word["end_s"])

        # end of the list / line
        if (
            len(current_line) >= max_line_length
            or ("\n" in word["word"])
            or index == len(aligned_lyrics) - 1
        ):
            # Remove double newlines, they look ugly
            # current_line = current_line.replace("\n\n", "\n")
            lines.append(current_line.strip())
            current_line = ""

            # # If there's text after the newline, make that text part of the next line
            if "\n" in word["word"]:
                current_line = " ".join(word["word"].split("\n")[1:])

            lyrics_line_timestamps.append(min(current_line_timestamps, default=0))
            # update to the prev line's last time, or empty
            current_line_timestamps = [max(current_line_timestamps)] if current_line_timestamps else []

    # this is a fallback action
    if all(t == 0 for t in lyrics_line_timestamps):
        # somehow the alignmnent just fails...
        # fall back to evenly spaced timestamps
        # better than nothing
        lyrics_line_timestamps = []
        for i in range(len(lines)):
            lyrics_line_timestamps.append(i / len(lines) * duration)

    return lines, lyrics_line_timestamps


# Useful for testing
if __name__ == "__main__":
    aligned_lyrics = [
        {"word": "[Verse]\nCaught ", "success": True, "start_s": 0.56, "end_s": 0.68, "p_align": 0.915},
        {
            "word": "in ate",
            "success": True,
            "start_s": 0.68,
            "end_s": 0.88,
            "p_align": 0.931,
        },
        {"word": "a ", "success": True, "start_s": 0.88, "end_s": 1.04, "p_align": 0.98},
        {"word": "web ", "success": True, "start_s": 1.04, "end_s": 1.28, "p_align": 0.951},
        {"word": "of ", "success": True, "start_s": 1.28, "end_s": 1.56, "p_align": 0.789},
    ]
    audio = Audio.from_file("/Users/martin/Downloads/031bbe48-0e0f-4310-b848-77cc513f74df_1.mp3")
    generate_video(
        audio,
        "/Users/martin/Downloads/test3.png",
        aligned_lyrics,
        Path("out_copy.mp4"),
        "Oh My Love",
        "pop acoustic",
        "olivermccann",
    )