from moviepy import VideoFileClip
import gradio as gr
import requests
import os
import time
from gradio_client import Client, handle_file
from pathlib import Path


# Store processed videos temporarily
PROCESSED_DIR = "processed_videos"
Path(PROCESSED_DIR).mkdir(exist_ok=True)

def main(url, parameters, progress=gr.Progress()):
    try:
        font_type, font_size, font_color, service, target, style, subject = parameters.split(',')
        progress(0, desc="پردازش شروع شد")
        yield "پردازش شروع شد", gr.DownloadButton.update(visible=False)

        # Transcription Stage
        progress(0.35, desc="تبدیل صوت به متن")
        yield "تبدیل صوت به متن", gr.DownloadButton.update(visible=False)
        
        transcribe_client = Client("rayesh/transcribe")
        try:
            job = transcribe_client.submit(url, api_name="/transcribe")
            while not job.done():
                time.sleep(0.5)
                
            results = job.outputs()[0]
            if len(results) != 3:
                raise ValueError(f"Expected 3 outputs, got {len(results)}")
                
            srt_file, video, mp3_file = results
        finally:
            transcribe_client.close()

        # Translation Stage
        progress(0.55, desc="در حال ترجمه")
        yield "در حال ترجمه", gr.DownloadButton.update(visible=False)
        
        translate_client = Client("rayesh/translate")
        try:
            job = translate_client.submit(
                handle_file(srt_file),
                3000,
                api_name="/translate"
            )
            while not job.done():
                time.sleep(0.3)
            subtitle_file = job.outputs()[0]
        finally:
            translate_client.close()

        # Video Processing Stage
        progress(0.75, desc="پردازش ویدئو")
        yield "درحال پردازش ویدئو", gr.DownloadButton.update(visible=False)
        
        video_client = Client("SPACERUNNER99/video_edite")
        try:
            job = video_client.submit(
                handle_file(subtitle_file),
                {"video": handle_file(video["video"])},
                font_color,
                font_type,
                font_size,
                handle_file(mp3_file),
                api_name="/video_edit"
            )
            while not job.done():
                time.sleep(0.5)
                
            output_video = job.outputs()[0]['video']
            
            # Save to processed directory with unique name
            final_path = str(Path(PROCESSED_DIR) / f"output_{int(time.time())}.mp4")
            Path(output_video).rename(final_path)
            
        finally:
            video_client.close()

        yield "پردازش کامل شد", gr.DownloadButton.update(value=final_path, visible=True)

    except Exception as e:
        raise gr.Error(f"خطا در پردازش: {str(e)}")

def after_download():
    return [gr.DownloadButton.update(visible=False), gr.Button.update(visible=True)]

with gr.Blocks() as demo:
    gr.Markdown("## ویدئو خود را آپلود کنید یا لینک وارد نمایید")
    with gr.Column():
        progress_output = gr.Textbox(label="وضعیت پردازش", visible=False)
        with gr.Row():
            video_file_input = gr.Text(label="لینک ویدئو")
            upload_btn = gr.UploadButton("فایل ویدئویی", file_types=["video/*"])
        parameters = gr.Text(label="تنظیمات زیرنویس (فونت,اندازه,رنگ)")
        create_btn = gr.Button("شروع پردازش")
        download_btn = gr.DownloadButton("دانلود ویدئو نهایی", visible=False)
    
    # Handle file upload
    upload_btn.upload(
        lambda f: [gr.Textbox.update(value=f.name), gr.DownloadButton.update(visible=False)],
        upload_btn,
        [video_file_input, download_btn]
    )
    
    # Processing handler
    create_btn.click(
        fn=main,
        inputs=[video_file_input, parameters],
        outputs=[progress_output, download_btn],
        concurrency_limit=4,
    )
    
    # Download completion handler
    download_btn.click(
        fn=after_download,
        inputs=None,
        outputs=[download_btn, create_btn]
    )

demo.launch(debug=True)