Update app.py

app.py

@@ -1,47 +1,16 @@
-from fastapi import FastAPI, HTTPException, Query
-from fastapi.responses import FileResponse
-from pydantic import BaseModel
+from fastapi import FastAPI, Query
+from fastapi.responses import FileResponse, JSONResponse
 from pytubefix import YouTube
 from pytubefix.cli import on_progress
+from moviepy.editor import AudioFileClip
 import os
-import csv
-import uuid
-from datetime import datetime
-from typing import List, Optional
 import re
 
-app = FastAPI(title="YouTube Download API", version="1.0.0")
+app = FastAPI(title="YouTube Downloader API")
 
-# Database file
-DB_FILE = "download_logs.csv"
-DOWNLOAD_DIR = "downloads"
-
-# Initialize CSV database
-def init_database():
-    """Initialize CSV database if it doesn't exist"""
-    os.makedirs(DOWNLOAD_DIR, exist_ok=True)
-    
-    if not os.path.exists(DB_FILE):
-        with open(DB_FILE, 'w', newline='', encoding='utf-8') as f:
-            writer = csv.writer(f)
-            writer.writerow(['id', 'username', 'youtube_url', 'choice_type', 'quality', 'timestamp', 'status'])
-
-def log_download(username: str, youtube_url: str, choice_type: str, quality: str, status: str):
-    """Log download to CSV database"""
-    download_id = str(uuid.uuid4())
-    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    
-    with open(DB_FILE, 'a', newline='', encoding='utf-8') as f:
-        writer = csv.writer(f)
-        writer.writerow([download_id, username, youtube_url, choice_type, quality, timestamp, status])
-    
-    return download_id
-
-def is_youtube_url(url: str) -> bool:
-    """Validate YouTube URL"""
+def is_youtube_url(url):
     if not url or not isinstance(url, str):
         return False
-    
     youtube_regex = (
         r'(https?://)?(www\.)?'
         '(youtube|youtu|youtube-nocookie)\.(com|be)/'
@@ -49,267 +18,95 @@ def is_youtube_url(url: str) -> bool:
     )
     return bool(re.match(youtube_regex, url))
 
-# Pydantic models
-class VideoChoice(BaseModel):
-    id: str
-    type: str  # "video" or "audio"
-    quality: str
-    resolution: Optional[str] = None
-    fps: Optional[int] = None
-    file_size: Optional[str] = None
-    mime_type: str
-
-class VideoInfoResponse(BaseModel):
-    title: str
-    duration: str
-    thumbnail: str
-    author: str
-    views: int
-    choices: List[VideoChoice]
-
-class DownloadRequest(BaseModel):
-    youtube_url: str
-    choice_id: str
-    username: str
+def estimate_compressed_size(original_size_mb, bitrate_factor=0.6):
+    if not original_size_mb:
+        return "?"
+    compressed = original_size_mb * bitrate_factor
+    return f"{round(compressed, 2)} MB (est.)"
 
-class DownloadResponse(BaseModel):
-    id: str
-    status: str
-    message: str
-    download_url: Optional[str] = None
-
-# API Endpoints
-
-@app.get("/")
-def read_root():
-    """Root endpoint with API documentation"""
-    return {
-        "message": "YouTube Download API",
-        "version": "1.0.0",
-        "endpoints": {
-            "/get-choices": "GET - Get available download options for a YouTube video",
-            "/download": "POST - Download video with selected options",
-            "/logs": "GET - View download logs (admin)",
-            "/health": "GET - Health check"
-        }
-    }
-
-@app.get("/get-choices", response_model=VideoInfoResponse)
-def get_video_choices(url: str = Query(..., description="YouTube video URL")):
-    """
-    Get all available download choices for a YouTube video
-    
-    Returns:
-    - Video information (title, duration, thumbnail, etc.)
-    - Available video qualities (MP4)
-    - Available audio qualities (MP3)
-    """
-    
+def get_download_options(url):
     if not is_youtube_url(url):
