import streamlit as st
from moviepy.editor import VideoFileClip
from pytube import YouTube
import os

def clean_youtube_url(url):
    """Optional: Remove extra playlist parameters from the YouTube URL."""
    from urllib.parse import urlparse, parse_qs, urlunparse
    parsed = urlparse(url)
    query = parse_qs(parsed.query)
    cleaned_query = {k: v for k, v in query.items() if k == "v"}
    cleaned_query_str = "&".join(f"{k}={v[0]}" for k, v in cleaned_query.items())
    cleaned_url = urlunparse((parsed.scheme, parsed.netloc, parsed.path, '', cleaned_query_str, ''))
    return cleaned_url

def is_valid_youtube_url(url):
    try:
        from urllib.parse import urlparse
        parsed = urlparse(url)
        return "youtube.com" in parsed.netloc or "youtu.be" in parsed.netloc
    except Exception:
        return False

def download_youtube_video(url):
    try:
        yt = YouTube(url)
        # Get the highest resolution progressive stream (includes audio)
        stream = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first()
        if not stream:
            st.error("No video stream found!")
            return None
        video_path = stream.download()
        return video_path
    except Exception as e:
        st.error(f"Error downloading video: {e}")
        return None

def save_uploaded_file(uploaded_file):
    video_path = os.path.join("uploaded_videos", uploaded_file.name)
    os.makedirs(os.path.dirname(video_path), exist_ok=True)
    with open(video_path, 'wb') as f:
        f.write(uploaded_file.getbuffer())
    return video_path

def process_and_show_downloads(video_path, mute_audio):
    audio_path = os.path.splitext(video_path)[0] + '.mp3'
    try:
        video_clip = VideoFileClip(video_path)
        # Extract audio
        video_clip.audio.write_audiofile(audio_path)
        # Mute the video if needed
        if mute_audio:
            muted_video_path = os.path.join("muted_videos", os.path.basename(video_path))
            os.makedirs(os.path.dirname(muted_video_path), exist_ok=True)
            video_clip.without_audio().write_videofile(muted_video_path)
            video_path = muted_video_path
        show_download_buttons(video_path, audio_path)
    finally:
        # Optionally remove the intermediate video file after processing
        if os.path.exists(video_path):
            os.remove(video_path)

def show_download_buttons(video_path, audio_path):
    st.video(video_path)
    with open(video_path, "rb") as file:
        st.download_button(label="Download Video", data=file, file_name=os.path.basename(video_path), mime="video/mp4")
    with open(audio_path, "rb") as file:
        st.download_button(label="Download Audio", data=file, file_name=os.path.basename(audio_path), mime="audio/mpeg")
        os.remove(audio_path)

def main():
    st.title("Video and Audio Separator with Pytube")

    with st.form("youtube_download_form"):
        video_url = st.text_input("Enter a YouTube video URL")
        mute_audio = st.checkbox("Mute Audio in Video", value=True)
        submit_button = st.form_submit_button(label="Download and Process YouTube Video")

    if submit_button and video_url:
        if not is_valid_youtube_url(video_url):
            st.error("Invalid YouTube URL. Please enter a proper video link.")
        else:
            cleaned_url = clean_youtube_url(video_url)
            st.info(f"Processing video from: {cleaned_url}")
            video_path = download_youtube_video(cleaned_url)
            if video_path:
                process_and_show_downloads(video_path, mute_audio)

    with st.form("file_upload_form"):
        uploaded_video_file = st.file_uploader("Or upload a video file", type=["mp4", "avi", "mov"])
        mute_audio_uploaded = st.checkbox("Mute Audio in Uploaded Video", value=True)
        process_button = st.form_submit_button(label="Process Uploaded Video")

    if process_button and uploaded_video_file is not None:
        video_path = save_uploaded_file(uploaded_video_file)
        process_and_show_downloads(video_path, mute_audio_uploaded)

if __name__ == "__main__":
    main()