import gradio as gr
from pytube import YouTube
from pytube.exceptions import RegexMatchError, VideoUnavailable
import os
import subprocess

def download_media(url, resolution, media_format):
    try:
        yt = YouTube(url)
        output_path = "downloads"
        os.makedirs(output_path, exist_ok=True)
        
        # Очищаем название файла от недопустимых символов
        title = "".join(c for c in yt.title if c.isalnum() or c in (' ', '_')).rstrip()
        
        if media_format == "mp3":
            audio_stream = yt.streams.get_audio_only()
            out_file = audio_stream.download(output_path=output_path, filename=title)
            base, _ = os.path.splitext(out_file)
            new_file = f"{base}.mp3"
            os.rename(out_file, new_file)
            return new_file
        
        else:  # mp4
            if resolution in ["1080p", "4k"]:
                # Для высоких разрешений
                video_stream = yt.streams.filter(res=resolution, file_extension="mp4", type="video").first()
                audio_stream = yt.streams.get_audio_only()
                
                if not video_stream or not audio_stream:
                    raise Exception("Required streams not found")
                
                video_path = video_stream.download(output_path=output_path, filename_prefix="video_")
                audio_path = audio_stream.download(output_path=output_path, filename_prefix="audio_")
                
                output_file = os.path.join(output_path, f"{title}.mp4")
                subprocess.run([
                    'ffmpeg', '-i', video_path, '-i', audio_path,
                    '-c:v', 'copy', '-c:a', 'aac', '-strict', 'experimental',
                    output_file
                ], check=True)
                
                os.remove(video_path)
                os.remove(audio_path)
                return output_file
            
            else:
                stream = yt.streams.filter(res=resolution, file_extension="mp4", progressive=True).first()
                if not stream:
                    raise Exception("Stream not found")
                return stream.download(output_path=output_path, filename=title + ".mp4")
                
    except Exception as e:
        print(f"Error: {str(e)}")
        return None  # Возвращаем None вместо строки с ошибкой

# Создаем интерфейс с отдельным компонентом для ошибок
with gr.Blocks() as app:
    gr.Markdown("## YouTube Video Downloader")
    with gr.Row():
        url = gr.Textbox(label="YouTube URL")
        resolution = gr.Dropdown(["144p", "360p", "720p", "1080p", "4k"], label="Resolution")
        media_format = gr.Dropdown(["mp4", "mp3"], label="Format")
    
    download_button = gr.Button("Download")
    output_file = gr.File(label="Downloaded File")
    error_output = gr.Textbox(label="Error Message", visible=False)

    def handle_download(url, resolution, media_format):
        result = download_media(url, resolution, media_format)
        if result and os.path.exists(result):
            return [result, gr.Textbox(visible=False)]
        return [None, gr.Textbox(value="Error: Failed to download. Possible reasons:\n- Invalid URL\n- Restricted content\n- Unavailable format", visible=True)]

    download_button.click(
        fn=handle_download,
        inputs=[url, resolution, media_format],
        outputs=[output_file, error_output]
    )

if __name__ == "__main__":
    app.launch(share=True)  # Добавлен параметр share=True