import yt_dlp
import gradio as gr
import os

DOWNLOAD_DIR = "downloads"
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

def clean_url(url: str) -> str:
    """Convert youtu.be short links to full format"""
    if "youtu.be" in url:
        return url.replace("youtu.be/", "www.youtube.com/watch?v=")
    return url

def get_video_info(url):
    try:
        url = clean_url(url)
        ydl_opts = {"quiet": True, "skip_download": True, "nocheckcertificate": True}
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=False)
            formats = []
            for f in info.get("formats", []):
                if f.get("vcodec") != "none":
                    resolution = f.get("format_note") or f"{f.get('height')}p"
                    if resolution:
                        formats.append({
                            "resolution": resolution,
                            "ext": f.get("ext"),
                            "format_id": f.get("format_id")
                        })
            # remove duplicates by resolution
            unique = {f["resolution"]: f for f in formats}
            return info.get("title", "Unknown"), list(unique.values())
    except Exception as e:
        return f"Error: {str(e)}", []

def download_video(url, quality):
    try:
        url = clean_url(url)
        title, formats = get_video_info(url)
        if not formats:
            return None, "No formats available."
        
        # find selected format
        format_id = None
        ext = "mp4"
        for f in formats:
            if f["resolution"] == quality:
                format_id = f["format_id"]
                ext = f["ext"]
                break
        
        if not format_id:
            return None, "Selected quality not available."
        
        # unique file path
        safe_title = "".join(c if c.isalnum() else "_" for c in title)[:50]
        filepath = os.path.join(DOWNLOAD_DIR, f"{safe_title}_{quality}.{ext}")
        
        ydl_opts = {
            "format": format_id,
            "outtmpl": filepath,
            "quiet": True,
            "nocheckcertificate": True
        }
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url])
        
        return filepath, f"Downloaded {title} in {quality}"
    except Exception as e:
        return None, f"Error: {str(e)}"

# ---------- Gradio UI ----------
with gr.Blocks() as demo:
    gr.Markdown("## 🎥 YouTube Video Downloader (Demo)")

    url_in = gr.Textbox(label="Enter YouTube URL")
    fetch_btn = gr.Button("Fetch Info")
    title_out = gr.Textbox(label="Video Title", interactive=False)
    quality_out = gr.Dropdown(label="Select Quality")
    
    download_btn = gr.Button("Download")
    result_msg = gr.Textbox(label="Status", interactive=False)
    download_file = gr.File(label="Download File")

    def fetch_info(url):
        title, formats = get_video_info(url)
        qualities = [f["resolution"] for f in formats]
        return title, gr.Dropdown.update(choices=qualities)

    fetch_btn.click(fetch_info, inputs=[url_in], outputs=[title_out, quality_out])
    download_btn.click(download_video, inputs=[url_in, quality_out], outputs=[download_file, result_msg])

demo.launch()