ohm_streaming/app/routers/router_downloads.py

"""
Download management routes for Ohm Stream Downloader API.
"""

import json
from typing import Optional
from pathlib import Path
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Query, Request, Response
from fastapi.templating import Jinja2Templates
from fastapi.responses import HTMLResponse, FileResponse

from app.download_manager import DownloadManager
from app.models import DownloadRequest, DownloadStatus
from app.models.auth import User
from app.routers.router_auth import get_current_user_from_token, get_optional_user

router = APIRouter(prefix="/api/downloads", tags=["downloads"])
templates = Jinja2Templates(directory="templates")


def get_download_manager() -> DownloadManager:
    from main import download_manager
    return download_manager


@router.get("")
async def get_downloads(
    request: Request,
    html: bool = Query(False),
    download_manager: DownloadManager = Depends(get_download_manager),
    current_user: Optional[User] = Depends(get_optional_user),
):
    """Get list of all download tasks. Returns HTML for HTMX requests."""
    is_htmx = request.headers.get("HX-Request") == "true" or request.headers.get("HX-Request")

    if current_user is None and (html or is_htmx):
        return templates.TemplateResponse(
            "components/login_prompt.html", {"request": request}
        )

    if current_user is None:
        raise HTTPException(status_code=401, detail="Authentication required")

    tasks = download_manager.get_all_tasks()

    if html or is_htmx:
        print(f"[DOWNLOADS] HTML Request. Found {len(tasks)} tasks.")
        return templates.TemplateResponse(
            "components/downloads_list.html",
            {"request": request, "tasks": tasks}
        )

    print(f"[DOWNLOADS] API Request. Returning JSON.")
    return {"downloads": tasks}


@router.post("")
async def create_download(
    download_request: DownloadRequest,
    download_manager: DownloadManager = Depends(get_download_manager),
    current_user=Depends(get_current_user_from_token),
):
    """Create a new download task"""
    return download_manager.create_task(download_request)


@router.get("/{task_id}")
async def get_download_status(
    task_id: str,
    download_manager: DownloadManager = Depends(get_download_manager),
    current_user: Optional[User] = Depends(get_optional_user),
):
    """Get status of a specific download task"""
    if current_user is None:
        raise HTTPException(status_code=401, detail="Authentication required")

    task = download_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    return task


@router.post("/{task_id}/pause")
async def pause_download(
    task_id: str,
    download_manager: DownloadManager = Depends(get_download_manager),
    current_user=Depends(get_current_user_from_token),
):
    """Pause a download task"""
    if download_manager.pause_download(task_id):
        return {"status": "success", "message": "Download paused"}
    raise HTTPException(status_code=400, detail="Failed to pause download")


@router.post("/{task_id}/resume")
async def resume_download(
    task_id: str,
    download_manager: DownloadManager = Depends(get_download_manager),
    current_user=Depends(get_current_user_from_token),
):
    """Resume a paused download task"""
    if download_manager.resume_download(task_id):
        return {"status": "success", "message": "Download resumed"}
    raise HTTPException(status_code=400, detail="Failed to resume download")


@router.delete("/{task_id}")
async def cancel_download(
    task_id: str,
    download_manager: DownloadManager = Depends(get_download_manager),
    current_user=Depends(get_current_user_from_token),
):
    """Cancel and delete a download task"""
    if hasattr(download_manager, "cancel_download"):
        if download_manager.cancel_download(task_id):
            return {"status": "success", "message": "Download cancelled"}
    
    if task_id in download_manager.tasks:
        del download_manager.tasks[task_id]
        return {"status": "success", "message": "Download removed"}
        
    raise HTTPException(status_code=400, detail="Failed to cancel download")


@router.get("/video/{task_id}")
async def stream_video(
    task_id: str,
    download_manager: DownloadManager = Depends(get_download_manager),
    current_user: Optional[User] = Depends(get_optional_user),
):
    """Stream a completed download as video"""
    if current_user is None:
        raise HTTPException(status_code=401, detail="Authentication required")
    task = download_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    if task.status != DownloadStatus.COMPLETED or not task.file_path:
        raise HTTPException(status_code=400, detail="Download not completed")
    file_path = Path(task.file_path)
    if not file_path.exists():
        raise HTTPException(status_code=404, detail="File not found on disk")
    media_types = {
        ".mp4": "video/mp4", ".mkv": "video/x-matroska", ".avi": "video/x-msvideo",
        ".webm": "video/webm", ".flv": "video/x-flv",
    }
    media_type = media_types.get(file_path.suffix.lower(), "video/mp4")
    return FileResponse(str(file_path), media_type=media_type)


@router.post("/{task_id}/retry")
async def retry_download(
    task_id: str,
    background_tasks: BackgroundTasks,
    response: Response,
    download_manager: DownloadManager = Depends(get_download_manager),
    current_user: User = Depends(get_current_user_from_token),
):
    """Retry a failed or cancelled download"""
    task = download_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    if task.status not in ("failed", "cancelled"):
        raise HTTPException(status_code=400, detail="Can only retry failed or cancelled downloads")
    task.status = DownloadStatus.PENDING
    task.progress = 0.0
    if hasattr(download_manager, "_process_download"):
        background_tasks.add_task(download_manager._process_download, task_id)
    response.headers["HX-Trigger"] = json.dumps(
        {"show-toast": {"message": "Telechargement relance", "type": "success"}}
    )
    return {"status": "retrying"}


@router.post("/cancel-all")
async def cancel_all_downloads(
    response: Response,
    download_manager: DownloadManager = Depends(get_download_manager),
    current_user: User = Depends(get_current_user_from_token),
):
    """Cancel all active downloads"""
    count = 0
    for tid, task in list(download_manager.tasks.items()):
        if task.status in ("downloading", "pending"):
            download_manager.cancel_download(tid) if hasattr(download_manager, "cancel_download") else download_manager.tasks.pop(tid, None)
            count += 1
    response.headers["HX-Trigger"] = json.dumps(
        {"show-toast": {"message": f"{count} telechargement(s) annule(s)", "type": "info"}}
    )
    return {"status": "cancelled", "count": count}


@router.post("/cleanup")
async def cleanup_completed(
    download_manager: DownloadManager = Depends(get_download_manager),
    current_user=Depends(get_current_user_from_token),
):
    """Remove all completed tasks from the list"""
    if hasattr(download_manager, "cleanup_tasks"):
        count = download_manager.cleanup_tasks()
        return {"status": "success", "message": f"Cleaned up {count} tasks"}
    
    to_delete = [tid for tid, t in download_manager.tasks.items() if t.status == "completed"]
    for tid in to_delete:
        del download_manager.tasks[tid]
    return {"status": "success", "message": f"Cleaned up {len(to_delete)} tasks"}