import streamlit as st
from pytubefix import YouTube
import tempfile
import os
from urllib.error import HTTPError
from typing import Optional, List, Dict, Tuple
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Page Configuration
st.set_page_config(
    page_title="Educational Video Downloader",
    page_icon="📚",
    layout="wide"
)

class VideoDownloader:
    def __init__(self):
        self.custom_user_agent = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                         "(KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
        }

    def get_youtube(self, url: str) -> Optional[YouTube]:
        """Initialize YouTube object with error handling"""
        try:
            return YouTube(url, self.custom_user_agent)
        except Exception as e:
            logger.error(f"Error initializing YouTube object: {e}")
            st.error(f"Error accessing video: {str(e)}")
            return None

    def process_streams(self, yt: YouTube, stream_type: str = 'video') -> Tuple[List, List, str]:
        """Process and return streams information"""
        try:
            if stream_type == 'video':
                # Convert streams to list first
                all_streams = list(yt.streams.filter(progressive=True, file_extension='mp4'))
                # Extract unique resolutions
                resolutions = []
                seen = set()
                for stream in all_streams:
                    res = getattr(stream, 'resolution', None)
                    if res and res not in seen:
                        resolutions.append(res)
                        seen.add(res)
                return all_streams, sorted(resolutions, reverse=True), yt.title
            else:
                # Convert streams to list first
                all_streams = list(yt.streams.filter(only_audio=True))
                # Extract unique bitrates
                bitrates = []
                seen = set()
                for stream in all_streams:
                    abr = getattr(stream, 'abr', None)
                    if abr and abr not in seen:
                        bitrates.append(abr)
                        seen.add(abr)
                return all_streams, sorted(bitrates, reverse=True), yt.title
        except Exception as e:
            logger.error(f"Error processing streams: {e}")
            st.error(f"Error processing available qualities: {str(e)}")
            return [], [], ""

    def download_content(self, stream, title: str, extension: str) -> None:
        """Download and provide content to user"""
        try:
            with tempfile.TemporaryDirectory() as temp_dir:
                filename = f"{title}.{extension}".replace(" ", "_")
                temp_path = os.path.join(temp_dir, filename)
                
                with st.spinner(f"Downloading {filename}..."):
                    stream.download(output_path=temp_dir, filename=filename)
                
                with open(temp_path, 'rb') as f:
                    st.download_button(
                        label=f"📥 Save {extension.upper()} File",
                        data=f.read(),
                        file_name=filename,
                        mime=f"video/{extension}" if extension == "mp4" else "audio/mp3"
                    )
        except Exception as e:
            logger.error(f"Error during download: {e}")
            st.error(f"Download failed: {str(e)}")

def main():
    downloader = VideoDownloader()
    
    # Sidebar
    st.sidebar.title("📚 Educational Video Tool")
    st.sidebar.markdown("""
    **Usage Guide:**
    1. Enter a valid educational video URL
    2. Select desired quality
    3. Download for offline study
    """)
    
    # Main Content
    st.title("Educational Content Downloader")
    st.markdown("*For educational and research purposes only*")
    
    # Disclaimer
    st.warning("⚠️ This tool is intended for educational content only. "
               "Please respect copyright and terms of service.")
    
    # Content Tabs
    tab1, tab2 = st.tabs(["Video Materials", "Audio Materials"])
    
    # Video Tab
    with tab1:
        st.subheader("Download Educational Videos")
        url = st.text_input("Enter Educational Video URL:", key="video_url")
        
        if url:
            yt = downloader.get_youtube(url)
            if yt:
                streams, qualities, title = downloader.process_streams(yt, 'video')
                if streams and qualities:
                    st.success(f"📺 Found: {title}")
                    quality = st.selectbox(
                        "Select Video Quality:",
                        qualities
                    )
                    
                    selected_stream = next(
                        (s for s in streams if s.resolution == quality),
                        None
                    )
                    
                    if selected_stream and st.button("Prepare Download", type="primary"):
                        downloader.download_content(
                            selected_stream,
                            title,
                            'mp4'
                        )
    
    # Audio Tab
    with tab2:
        st.subheader("Download Audio Materials")
        url = st.text_input("Enter Educational Video URL:", key="audio_url")
        
        if url:
            yt = downloader.get_youtube(url)
            if yt:
                streams, qualities, title = downloader.process_streams(yt, 'audio')
                if streams and qualities:
                    st.success(f"🎵 Found: {title}")
                    quality = st.selectbox(
                        "Select Audio Quality:",
                        qualities
                    )
                    
                    selected_stream = next(
                        (s for s in streams if s.abr == quality),
                        None
                    )
                    
                    if selected_stream and st.button("Prepare Audio Download", type="primary"):
                        downloader.download_content(
                            selected_stream,
                            title,
                            'mp3'
                        )

if __name__ == "__main__":
    main()