import os
import requests
import gradio as gr

from time import time
from pytube import YouTube
from moviepy.editor import *
from moviepy.video.tools.subtitles import SubtitlesClip


GLADIA_API_KEY = os.environ.get("GLADIA_API_KEY")


headers = {
    "accept": "application/json",
    "x-gladia-key": GLADIA_API_KEY,
}

ACCEPTED_LANGUAGE_BEHAVIOUR = [
    "manual",
    "automatic single language",
    "automatic multiple languages",
]


def __download_youtube_video(youtube_url: str, path):

    youtube_video = YouTube(youtube_url)

    video_stream = youtube_video.streams.filter(res="720p").first()

    if not video_stream:
        video_stream = youtube_video.streams.filter(res="480p").first()

    if not video_stream:
        video_stream = youtube_video.streams.filter(res="360p").first()

    video_stream.download(filename=path)


def __apply_srt_to_video(path_to_srt: str, path_to_video: str, output_path: str):

    def __generator(raw_txt):

        MIN_LINE_SIZE = 80
        MAX_LINE_SIZE = 140

        text_to_display = []

        while raw_txt:
            idx = MIN_LINE_SIZE if MIN_LINE_SIZE < len(raw_txt) else len(raw_txt)

            while idx <= len(raw_txt) and raw_txt[idx - 1] not in ('\n', ' ', ',', ';', '.', '-') and idx < MAX_LINE_SIZE:
                idx += 1

            text_to_display.append(raw_txt[:idx])
            raw_txt = raw_txt[idx:]

        text_to_display = "\n".join(text_to_display)

        return TextClip(text_to_display, font='Arial', fontsize=30, color='white')

    subs = SubtitlesClip(path_to_srt, __generator)
    subtitles = SubtitlesClip(subs, __generator)

    video = VideoFileClip(path_to_video)

    # Combine video and subtitle clips
    result = CompositeVideoClip([video, subtitles.set_pos(('center','bottom'))])

    # Write the result to a new video file
    result.write_videofile(output_path)


def youtube_subtitler(youtube_url: str):
    files_to_remove = []

    try:
        path_to_srt = f"/tmp/{str(time())}.srt"

        files = {
            'audio_url': (None, youtube_url),
            'output_format': (None, "srt"),
            'language_behaviour': (None, 'automatic multiple languages'),
        }

        response = requests.post(
            'https://api.gladia.io/audio/text/audio-transcription/?output_format=srt',
            headers=headers,
            files=files
        )

        if response.status_code != 200:
            raise RuntimeError(f"Failed to transcribe the video: {response.content}")

        with open(path_to_srt, "w+") as f:
            f.write(response.json()["prediction"])

        files_to_remove.append(path_to_srt)

        path_to_video = f"/tmp/{str(time())}.mp4"

        __download_youtube_video(youtube_url=youtube_url, path=path_to_video)

        files_to_remove.append(path_to_video)

        output_path = f"/tmp/{str(time())}.mp4"

        __apply_srt_to_video(
            path_to_srt=path_to_srt,
            path_to_video=path_to_video,
            output_path=output_path
        )

    finally:
        for file in files_to_remove:
            os.remove(file)

    return output_path


iface = gr.Interface(
    title="Gladia.io Youtube Subtitler",
    description="""Burn subtitles into your youtube video using Gladia's Audio-Transcription solution.
    <br/><br/>
    You are more than welcome to join us on [Slack](https://gladia-io.slack.com)
    and don't forget to get your own API key on [Gladia.io](https://gladia.io/) during the free alpha !
    """,
    fn=youtube_subtitler,
    inputs=[
        gr.Text(label="Youtube URL"),
    ], outputs=[
        gr.Video(format="mp4"),
    ],
)

print("Starting demo")

iface.launch(server_name="0.0.0.0", server_port=7860)