app.py

from __future__ import annotations

import os
import glob
from typing import Optional, Literal

from pydantic import BaseModel, Field, HttpUrl

from fastmcp import FastMCP


mcp = FastMCP(
    name="youtube-audio",
    host="0.0.0.0",
    port=7860,
)


class DownloadResult(BaseModel):
    url: HttpUrl = Field(..., description="Original YouTube URL")
    title: Optional[str] = Field(None, description="Video title if available")
    filepath: str = Field(..., description="Absolute path to the downloaded audio file in the container")
    ext: str = Field(..., description="Audio file extension, e.g., mp3")


@mcp.tool(description="Download audio from a YouTube video URL and return the local file path inside the container.")
def download_youtube_audio(url: HttpUrl, audio_format: Literal["mp3", "m4a", "wav", "aac", "opus"] = "mp3") -> DownloadResult:
    """
    - url: A direct YouTube video URL
    - audio_format: Desired audio format (requires ffmpeg in the container)

    The file will be saved under /app/downloads. Ensure the container has write access.

    Networking:
    - If outbound DNS is restricted, set env DNS_SERVERS (space-separated, e.g. "8.8.8.8 1.1.1.1").
    - To route via a proxy, set env YT_PROXY (e.g. http://user:pass@proxy:port).
    """
    # Ensure output directory exists
    output_dir = "/app/downloads"
    os.makedirs(output_dir, exist_ok=True)

    try:
        import yt_dlp as ytdlp
    except Exception:
        raise RuntimeError("yt-dlp is required. Ensure it is listed in requirements.txt and installed.")

    ydl_opts = {
        "format": "bestaudio/best",
        "outtmpl": os.path.join(output_dir, "%(title).200s [%(id)s].%(ext)s"),
        "postprocessors": [
            {
                "key": "FFmpegExtractAudio",
                "preferredcodec": audio_format,
                "preferredquality": "0",
            }
        ],
        "noplaylist": True,
        "quiet": True,
        "nocheckcertificate": True,
    }

    proxy_url = os.environ.get("YT_PROXY")
    if proxy_url:
        ydl_opts["proxy"] = proxy_url

    info_title: Optional[str] = None
    downloaded_id: Optional[str] = None

    with ytdlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(str(url), download=True)
        info_title = info.get("title") if isinstance(info, dict) else None
        downloaded_id = info.get("id") if isinstance(info, dict) else None

    # Resolve the final filename after post-processing
    final_path: Optional[str] = None
    if downloaded_id:
        pattern = os.path.join(output_dir, f"*[{downloaded_id}].{audio_format}")
        matches = glob.glob(pattern)
        if matches:
            final_path = os.path.abspath(matches[0])

    if not final_path:
        # Fallback: best-effort to find any file with the selected extension modified recently
        candidates = sorted(
            glob.glob(os.path.join(output_dir, f"*.{audio_format}")),
            key=lambda p: os.path.getmtime(p),
            reverse=True,
        )
        if candidates:
            final_path = os.path.abspath(candidates[0])

    if not final_path:
        raise RuntimeError("Audio file not found after download. Check logs and ffmpeg availability.")

    return DownloadResult(url=url, title=info_title, filepath=final_path, ext=audio_format)


if __name__ == "__main__":
    mcp.run(transport="http")