# Code produces bark videos?
import os
import tempfile
import shutil
import subprocess
import textwrap
from typing import Any

from matplotlib import font_manager
import matplotlib.pyplot as plt
import numpy as np

# STUPID HACK TO MAKE IT WORK, we need to depreciate this dependency
try:
    from gradio import processing_utils, utils
except ImportError:
    processing_utils = None
    utils = None
import PIL
import PIL.Image

import requests
import giphy_client
from giphy_client.rest import ApiException
from moviepy.editor import *
from pydub import AudioSegment
import imageio
from PIL import Image, ImageDraw, ImageSequence, ImageFont

import random


def find_gif(api_key, keyword):
    api_instance = giphy_client.DefaultApi()
    try:
        output_file_name = "background_gif.gif"
        random.seed()
        offset = random.randint(0, 100)
        response = api_instance.gifs_search_get(api_key, keyword, limit=1, rating="g", offset=offset)
        print(response)
        gif_url = response.data[0].images.original.url
        gif_file = requests.get(gif_url, stream=True)
        if gif_file.status_code == 200:
            with open(output_file_name, "wb") as f:
                for chunk in gif_file.iter_content(1024):
                    f.write(chunk)
        print(f"Chosen GIF URL: {gif_url}")
        # ToDo: Check the dimensions of the gif_file (they need to be even and within twitter's range). Regenerate if not.
        return output_file_name
    except ApiException as e:
        print(f"Exception when calling DefaultApi->gifs_search_get: {e}\n")


def add_text_overlay_to_gif(gif_filepath, text):
    # Load the gif
    gif = Image.open(gif_filepath)

    # Prepare to add text to each frame
    frames = []
    for frame in ImageSequence.Iterator(gif):
        # Duplicate the frame so we don't modify the original image
        new_frame = frame.copy()

        # Add a semi-transparent rectangle
        rectangle = Image.new("RGBA", new_frame.size, (0, 0, 0, 250))
        new_frame.paste(rectangle, mask=rectangle)

        # Add text to the rectangle
        draw = ImageDraw.Draw(new_frame)
        font = ImageFont.load_default()  # Use default font, you may want to provide your own
        draw.text(
            (10, 10), text, fill="white", font=font
        )  # Adjust coordinates and fill color as desired

        frames.append(new_frame)

    # Save new gif
    frames[1].save(gif_filepath, save_all=True, append_images=frames[1:], loop=0)