-        raise HTTPException(status_code=400, detail="Invalid YouTube URL")
-    
+        return None, "Invalid YouTube URL"
+    yt = YouTube(url, on_progress_callback=on_progress)
+    streams = yt.streams.order_by('resolution').desc()
+    choices = []
+    for s in streams:
+        res = s.resolution or "Audio only"
+        type_ = "video+audio" if s.is_progressive else ("video only" if s.includes_video_track else "audio only")
+        ext = s.mime_type.split("/")[-1]
+        size_mb = round(s.filesize / 1048576, 2) if s.filesize else None
+        est_size = estimate_compressed_size(size_mb)
+        choices.append({
+            "label": f"{res} | {type_} | {ext} | {est_size}",
+            "resolution": res,
+            "type": type_,
+            "extension": ext
+        })
+    choices.append({"label": "Convert to MP3 – High Quality (320kbps)", "type": "mp3_high"})
+    choices.append({"label": "Convert to MP3 – Medium Quality (192kbps)", "type": "mp3_medium"})
+    choices.append({"label": "Convert to MP3 – Low Quality (128kbps)", "type": "mp3_low"})
+    return choices, yt.title
+
+def convert_to_mp3(temp_path, title, quality):
+    output_dir = "downloads"
+    os.makedirs(output_dir, exist_ok=True)
+    mp3_path = os.path.join(output_dir, f"{title} ({quality}).mp3")
+    clip = AudioFileClip(temp_path)
+    bitrate = {"High": "320k", "Medium": "192k", "Low": "128k"}[quality]
+    clip.write_audiofile(mp3_path, bitrate=bitrate)
+    clip.close()
+    os.remove(temp_path)
+    return mp3_path
+
+@app.get("/get_choices")
+def api_get_choices(url: str = Query(..., description="YouTube video URL")):
     try:
-        yt = YouTube(url, on_progress_callback=on_progress)
-        
-        # Get video info
-        title = yt.title
-        duration = yt.length
-        duration_str = f"{duration // 60}:{duration % 60:02d}"
-        thumbnail = yt.thumbnail_url
-        author = yt.author
-        views = yt.views
-        
-        choices = []
-        
-        # Get video streams (MP4)
-        video_streams = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc()
-        
-        for idx, stream in enumerate(video_streams):
-            file_size_mb = stream.filesize / (1024 * 1024) if stream.filesize else 0
-            choices.append(VideoChoice(
-                id=f"video_{idx}_{stream.itag}",
-                type="video",
-                quality=f"{stream.resolution} - {stream.fps}fps",
-                resolution=stream.resolution,
-                fps=stream.fps,
-                file_size=f"{file_size_mb:.2f} MB",
-                mime_type=stream.mime_type
-            ))
-        
-        # Get audio-only streams (MP3)
-        audio_streams = yt.streams.filter(only_audio=True).order_by('abr').desc()
-        
-        for idx, stream in enumerate(audio_streams):
-            file_size_mb = stream.filesize / (1024 * 1024) if stream.filesize else 0
-            choices.append(VideoChoice(
-                id=f"audio_{idx}_{stream.itag}",
-                type="audio",
-                quality=stream.abr if stream.abr else "Unknown",
-                file_size=f"{file_size_mb:.2f} MB",
-                mime_type=stream.mime_type
-            ))
-        
-        return VideoInfoResponse(
-            title=title,
-            duration=duration_str,
-            thumbnail=thumbnail,
-            author=author,
-            views=views,
-            choices=choices
-        )
-        
+        choices, title = get_download_options(url)
+        if not choices:
+            return JSONResponse(status_code=400, content={"error": title})
+        return {"title": title, "choices": choices}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error fetching video info: {str(e)}")
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
-@app.post("/download", response_model=DownloadResponse)
-def download_video(request: DownloadRequest):
-    """
-    Download video with selected quality/format
-    
-    Body:
-    - youtube_url: YouTube video URL
-    - choice_id: Choice ID from /get-choices endpoint
-    - username: Username for logging
-    """
-    
-    if not is_youtube_url(request.youtube_url):
-        raise HTTPException(status_code=400, detail="Invalid YouTube URL")
-    
-    if not request.username:
-        raise HTTPException(status_code=400, detail="Username is required")
-    
+@app.get("/download")
+def api_download(url: str = Query(...), choice: str = Query(...)):
     try:
