from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
import yt_dlp
import base64
import os

app = FastAPI(
    title="YouTube Download API",
    description="API for retrieving YouTube audio/video download links using yt-dlp.",
    version="1.0.0",
    contact={
        "name": "Your Name",
        "url": "https://github.com/yourusername",
        "email": "your@email.com",
    }
)

class VideoRequest(BaseModel):
    video_id: str = Field(..., example="dQw4w9WgXcQ", description="YouTube video ID (e.g., dQw4w9WgXcQ)")

class AudioResponse(BaseModel):
    status: str = Field(..., example="success")
    audio_url: str = Field(..., description="Base64 encoded direct audio URL")

class VideoResponse(BaseModel):
    status: str = Field(..., example="success")
    video_url: str = Field(..., description="Base64 encoded direct video URL")

COOKIES_PATH = "cookies.txt"

@app.post("/audio/", response_model=AudioResponse, summary="Get YouTube Audio URL")
def get_audio_link(req: VideoRequest):
    """
    Returns a downloadable audio link for a given YouTube video ID.
    The URL is base64-encoded.
    """
    url = f"https://www.youtube.com/watch?v={req.video_id}"
    ydl_opts = {
        'quiet': True,
        'format': 'bestaudio/best',
        'skip_download': True,
        'forceurl': True,
        'cookies': COOKIES_PATH if os.path.isfile(COOKIES_PATH) else None,
        'cache_dir': False
    }
    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=False)
            for f in info['formats']:
                if f.get('acodec') != 'none' and f.get('vcodec') == 'none':
                    audio_url = f['url']
                    b64 = base64.b64encode(audio_url.encode()).decode()
                    return {"status": "success", "audio_url": b64}
        raise Exception("No audio found")
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.post("/video/", response_model=VideoResponse, summary="Get YouTube Video URL")
def get_video_link(req: VideoRequest):
    """
    Returns a downloadable video link (max 720p) for a given YouTube video ID.
    The URL is base64-encoded.
    """
    url = f"https://www.youtube.com/watch?v={req.video_id}"
    ydl_opts = {
        'quiet': True,
        'format': 'best[height<=720][width<=1280]',
        'skip_download': True,
        'forceurl': True,
        'cookies': COOKIES_PATH if os.path.isfile(COOKIES_PATH) else None,
        'cache_dir': False
    }
    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=False)
            for f in info['formats']:
                if f.get('vcodec') != 'none':
                    video_url = f['url']
                    b64 = base64.b64encode(video_url.encode()).decode()
                    return {"status": "success", "video_url": b64}
        raise Exception("No video found")
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))