def create_video(mp3_filepath, gif_filepath, output_filepath):
    # Get the duration of the audio file
    audio = AudioSegment.from_mp3(mp3_filepath)
    duration = len(audio) / 1000  # pydub calculates in millisec
    print("mp3 duration ", duration)

    # Load gif using imageio
    gif = imageio.get_reader(gif_filepath)

    # Get the fps to calculate duration of the gif, default to 10 if not available
    fps = gif.get_meta_data().get("fps", 10)

    # Calculate the number of frames in the gif
    num_frames = sum(1 for _ in gif)
    print("num frames ", num_frames)

    # Calculate duration of the gif
    gif_duration = num_frames / fps

    # Calculate the number of times gif loops
    loops = int(duration // gif_duration)

    # Calculate the duration of each frame for the writer
    frame_duration = 1 / fps * 1000  # Convert to milliseconds

    # Create temp gif file
    temp_gif_filepath = "temp.gif"

    # Loop the gif
    with imageio.get_writer(temp_gif_filepath, mode="I", duration=frame_duration) as writer:
        for _ in range(loops):
            for frame in gif:
                writer.append_data(frame)

    # Using ffmpeg to combine gif and audio into a video
    cmd = f"ffmpeg -y -stream_loop -1 -i {gif_filepath} -i {mp3_filepath} -c:v libx264 -pix_fmt yuv420p -shortest {output_filepath}"

    subprocess.check_call(cmd)
    # Delete temp gif file
    os.remove(temp_gif_filepath)

    return os.path.abspath(output_filepath)


def make_waveform(
    audio: str | tuple[int, np.ndarray],
    *,
    bg_color: str = "#f3f4f6",
    bg_image: str | None = None,
    fg_alpha: float = 0.80,
    bars_color: str | tuple[str, str] = ("#fbbf24", "#ea580c"),
    bar_count: int = 50,
    bar_width: float = 0.6,
    text: str = "",
) -> str:
    """
    Generates a waveform video from an audio file. Useful for creating an easy to share audio visualization. The output should be passed into a `gr.Video` component.
    Parameters:
        audio: Audio file path or tuple of (sample_rate, audio_data)
        bg_color: Background color of waveform (ignored if bg_image is provided)
        bg_image: Background image of waveform
        fg_alpha: Opacity of foreground waveform
        bars_color: Color of waveform bars. Can be a single color or a tuple of (start_color, end_color) of gradient
        bar_count: Number of bars in waveform
        bar_width: Width of bars in waveform. 1 represents full width, 0.5 represents half width, etc.
    Returns:
        A filepath to the output video in mp4 format.
    """
    if isinstance(audio, str):
        audio_file = audio
        audio = processing_utils.audio_from_file(audio)
    else:
        tmp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)

        processing_utils.audio_to_file(audio[0], audio[1], tmp_wav.name, format="wav")
        audio_file = tmp_wav.name

    if not os.path.isfile(audio_file):
        raise ValueError("Audio file not found.")

    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg not found.")

    duration = round(len(audio[1]) / audio[0], 4)

    # Helper methods to create waveform
    def hex_to_rgb(hex_str):
        return [int(hex_str[i : i + 2], 16) for i in range(1, 6, 2)]

    def get_color_gradient(c1, c2, n):
        assert n > 1
        c1_rgb = np.array(hex_to_rgb(c1)) / 255
        c2_rgb = np.array(hex_to_rgb(c2)) / 255
        mix_pcts = [x / (n - 1) for x in range(n)]
        rgb_colors = [((1 - mix) * c1_rgb + (mix * c2_rgb)) for mix in mix_pcts]
        return ["#" + "".join(f"{int(round(val * 255)):02x}" for val in item) for item in rgb_colors]

    # Reshape audio to have a fixed number of bars
    samples = audio[1]
    if len(samples.shape) > 1:
        samples = np.mean(samples, 1)
    bins_to_pad = bar_count - (len(samples) % bar_count)
    samples = np.pad(samples, [(0, bins_to_pad)])
    samples = np.reshape(samples, (bar_count, -1))
    samples = np.abs(samples)
    samples = np.max(samples, 1)

    width = 800  # in pixels
    height = 400  # in pixels

    with utils.MatplotlibBackendMananger():
        plt.clf()
        # Plot waveform
        color = (
            bars_color
            if isinstance(bars_color, str)
            else get_color_gradient(bars_color[0], bars_color[1], bar_count)
        )

        # Create a figure and axes with the specified dimensions
        fig, ax = plt.subplots(figsize=(width / 100, height / 100), dpi=100)

        ax.bar(
            np.arange(0, bar_count),
            samples * 2,
            bottom=(-1 * samples),
            width=bar_width,
            color=color,
        )

        wrapped_text = "\n".join(textwrap.fill(t.strip(), width=80) for t in text.split("\n"))

        ax.text(
            0.5,
            0.5,
            wrapped_text,
            transform=ax.transAxes,
            ha="center",
            va="center",
            linespacing=1.7,
            fontsize=11 if len(text) < 200 else 9,
            fontproperties=font_manager.FontProperties(
                fname="/suno/models/assets/NotoSansJP-Medium.ttf"
            ),
            bbox=dict(pad=20, facecolor="white", alpha=0.7),
        )

        ax.axis("off")
        ax.margins(x=0)
        tmp_img = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        savefig_kwargs: dict[str, Any] = {"bbox_inches": "tight"}
        if bg_image is not None:
            savefig_kwargs["transparent"] = True
        else:
            savefig_kwargs["facecolor"] = bg_color
        plt.savefig(tmp_img.name, **savefig_kwargs)
        waveform_img = PIL.Image.open(tmp_img.name)
        waveform_img = waveform_img.resize((width, height))

        # Composite waveform with background image
        if bg_image is not None:
            waveform_array = np.array(waveform_img)
            waveform_array[:, :, 3] = waveform_array[:, :, 3] * fg_alpha
            waveform_img = PIL.Image.fromarray(waveform_array)

            bg_img = PIL.Image.open(bg_image)
            waveform_width, waveform_height = waveform_img.size
            bg_width, bg_height = bg_img.size
            if waveform_width != bg_width:
                bg_img = bg_img.resize(
                    (waveform_width, 2 * int(bg_height * waveform_width / bg_width / 2))
                )
                bg_width, bg_height = bg_img.size
            composite_height = max(bg_height, waveform_height)
            composite = PIL.Image.new("RGBA", (waveform_width, composite_height), "#FFFFFF")
            composite.paste(bg_img, (0, composite_height - bg_height))
            composite.paste(waveform_img, (0, composite_height - waveform_height), waveform_img)
            composite.save(tmp_img.name)
            img_width, img_height = composite.size
        else:
            img_width, img_height = waveform_img.size
            waveform_img.save(tmp_img.name)

    # Convert waveform to video with ffmpeg
    output_mp4 = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)

    ffmpeg_cmd = [
        ffmpeg,
        "-loop",
        "1",
        "-i",
        tmp_img.name,
        "-i",
        audio_file,
        "-vf",
        f"color=c=#AAAAAA33:s={img_width}x{40}[bar];[0][bar]overlay=-w+(w/{duration})*t:H-h:shortest=1",
        "-t",
        str(duration),
        "-y",
        output_mp4.name,
    ]

    subprocess.check_call(ffmpeg_cmd)
    return output_mp4.name