-        yt = YouTube(request.youtube_url, on_progress_callback=on_progress)
-        
-        # Parse choice_id to get itag
-        choice_parts = request.choice_id.split('_')
-        if len(choice_parts) < 3:
-            raise HTTPException(status_code=400, detail="Invalid choice_id format")
-        
-        choice_type = choice_parts[0]  # "video" or "audio"
-        itag = int(choice_parts[2])
-        
-        # Get the stream by itag
-        stream = yt.streams.get_by_itag(itag)
-        
+        if not is_youtube_url(url):
+            return JSONResponse(status_code=400, content={"error": "Invalid YouTube URL"})
+
+        yt = YouTube(url, on_progress_callback=on_progress)
+        output_dir = "downloads"
+        os.makedirs(output_dir, exist_ok=True)
+
+        if "mp3" in choice.lower():
+            stream = yt.streams.filter(only_audio=True).first()
+            if not stream:
+                return JSONResponse(status_code=400, content={"error": "No audio stream available"})
+            temp_path = stream.download(output_path=output_dir, filename="temp.mp4")
+            if "high" in choice.lower():
+                mp3_path = convert_to_mp3(temp_path, yt.title, "High")
+            elif "medium" in choice.lower():
+                mp3_path = convert_to_mp3(temp_path, yt.title, "Medium")
+            else:
+                mp3_path = convert_to_mp3(temp_path, yt.title, "Low")
+            return FileResponse(mp3_path, filename=os.path.basename(mp3_path))
+
+        res = choice.split(" | ")[0]
+        ext = choice.split(" | ")[2]
+        stream = yt.streams.filter(res=res, mime_type=f"video/{ext}").first()
+        if not stream:
+            stream = yt.streams.filter(res=res).first()
         if not stream:
-            raise HTTPException(status_code=404, detail="Selected stream not found")
-        
-        # Download the file
-        output_path = stream.download(output_path=DOWNLOAD_DIR)
-        
-        # If audio, convert extension to mp3 (rename)
-        if choice_type == "audio":
-            base_name = os.path.splitext(output_path)[0]
-            new_path = f"{base_name}.mp3"
-            os.rename(output_path, new_path)
-            output_path = new_path
-        
-        # Log to database
-        quality = f"{stream.resolution}" if choice_type == "video" else f"{stream.abr}"
-        download_id = log_download(
-            username=request.username,
-            youtube_url=request.youtube_url,
-            choice_type=choice_type,
-            quality=quality,
-            status="success"
-        )
-        
-        # Get filename for download URL
-        filename = os.path.basename(output_path)
-        
-        return DownloadResponse(
-            id=download_id,
-            status="success",
-            message="Download completed successfully",
-            download_url=f"/download-file/{filename}"
-        )
-        
+            return JSONResponse(status_code=400, content={"error": "Selected quality not available"})
+        output_path = stream.download(output_path=output_dir)
+        return FileResponse(output_path, filename=os.path.basename(output_path))
     except Exception as e:
-        # Log failed download
-        log_download(
-            username=request.username,
-            youtube_url=request.youtube_url,
-            choice_type="unknown",
-            quality="unknown",
-            status=f"failed: {str(e)}"
-        )
-        
-        raise HTTPException(status_code=500, detail=f"Download error: {str(e)}")
-
-@app.get("/download-file/{filename}")
-def download_file(filename: str):
-    """
-    Download the actual file
-    """
-    file_path = os.path.join(DOWNLOAD_DIR, filename)
-    
-    if not os.path.exists(file_path):
-        raise HTTPException(status_code=404, detail="File not found")
-    
-    return FileResponse(
-        path=file_path,
-        filename=filename,
-        media_type='application/octet-stream'
-    )
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
-@app.get("/logs")
-def get_logs(limit: int = Query(100, description="Number of logs to retrieve")):
-    """
-    Get download logs from CSV database
-    """
-    if not os.path.exists(DB_FILE):
-        return {"logs": [], "message": "No logs found"}
-    
-    logs = []
-    with open(DB_FILE, 'r', newline='', encoding='utf-8') as f:
-        reader = csv.DictReader(f)
-        for row in reader:
-            logs.append(row)
-    
-    # Return last N logs
-    return {
-        "total": len(logs),
-        "logs": logs[-limit:] if len(logs) > limit else logs
-    }
-
-@app.get("/logs/user/{username}")
-def get_user_logs(username: str):
-    """
-    Get download logs for a specific user
-    """
-    if not os.path.exists(DB_FILE):
-        return {"logs": [], "message": "No logs found"}
-    
-    logs = []
-    with open(DB_FILE, 'r', newline='', encoding='utf-8') as f:
-        reader = csv.DictReader(f)
-        for row in reader:
-            if row['username'] == username:
-                logs.append(row)
-    
-    return {
-        "username": username,
-        "total": len(logs),
-        "logs": logs
-    }
-
-@app.get("/health")
-def health_check():
-    """Health check endpoint"""
-    return {
-        "status": "healthy",
-        "database": "connected" if os.path.exists(DB_FILE) else "not initialized",
-        "timestamp": datetime.now().isoformat()
-    }
-
-# Initialize database on startup
-@app.on_event("startup")
-def startup_event():
-    init_database()
-    print("✅ Database initialized")
-    print("✅ YouTube Download API is ready!")
+@app.get("/")
+def home():
+    return {"message": "YouTube Downloader API running! Use /get_choices and /download endpoints."}
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860) 
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)