app.py

from fastapi import FastAPI, HTTPException, Query, BackgroundTasks
from pydantic import BaseModel
import subprocess
import uuid
import os
import json
import shutil

app = FastAPI()

tasks = {}

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:
        cmd = [
            "yt-dlp",
            "--no-warnings",
            "--skip-download",
            "--print-json",
            url
        ]
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            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 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")