from pytubefix import YouTube
from pytubefix.cli import on_progress
import time
import math
import gradio
import ffmpeg
from faster_whisper import WhisperModel
import requests
import json
import arabic_reshaper # pip install arabic-reshaper
from bidi.algorithm import get_display # pip install python-bidi
from moviepy.editor import VideoFileClip, TextClip, CompositeVideoClip
import pysrt
api_key = "268976:66f4f58a2a905"




# Define your functions here
def yt_download(url):
    yt = YouTube(url)
    print(yt.title)
    video_path = f"{yt.title}.mp4"
    ys = yt.streams.get_highest_resolution()
    print(ys)
    ys.download()
    return video_path, yt.title

def insta_download(shortcode, id_1):
    url = "https://api.one-api.ir/instagram/v1/post/"
    request_body = {"shortcode": shortcode, "id": id_1}
    headers = {"one-api-token": api_key, "Content-Type": "application/json"}
    response = requests.post(url, headers=headers, json=request_body)
    if response.status_code == 200:
        result = response.json()
        return result['result']
    else:
        print(f"Error: {response.status_code}, {response.text}")
        return None

def extract_audio(input_video_name):
    extracted_audio = f"audio-{input_video_name}.wav"
    stream = ffmpeg.input(input_video)
    stream = ffmpeg.output(stream, extracted_audio)
    ffmpeg.run(stream, overwrite_output=True)
    return extracted_audio

def transcribe(audio):
    model = WhisperModel("tiny")
    segments, info = model.transcribe(audio)
    segments = list(segments)
    for segment in segments:
        print("[%.2fs -> %.2fs] %s" % (segment.start, segment.end, segment.text))
    return segments

def format_time(seconds):
    hours = math.floor(seconds / 3600)
    seconds %= 3600
    minutes = math.floor(seconds / 60)
    seconds %= 60
    milliseconds = round((seconds - math.floor(seconds)) * 1000)
    seconds = math.floor(seconds)
    formatted_time = f"{hours:02d}:{minutes:02d}:{seconds:01d},{milliseconds:03d}"
    return formatted_time

def generate_subtitle_file(language, segments, input_video_name):
    subtitle_file = f"sub-{input_video_name}.{language}.srt"
    text = ""
    for index, segment in enumerate(segments):
        segment_start = format_time(segment.start)
        segment_end = format_time(segment.end)
        text += f"{str(index+1)} \n"
        text += f"{segment_start} --> {segment_end} \n"
        text += f"{segment.text} \n"
        text += "\n"
    f = open(subtitle_file, "w", encoding='utf8')
    f.write(text)
    f.close()
    return subtitle_file

def read_srt_file(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            srt_content = file.read()
            return srt_content
    except FileNotFoundError:
        print(f"The file {file_path} was not found.")
    except Exception as e:
        print(f"An error occurred: {e}")

def translate_text(api_key, source_lang, target_lang, text):
    url = "https://api.one-api.ir/translate/v1/google/"
    request_body = {"source": source_lang, "target": target_lang, "text": text}
    headers = {"one-api-token": api_key, "Content-Type": "application/json"}
    response = requests.post(url, headers=headers, json=request_body)
    if response.status_code == 200:
        result = response.json()
        return result['result']
    else:
        print(f"Error: {response.status_code}, {response.text}")
        return None

def write_google(google_translate):
    google = "google_translate.srt"
    with open(google, 'w', encoding="utf-8") as f:
        f.write(google_translate)

def time_to_seconds(time_obj):
    return time_obj.hours * 3600 + time_obj.minutes * 60 + time_obj.seconds + time_obj.milliseconds / 1000

def create_subtitle_clips(subtitles, videosize, fontsize=24, font='/content/arial-unicode-ms.ttf', color='yellow', debug=False):
    subtitle_clips = []
    for subtitle in subtitles:
        start_time = time_to_seconds(subtitle.start)
        end_time = time_to_seconds(subtitle.end)
        duration = end_time - start_time
        video_width, video_height = videosize
        reshaped_text = arabic_reshaper.reshape(subtitle.text)
        bidi_text = get_display(reshaped_text)
        text_clip = TextClip(bidi_text, fontsize=fontsize, font=font, color=color, bg_color='black', size=(video_width*3/4, None), method='label', align='West').set_start(start_time).set_duration(duration)
        subtitle_x_position = 'center'
        subtitle_y_position = video_height * 4 / 5
        text_position = (subtitle_x_position, subtitle_y_position)
        subtitle_clips.append(text_clip.set_position(text_position))
    return subtitle_clips

def process_video(url, api_key):
    input_video, title = yt_download(url)
    input_video_name = input_video.replace(".mp4", "")
    extracted_audio = extract_audio(input_video_name)
    segments = transcribe(audio=extracted_audio)
    language = "fa"
    subtitle_file = generate_subtitle_file(language=language, segments=segments, input_video_name=input_video_name)
    source_language = "en"
    target_language = "fa"
    srt_string = read_srt_file(subtitle_file)
    google_translate = translate_text(api_key, source_language, target_language, srt_string)
    write_google(google_translate)
    video = VideoFileClip(input_video)
    subtitles = pysrt.open("/content/google_translate.srt", encoding="utf-8")
    output_video_file = input_video_name + '_subtitled' + ".mp4"
    subtitle_clips = create_subtitle_clips(subtitles, video.size)
    final_video = CompositeVideoClip([video] + subtitle_clips)
    final_video.write_videofile(output_video_file, codec="libx264", audio_codec="aac")
    return output_video_file

def download_file(file_path):
    return gr.File.update(file_path)

iface = gr.Interface(fn=process_video, inputs=["text", "text"], outputs="file")

iface.launch()