Update app.py

app.py

@@ -1,41 +1,90 @@
-from fastapi import FastAPI, HTTPException, Query
-from fastapi.responses import RedirectResponse
+from fastapi import FastAPI, HTTPException, Query, BackgroundTasks
+from pydantic import BaseModel
 import subprocess
-import urllib.parse
+import uuid
+import os
+import json
+import shutil
 
-app = FastAPI(
-    title="YouTube Audio Streamer",
-    description="Stream best audio of a YouTube video using yt-dlp.",
-    version="1.0.2",
-    docs_url="/docs",
-    redoc_url="/redoc"
-)
+app = FastAPI()
 
-@app.get("/", tags=["Health"])
-async def root():
-    return {"status": "FastAPI working", "message": "YT-DLP backend ready"}
+tasks = {}
 
-@app.get("/stream", tags=["Stream"])
-async def stream_audio(url: str = Query(..., description="Full YouTube video URL")):
-    decoded_url = urllib.parse.unquote(url)
-
-    if "youtube.com/watch?v=" not in decoded_url and "youtu.be/" not in decoded_url:
-        raise HTTPException(status_code=400, detail="Invalid YouTube URL.")
-
-    # Remove URL query params (yt-dlp can fail with things like ?si=xxx)
-    base_url = decoded_url.split("&")[0].split("?")[0]
+class DownloadRequest(BaseModel):
+    url: str
+    format_id: str  # yt-dlp format code
+    output_name: str = None
+    convert_to: str = None  # e.g. mp3 or mp4
 
+@app.get("/info")
+async def get_video_info(url: str = Query(...)):
     try:
-        result = subprocess.run(
-            ["yt-dlp", "-f", "bestaudio", "--get-url", base_url],
-            capture_output=True, text=True
-        )
+        cmd = [
+            "yt-dlp",
+            "--no-warnings",
+            "--skip-download",
+            "--print-json",
+            url
+        ]
+        result = subprocess.run(cmd, capture_output=True, text=True)
         if result.returncode != 0:
-            raise HTTPException(status_code=500, detail=f"yt-dlp error: {result.stderr.strip()}")
-
-        audio_url = result.stdout.strip()
-        if not audio_url:
-            raise HTTPException(status_code=404, detail="No audio URL found.")
-        return RedirectResponse(audio_url)
+            raise Exception(result.stderr)
+        data = json.loads(result.stdout)
+        return {
+            "title": data.get("title"),
+            "thumbnail": data.get("thumbnail"),
+            "formats": [
+                {"id": f["format_id"], "ext": f["ext"], "note": f.get("format_note"), "filesize": f.get("filesize")}
+                for f in data.get("formats", [])
+                if f.get("vcodec") != "none" or f.get("acodec") != "none"
+            ]
+        }
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error getting info: {e}")
+
+@app.post("/download")
+async def download_media(request: DownloadRequest, background_tasks: BackgroundTasks):
+    task_id = str(uuid.uuid4())
+    tasks[task_id] = {"status": "queued"}
+
+    def _do_download():
+        try:
+            output_file = f"{request.output_name or 'video'}_{task_id}"
+            if request.convert_to:
+                output_file += f".{request.convert_to}"
+            else:
+                output_file += ".%(ext)s"
+
+            cmd = [
+                "yt-dlp",
+                "-f", request.format_id,
+                "-o", f"/tmp/{output_file}",
+                request.url
+            ]
+            if request.convert_to:
+                cmd += ["--recode-video", request.convert_to]
+
+            tasks[task_id]["status"] = "downloading"
+            subprocess.run(cmd, capture_output=True, text=True)
+            tasks[task_id]["status"] = "completed"
+            tasks[task_id]["file"] = f"/tmp/{output_file}"
+        except Exception as e:
+            tasks[task_id]["status"] = "error"
+            tasks[task_id]["error"] = str(e)
+
+    background_tasks.add_task(_do_download)
+    return {"task_id": task_id}
+
+@app.get("/progress/{task_id}")
+def check_progress(task_id: str):
+    if task_id not in tasks:
+        raise HTTPException(status_code=404, detail="Task ID not found")
+    return tasks[task_id]
+
+@app.get("/cancel/{task_id}")
+def cancel_task(task_id: str):
+    # Basic cancellation support stub (real cancellation needs async task mgmt)
+    if task_id in tasks and tasks[task_id]["status"] not in ["completed", "error"]:
+        tasks[task_id]["status"] = "cancelled"
+        return {"detail": "Task cancelled"}
+    raise HTTPException(status_code=404, detail="Task not found or already done")