import os
import moviepy.editor as mp
from pytube import YouTube
import gradio as gr

def download_video(url, quality, crop_start, crop_end):
    yt = YouTube(url)
    stream = yt.streams.filter(only_audio=False).first()
    stream.on_progress = lambda stream, chunk, bytes_remaining: None
    stream.download(output_path="downloads")

    # Wait until the download is complete
    while not os.path.exists(os.path.join("downloads", yt.title + ".mp4")):
        pass

    video_path = os.path.join("downloads", yt.title + ".mp4")
    clip = mp.VideoFileClip(video_path)

    if crop_start and crop_end:
        clip = clip.subclip(crop_start, crop_end)

    audio_path = os.path.join("downloads", yt.title + ".mp3")
    clip.audio.write_audiofile(audio_path)

    return {"video": video_path, "audio": audio_path}

def main():
    iface = gr.Interface(
        fn=download_video,
        inputs=[
            gr.Textbox(label="YouTube URL"),
            gr.Radio(label="Quality", choices=["144p", "240p", "360p", "480p", "720p", "1080p"]),
            gr.Number(label="Crop start (seconds)"),
            gr.Number(label="Crop end (seconds)"),
        ],
        outputs=[
            gr.File(label="Video"),
            gr.File(label="Audio"),
        ],
        title="YouTube Audio and Video Downloader with Cropping",
        description="Enter a YouTube URL, select a quality, and optionally enter crop start and end times to download the video and audio with cropping.",
    )
    iface.launch()

if __name__ == "__main__":
    main()